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; }