def _oversegmentation_impl_masked(self, input_, mask): assert input_.ndim == 4 assert input_.ndim == 4 assert len(input_) == len( self.offsets), "%s, %i" % (str(input_.shape), len(self.offsets)) if self.invert_repulsive_channels: input_[self.seperating_channel:] *= -1 input_[self.seperating_channel:] += 1 # we mask by setting the local / attractive affinity channels to 1, so they will never be drawn # TODO is this correct ? input_[:3, mask] = 1 # TODO make sure the pre-sorting works # sorted_edges = self._sorted_edges(input_) sorted_edges = np.argsort(input_.ravel()) # run the mst watershed mst = cmst.ConstrainedWatershed(np.array(input_.shape[1:]), self.offsets, self.seperating_channel, np.array([1, 1, 1])) # self.strides) # need [1, 1, 1] strides if we subsample before sorting # don't need this if we subsample before sorting # if self.randomize_bounds: # mst.compute_randomized_bounds() mst.repulsive_ucc_mst_cut(sorted_edges, 0) segmentation = mst.get_flat_label_image().reshape(input_.shape[1:]) return segmentation, int(segmentation.max())
def _oversegmentation_impl(self, input_): assert input_.ndim == 4 assert len(input_) == len( self.offsets), "%s, %i" % (str(input_.shape), len(self.offsets)) if self.invert_repulsive_channels: input_[self.seperating_channel:] *= -1 input_[self.seperating_channel:] += 1 # TODO make sure the pre-sorting works # sorted_edges = self._sorted_edges(input_) sorted_edges = np.argsort(input_.ravel()) # run the mst watershed mst = cmst.ConstrainedWatershed(np.array(input_.shape[1:]), self.offsets, self.seperating_channel, np.array(self.strides)) # np.array([1, 1, 1])) # self.strides) # need [1, 1, 1] strides if we subsample before sorting # don't need this if we subsample before sorting # if self.randomize_bounds: # mst.compute_randomized_bounds() mst.repulsive_ucc_mst_cut(sorted_edges, 0) segmentation = mst.get_flat_label_image().reshape(input_.shape[1:]) return segmentation
def mutex_segmentation(affs, offsets): import constrained_mst as cmst mst = cmst.ConstrainedWatershed(np.array(affs.shape[1:]), np.array(offsets), 3, np.array([1, 1, 1])) affs[3:] *= -1 affs[3:] += 1 print("Computing mws ...") t0 = time.time() sorted_edges = np.argsort(affs.ravel()) mst.repulsive_ucc_mst_cut(sorted_edges, 0) print("... in %f s" % (time.time() - t0,)) seg = mst.get_flat_label_image().reshape(affs.shape[1:]) return seg
def run_mst(affinities, offsets, stride, seperating_channel=2, invert_dam_channels=True, randomize_bounds=True): import constrained_mst as cmst assert len(affinities) == len(offsets), "%s, %i" % (str( affinities.shape), len(offsets)) affinities_ = np.require(affinities.copy(), requirements='C') if invert_dam_channels: affinities_[seperating_channel:] *= -1 affinities_[seperating_channel:] += 1 sorted_edges = np.argsort(affinities_.ravel()) # run the mst watershed vol_shape = affinities_.shape[1:] mst = cmst.ConstrainedWatershed(np.array(vol_shape), offsets, seperating_channel, stride) if randomize_bounds: mst.compute_randomized_bounds() mst.repulsive_ucc_mst_cut(sorted_edges, 0) segmentation = mst.get_flat_label_image().reshape(vol_shape) return segmentation
def damws_superpixel(self, affinities): assert affinities.shape[0] >= len(self.offsets) # dam channels if ncessary if self.invert_dam_channels: affinities_ = affinities.copy() affinities_[self.seperating_channel:] *= -1 affinities_[self.seperating_channel:] += 1 else: affinities_ = affinities # sort all edges sorted_edges = np.argsort(affinities_.ravel()) # run the mst watershed vol_shape = affinities_.shape[1:] mst = cmst.ConstrainedWatershed(np.array(vol_shape), self.offsets, self.seperating_channel, np.array(self.stride)) mst.repulsive_ucc_mst_cut(sorted_edges, 0) if self.randomize_bounds: mst.compute_randomized_bounds() segmentation = mst.get_flat_label_image().reshape(vol_shape) max_label = segmentation.max() return segmentation, max_label
def mws_segmentation(embedding, offsets='default-3D', affinity_measure=logistic_similarity, pass_offset=False, ATT_C=None, repulsive_strides=None, percentile=5, return_affinities=False, attraction_factor=1, z_delay=0): assert cmst is not None, 'need constrained_mst for mws_segmentation' offsets = get_offsets(offsets) if ATT_C is None: ATT_C = len(offsets[0]) emb_shape = embedding.shape img_shape = embedding.shape[-len(offsets[0]):] n_img_dims = len(img_shape) if repulsive_strides is None: repulsive_strides = (1, ) * (n_img_dims - 2) + (8, 8) repulsive_strides = np.array(repulsive_strides, dtype=int) # actually not needed, huge offsets are fine # for off in offsets: # assert all(abs(o) < s for o, s in zip(off, emb.shape[-len(off):])), \ # f'offset {off} is to big for image of shape {img_shape}' if affinity_measure is not None: affinities = embedding_to_affinities(embedding, offsets=offsets, affinity_measure=affinity_measure, pass_offset=pass_offset) else: affinities = embedding affinities = affinities.contiguous().view((-1, len(offsets)) + emb_shape[-n_img_dims:]) if percentile is not None: affinities -= np.percentile(affinities, percentile) affinities[:, :ATT_C] *= -1 else: affinities[:, :ATT_C] *= -1 affinities[:, :ATT_C] += 1 affinities[:, :ATT_C] *= attraction_factor if z_delay != 0: affinities[:, (offsets[:, 0] != 0).astype(np.uint8)] += z_delay result = [] for aff in affinities: dws = cmst.ConstrainedWatershed(np.array(img_shape), offsets, ATT_C, repulsive_strides) sorted_edges = np.argsort(aff, axis=None) dws.repulsive_ucc_mst_cut(sorted_edges, 0) seg = label(dws.get_flat_label_image().reshape(img_shape)) seg = np.random.permutation(seg.max() + 1)[seg] result.append(seg) result = np.stack( result, axis=-(n_img_dims + 1)).reshape(emb_shape[:-n_img_dims - 1] + emb_shape[-n_img_dims:]) if return_affinities: return torch.from_numpy(result), affinities else: return torch.from_numpy(result)