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)