From b701afe363330b8d7fd2d966d999196e279b283e Mon Sep 17 00:00:00 2001 From: Laurent Fainsin Date: Mon, 12 Sep 2022 09:28:11 +0200 Subject: [PATCH] feat: better image logging Former-commit-id: 5ce04a5534cef72a3815be13cd79731800b7419f [formerly 47ab70b94c3f4a696b4e5a131e087660dba0b8ba] Former-commit-id: 81858cf5970cdecdbd0a151e3547a25857e3e958 --- src/utils/callback.py | 79 +++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 76 insertions(+), 3 deletions(-) diff --git a/src/utils/callback.py b/src/utils/callback.py index 37366e3..b00c956 100644 --- a/src/utils/callback.py +++ b/src/utils/callback.py @@ -7,16 +7,54 @@ columns = [ ] class_labels = { 1: "sphere", + 2: "sphere_gt", } class TableLog(Callback): + def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx): + if batch_idx == 0: + rows = [] + + # unpacking + images, targets = batch + + for i, (image, target) in enumerate( + zip( + images, + targets, + ) + ): + rows.append( + [ + i, + wandb.Image( + image.cpu(), + masks={ + "ground_truth": { + "mask_data": (target["masks"].cpu().sum(dim=0) > 0.5).int().numpy() * 2, + "class_labels": class_labels, + }, + }, + ), + ] + ) + + wandb.log( + { + "train/predictions": wandb.Table( + columns=columns, + data=rows, + ) + } + ) + def on_validation_epoch_start(self, trainer, pl_module): self.rows = [] def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx): - # unpacking if batch_idx == 2: + # unpacking images, targets = batch for i, (image, target, pred) in enumerate( @@ -26,6 +64,37 @@ class TableLog(Callback): outputs, ) ): + box_data_gt = [ + { + "position": { + "minX": int(target["boxes"][j][0]), + "minY": int(target["boxes"][j][1]), + "maxX": int(target["boxes"][j][2]), + "maxY": int(target["boxes"][j][3]), + }, + "domain": "pixel", + "class_id": 2, + "class_labels": class_labels, + } + for j in range(len(target["labels"])) + ] + + box_data = [ + { + "position": { + "minX": int(pred["boxes"][j][0]), + "minY": int(pred["boxes"][j][1]), + "maxX": int(pred["boxes"][j][2]), + "maxY": int(pred["boxes"][j][3]), + }, + "domain": "pixel", + "class_id": 1, + "box_caption": f"{pred['scores'][j]:0.3f}", + "class_labels": class_labels, + } + for j in range(len(pred["labels"])) + ] + self.rows.append( [ i, @@ -33,14 +102,18 @@ class TableLog(Callback): image.cpu(), masks={ "ground_truth": { - "mask_data": (target["masks"].cpu().sum(dim=0) > 0.5).int().numpy(), + "mask_data": target["masks"].cpu().sum(dim=0).int().numpy() * 2, "class_labels": class_labels, }, "predictions": { - "mask_data": (pred["masks"].cpu().sum(dim=0) > 0.5).int().numpy(), + "mask_data": pred["masks"].cpu().sum(dim=0).int().numpy(), "class_labels": class_labels, }, }, + boxes={ + "ground_truth": {"box_data": box_data_gt}, + "predictions": {"box_data": box_data}, + }, ), ] )