def _draw_masks(
self,
writer: SummaryWriter,
global_step: int,
image_over_predicted_mask: np.ndarray,
image_over_gt_mask: Optional[np.ndarray] = None,
) -> None:
"""
Draw image over mask to tensorboard
Args:
writer: loader writer
global_step: global step
image_over_predicted_mask: image over predicted mask
image_over_gt_mask: image over ground truth mask
"""
if image_over_gt_mask is not None:
writer.add_image(
f"{self.log_name} Ground Truth",
image_over_gt_mask,
global_step=global_step,
dataformats="HWC",
)
writer.add_image(
f"{self.log_name} Prediction",
image_over_predicted_mask,
global_step=global_step,
dataformats="HWC",
)
def __call__(self, writer: SummaryWriter, tag, sample, idx):
if not isinstance(sample, torch.Tensor):
sample = torch.tensor(sample)
if sample.shape[-1] == 3: # H, W, C
sample = sample.permute(2, 0, 1)
X = self.forward_pass_preprocess(sample).unsqueeze(0)
device = next(self.model.parameters()).device
X = X.to(device)
logits = self.model(X)
res = self.grad_cam.attribute(X, logits.argmax())
res = res.squeeze().detach().cpu().numpy()
res = (res - res.min())
res /= res.max()
res = (res * 255).astype(np.uint8)
writer.add_image(f"{tag}_gradcam", res, global_step=idx)
def img(writer: SummaryWriter, sample: Tuple, prefix: str, task_name: str, key: str):
X, y = sample
writer.add_image(f'{prefix}_{task_name}_images', X[key])
def img_publisher(writer: SummaryWriter, tag: str, sample: Any, idx: int):
writer.add_image(f'{tag}_images', sample, global_step=idx)