affinity_graph_ptr<T> mult_aff( const affinity_graph_ptr<T>& aff, int n )
{
std::ptrdiff_t xdim = aff->shape()[0];
std::ptrdiff_t ydim = aff->shape()[1];
std::ptrdiff_t zdim = aff->shape()[2];
affinity_graph_ptr<T> r(new affinity_graph<T>
(boost::extents[xdim*n][ydim*n][zdim*n][3],
boost::fortran_storage_order()));
for ( std::ptrdiff_t z = 0; z < n; ++z )
for ( std::ptrdiff_t y = 0; y < n; ++y )
for ( std::ptrdiff_t x = 0; x < n; ++x )
{
typedef boost::multi_array_types::index_range range;
(*r)[boost::indices[range(x*xdim,(x+1)*xdim)]
[range(y*ydim,(y+1)*ydim)]
[range(z*zdim,(z+1)*zdim)]
[range(0,3)]] = *aff;
}
return r;
}
inline region_graph_ptr<ID,F>
get_region_graph( const affinity_graph_ptr<F>& aff_ptr,
const volume_ptr<ID> seg_ptr,
std::size_t max_segid)
{
std::ptrdiff_t xdim = aff_ptr->shape()[0];
std::ptrdiff_t ydim = aff_ptr->shape()[1];
std::ptrdiff_t zdim = aff_ptr->shape()[2];
volume<ID>& seg = *seg_ptr;
affinity_graph<F> aff = *aff_ptr;
region_graph_ptr<ID,F> rg_ptr( new region_graph<ID,F> );
region_graph<ID,F>& rg = *rg_ptr;
std::vector<std::map<ID,F>> edges(max_segid+1);
for ( std::ptrdiff_t z = 0; z < zdim; ++z )
for ( std::ptrdiff_t y = 0; y < ydim; ++y )
for ( std::ptrdiff_t x = 0; x < xdim; ++x )
{
if ( (x > 0) && seg[x][y][z] && seg[x-1][y][z] )
{
auto mm = std::minmax(seg[x][y][z], seg[x-1][y][z]);
F& curr = edges[mm.first][mm.second];
curr = std::max(curr, aff[x][y][z][0]);
}
if ( (y > 0) && seg[x][y][z] && seg[x][y-1][z] )
{
auto mm = std::minmax(seg[x][y][z], seg[x][y-1][z]);
F& curr = edges[mm.first][mm.second];
curr = std::max(curr, aff[x][y][z][1]);
}
if ( (z > 0) && seg[x][y][z] && seg[x][y][z-1] )
{
auto mm = std::minmax(seg[x][y][z], seg[x][y][z-1]);
F& curr = edges[mm.first][mm.second];
curr = std::max(curr, aff[x][y][z][2]);
}
}
for ( ID id1 = 1; id1 <= max_segid; ++id1 )
{
for ( const auto& p: edges[id1] )
{
rg.emplace_back(p.second, id1, p.first);
}
}
std::cout << "Region graph size: " << rg.size() << std::endl;
std::sort(std::begin(rg), std::end(rg),
std::greater<std::tuple<F,ID,ID>>());
std::cout << "Sorted" << std::endl;
return rg_ptr;
}
inline std::pair<volume_ptr<ID>, ID>
simple_watershed( const affinity_graph_ptr<F>& aff_ptr,
const L& lowv,
const H& highv,
std::vector<std::size_t>& counts )
{
typedef F float_type;
typedef watershed_traits<ID> traits;
float_type low = static_cast<float_type>(lowv);
float_type high = static_cast<float_type>(highv);
std::ptrdiff_t xdim = aff_ptr->shape()[0];
std::ptrdiff_t ydim = aff_ptr->shape()[1];
std::ptrdiff_t zdim = aff_ptr->shape()[2];
std::ptrdiff_t size = xdim * ydim * zdim;
volume_ptr<ID> seg_ptr( new volume<ID>(boost::extents[xdim][ydim][zdim],
boost::fortran_storage_order()));
affinity_graph<F>& aff = *aff_ptr;
volume<ID>& seg = *seg_ptr;
ID* seg_raw = seg_ptr->data();
for ( std::ptrdiff_t z = 0; z < zdim; ++z )
for ( std::ptrdiff_t y = 0; y < ydim; ++y )
for ( std::ptrdiff_t x = 0; x < xdim; ++x )
{
ID& id = seg[x][y][z] = 0;
F negx = (x>0) ? aff[x][y][z][0] : low;
F negy = (y>0) ? aff[x][y][z][1] : low;
F negz = (z>0) ? aff[x][y][z][2] : low;
F posx = (x<(xdim-1)) ? aff[x+1][y][z][0] : low;
F posy = (y<(ydim-1)) ? aff[x][y+1][z][1] : low;
F posz = (z<(zdim-1)) ? aff[x][y][z+1][2] : low;
F m = std::max({negx,negy,negz,posx,posy,posz});
if ( m > low )
{
if ( negx == m || negx >= high ) { id |= 0x01; }
if ( negy == m || negy >= high ) { id |= 0x02; }
if ( negz == m || negz >= high ) { id |= 0x04; }
if ( posx == m || posx >= high ) { id |= 0x08; }
if ( posy == m || posy >= high ) { id |= 0x10; }
if ( posz == m || posz >= high ) { id |= 0x20; }
}
}
const std::ptrdiff_t dir[6] = { -1, -xdim, -xdim*ydim, 1, xdim, xdim*ydim };
const ID dirmask[6] = { 0x01, 0x02, 0x04, 0x08, 0x10, 0x20 };
const ID idirmask[6] = { 0x08, 0x10, 0x20, 0x01, 0x02, 0x04 };
// get plato escapers
std::vector<std::ptrdiff_t> bfs;
for ( std::ptrdiff_t idx = 0; idx < size; ++idx )
{
for ( std::ptrdiff_t d = 0; d < 6; ++d )
{
if ( seg_raw[idx] & dirmask[d] )
{
if ( !(seg_raw[idx+dir[d]] & idirmask[d]) )
{
seg_raw[idx] |= 0x40;
bfs.push_back(idx);
d = 6; // break;
}
}
}
}
std::size_t bfs_index = 0;
while ( bfs_index < bfs.size() )
{
std::ptrdiff_t idx = bfs[bfs_index];
ID to_set = 0;
for ( std::ptrdiff_t d = 0; d < 6; ++d )
{
if ( seg_raw[idx] & dirmask[d] )
{
if ( seg_raw[idx+dir[d]] & idirmask[d] )
{
if ( !( seg_raw[idx+dir[d]] & 0x40 ) )
{
bfs.push_back(idx+dir[d]);
seg_raw[idx+dir[d]] |= 0x40;
}
}
else
{
to_set = dirmask[d];
}
}
}
seg_raw[idx] = to_set;
++bfs_index;
}
bfs.clear();
//std::vector<std::size_t> counts({0});
counts.resize(1);
counts[0] = 0;
ID next_id = static_cast<ID>(1);
for ( std::ptrdiff_t idx = 0; idx < size; ++idx )
{
if ( seg_raw[idx] == 0 )
{
seg_raw[idx] |= traits::high_bit;
++counts[0];
}
if ( !( seg_raw[idx] & traits::high_bit ) && seg_raw[idx] )
{
bfs.push_back(idx);
bfs_index = 0;
seg_raw[idx] |= 0x40;
while ( bfs_index < bfs.size() )
{
std::ptrdiff_t me = bfs[bfs_index];
for ( std::ptrdiff_t d = 0; d < 6; ++d )
{
if ( seg_raw[me] & dirmask[d] )
{
std::ptrdiff_t him = me + dir[d];
if ( seg_raw[him] & traits::high_bit )
{
counts[ seg_raw[him] & ~traits::high_bit ]
+= bfs.size();
for ( auto& it: bfs )
{
seg_raw[it] = seg_raw[him];
}
bfs.clear();
d = 6; // break
}
else if ( !( seg_raw[him] & 0x40 ) )
{
seg_raw[him] |= 0x40;
bfs.push_back( him );
}
}
}
++bfs_index;
}
if ( bfs.size() )
{
counts.push_back( bfs.size() );
for ( auto& it: bfs )
{
seg_raw[it] = traits::high_bit | next_id;
}
++next_id;
bfs.clear();
}
}
}
std::cout << "found: " << (next_id-1) << " components\n";
for ( std::ptrdiff_t idx = 0; idx < size; ++idx )
{
seg_raw[idx] &= traits::mask;
}
return std::make_pair(seg_ptr, next_id-1);
}
inline std::tuple< volume_ptr<ID>, std::vector<std::size_t> >
watershed( const affinity_graph_ptr<F>& aff_ptr, const L& lowv, const H& highv )
{
using affinity_t = F;
using id_t = ID;
using traits = watershed_traits<id_t>;
affinity_t low = static_cast<affinity_t>(lowv);
affinity_t high = static_cast<affinity_t>(highv);
std::ptrdiff_t xdim = aff_ptr->shape()[0];
std::ptrdiff_t ydim = aff_ptr->shape()[1];
std::ptrdiff_t zdim = aff_ptr->shape()[2];
std::ptrdiff_t size = xdim * ydim * zdim;
std::tuple< volume_ptr<id_t>, std::vector<std::size_t> > result
( volume_ptr<id_t>( new volume<id_t>(boost::extents[xdim][ydim][zdim],
boost::fortran_storage_order())),
std::vector<std::size_t>(1) );
auto& counts = std::get<1>(result);
counts[0] = 0;
affinity_graph<F>& aff = *aff_ptr;
volume<id_t>& seg = *std::get<0>(result);
id_t* seg_raw = seg.data();
for ( std::ptrdiff_t z = 0; z < zdim; ++z )
for ( std::ptrdiff_t y = 0; y < ydim; ++y )
for ( std::ptrdiff_t x = 0; x < xdim; ++x )
{
id_t& id = seg[x][y][z] = 0;
F negx = (x>0) ? aff[x][y][z][0] : low;
F negy = (y>0) ? aff[x][y][z][1] : low;
F negz = (z>0) ? aff[x][y][z][2] : low;
F posx = (x<(xdim-1)) ? aff[x+1][y][z][0] : low;
F posy = (y<(ydim-1)) ? aff[x][y+1][z][1] : low;
F posz = (z<(zdim-1)) ? aff[x][y][z+1][2] : low;
F m = std::max({negx,negy,negz,posx,posy,posz});
if ( m > low )
{
if ( negx == m || negx >= high ) { id |= 0x01; }
if ( negy == m || negy >= high ) { id |= 0x02; }
if ( negz == m || negz >= high ) { id |= 0x04; }
if ( posx == m || posx >= high ) { id |= 0x08; }
if ( posy == m || posy >= high ) { id |= 0x10; }
if ( posz == m || posz >= high ) { id |= 0x20; }
}
}
const std::ptrdiff_t dir[6] = { -1, -xdim, -xdim*ydim, 1, xdim, xdim*ydim };
const id_t dirmask[6] = { 0x01, 0x02, 0x04, 0x08, 0x10, 0x20 };
const id_t idirmask[6] = { 0x08, 0x10, 0x20, 0x01, 0x02, 0x04 };
// get plato corners
std::vector<std::ptrdiff_t> bfs;
for ( std::ptrdiff_t idx = 0; idx < size; ++idx )
{
for ( std::ptrdiff_t d = 0; d < 6; ++d )
{
if ( seg_raw[idx] & dirmask[d] )
{
if ( !(seg_raw[idx+dir[d]] & idirmask[d]) )
{
seg_raw[idx] |= 0x40;
bfs.push_back(idx);
d = 6; // break;
}
}
}
}
// divide the plateaus
std::size_t bfs_index = 0;
while ( bfs_index < bfs.size() )
{
std::ptrdiff_t idx = bfs[bfs_index];
id_t to_set = 0;
for ( std::ptrdiff_t d = 0; d < 6; ++d )
{
if ( seg_raw[idx] & dirmask[d] )
{
if ( seg_raw[idx+dir[d]] & idirmask[d] )
{
if ( !( seg_raw[idx+dir[d]] & 0x40 ) )
{
bfs.push_back(idx+dir[d]);
seg_raw[idx+dir[d]] |= 0x40;
}
}
else
{
to_set = dirmask[d];
}
}
}
seg_raw[idx] = to_set;
++bfs_index;
}
bfs.clear();
// main watershed logic
id_t next_id = 1;
for ( std::ptrdiff_t idx = 0; idx < size; ++idx )
{
if ( seg_raw[idx] == 0 )
{
seg_raw[idx] |= traits::high_bit;
++counts[0];
}
if ( !( seg_raw[idx] & traits::high_bit ) && seg_raw[idx] )
{
bfs.push_back(idx);
bfs_index = 0;
seg_raw[idx] |= 0x40;
while ( bfs_index < bfs.size() )
{
std::ptrdiff_t me = bfs[bfs_index];
for ( std::ptrdiff_t d = 0; d < 6; ++d )
{
if ( seg_raw[me] & dirmask[d] )
{
std::ptrdiff_t him = me + dir[d];
if ( seg_raw[him] & traits::high_bit )
{
counts[ seg_raw[him] & ~traits::high_bit ]
+= bfs.size();
for ( auto& it: bfs )
{
seg_raw[it] = seg_raw[him];
}
bfs.clear();
d = 6; // break
}
else if ( !( seg_raw[him] & 0x40 ) )
{
seg_raw[him] |= 0x40;
bfs.push_back( him );
}
}
}
++bfs_index;
}
if ( bfs.size() )
{
counts.push_back( bfs.size() );
for ( auto& it: bfs )
{
seg_raw[it] = traits::high_bit | next_id;
}
++next_id;
bfs.clear();
}
}
}
std::cout << "found: " << (next_id-1) << " components\n";
for ( std::ptrdiff_t idx = 0; idx < size; ++idx )
{
seg_raw[idx] &= traits::mask;
}
return result;
}
