|
|
from lightning.pytorch.callbacks import Callback
|
|
|
from pytorch_lightning.utilities import rank_zero_only
|
|
|
import wandb
|
|
|
|
|
|
|
|
|
class LogPredictionSamplesCallback(Callback):
|
|
|
def __init__(self, wandb_logger, n_images=8):
|
|
|
self.n_images = n_images
|
|
|
self.wandb_logger = wandb_logger
|
|
|
super().__init__()
|
|
|
|
|
|
@rank_zero_only
|
|
|
def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx):
|
|
|
if batch_idx < 1 and trainer.is_global_zero:
|
|
|
n = self.n_images
|
|
|
x = batch["image"][:n].float()
|
|
|
h, w = x.shape[-2:]
|
|
|
y = batch["mask"][:n]
|
|
|
pred = outputs[:n]
|
|
|
pred = pred.argmax(dim=1)
|
|
|
|
|
|
if len(y.shape) == 4:
|
|
|
y = y.squeeze(1)
|
|
|
if len(pred.shape) == 4:
|
|
|
pred = pred.squeeze(1)
|
|
|
y = y.clamp(0, 2)
|
|
|
columns = ["image"]
|
|
|
class_labels = {0: "Background", 1: "Red", 2: "Green"}
|
|
|
|
|
|
data = [
|
|
|
[
|
|
|
wandb.Image(
|
|
|
x_i,
|
|
|
masks={
|
|
|
"Prediction": {
|
|
|
"mask_data": p_i.cpu().numpy(),
|
|
|
"class_labels": class_labels,
|
|
|
},
|
|
|
"Groundtruth": {
|
|
|
"mask_data": y_i.cpu().numpy(),
|
|
|
"class_labels": class_labels,
|
|
|
},
|
|
|
},
|
|
|
)
|
|
|
]
|
|
|
for x_i, y_i, p_i in list(zip(x, y, pred))
|
|
|
]
|
|
|
self.wandb_logger.log_table(
|
|
|
data=data, key=f"Validation Batch {batch_idx}", columns=columns
|
|
|
)
|
|
|
|
