def on_batch_end(self, state: RunnerState):
lm = state.loader_name
names = state.input.get(self.name_key, [])
features = state.input[self.input_key]
logits = state.output[self.output_key]
logits = torch.unsqueeze_(logits, dim=1) \
if len(logits.shape) < 4 \
else logits
if self.mask_type == "soft":
probs = F.sigmoid(logits)
else:
probs = F.softmax(logits, dim=1)
features = features.detach().cpu().numpy()
features = np.transpose(features, (0, 2, 3, 1))
probs = probs.detach().cpu().numpy()
probs = np.transpose(probs, (0, 2, 3, 1))
# colors = self._get_spaced_colors2(n_colors=probs.shape[3])
for i in range(probs.shape[0]):
img = np.uint8(255 * (self.std * features[i] + self.mean))
try:
suffix = names[i]
except IndexError:
suffix = f"{self.counter:06d}"
self.counter += 1
# shw = img.copy()
masks = []
for t in range(probs.shape[3]):
mask = probs[i, :, :, t] > self.threshold \
if self.threshold is not None \
else probs[i, :, :, t]
mask = mask.astype(np.float32)
if self.dump_mask:
mask_ = np.concatenate([np.expand_dims(mask, -1)] * 3,
axis=-1).astype(np.float32) * 255
filename_ = f"{self.out_prefix}/{lm}/{suffix}_{t}.jpg"
cv2.imwrite(filename_, mask_)
masks.append(mask)
# mask = mask - erosion(mask, disk(4))
# shw[mask > 0.5] = colors[t]
shw = binary_mask_to_overlay_image(img.copy(), masks)
filename = f"{self.out_prefix}/{lm}/{suffix}.jpg"
cv2.imwrite(filename, shw[:, :, ::-1])
def on_batch_end(self, runner: IRunner):
"""Batch end hook.
Args:
runner (IRunner): current runner
"""
lm = runner.loader_name
names = runner.input.get(self.name_key, [])
features = runner.input[self.input_key].detach().cpu()
images = utils.tensor_to_ndimage(features)
logits = runner.output[self.output_key]
logits = (
torch.unsqueeze_(logits, dim=1)
if len(logits.shape) < 4
else logits
)
if self.mask_type == "soft":
probabilities = torch.sigmoid(logits)
else:
probabilities = F.softmax(logits, dim=1)
probabilities = probabilities.detach().cpu().numpy()
masks = []
for probability in probabilities:
mask = np.zeros_like(probability[0], dtype=np.int32)
for i, ch in enumerate(probability):
mask[ch >= self.threshold] = i + 1
masks.append(mask)
for index, (image, mask) in enumerate(zip(images, masks)):
try:
suffix = names[index]
except IndexError:
suffix = f"{self.counter:06d}"
self.counter += 1
mask = label2rgb(mask, bg_label=0)
image = (
image * (1 - self.mask_strength) + mask * self.mask_strength
)
image = (image * 255).clip(0, 255).round().astype(np.uint8)
filename = f"{self.out_prefix}/{lm}/{suffix}.jpg"
imageio.imwrite(filename, image)
def on_batch_end(self, state):
lm = state.loader_mode
features = state.input[self.input_key]
logits = state.output[self.output_key]
logits = torch.unsqueeze_(logits, dim=1) \
if len(logits.shape) < 4 \
else logits
if self.mask_type == "soft":
probs = F.sigmoid(logits)
else:
probs = F.softmax(logits, dim=1)
features = features.detach().cpu().numpy()
features = np.transpose(features, (0, 2, 3, 1))
probs = probs.detach().cpu().numpy()
probs = np.transpose(probs, (0, 2, 3, 1))
for i in range(probs.shape[0]):
img = np.uint8(255 * (self.std * features[i] + self.mean))
filename = f"{self.out_prefix}/{lm}/{self.counter}.jpg"
cv2.imwrite(filename, img)
for t in range(probs.shape[-1]):
mask = probs[i, :, :, t] > self.threshold \
if self.threshold is not None \
else probs[i, :, :, t]
mask = np.float32(np.expand_dims(mask, -1))
masked_img = img * mask
# @TODO: better naming
filename = f"{self.out_prefix}/{lm}/{self.counter}_{t}.jpg"
cv2.imwrite(filename, masked_img)
self.counter += 1
