From f0fbf7917b942e6bc048106ac45dba61970e822a Mon Sep 17 00:00:00 2001
From: "s.islam" <s.islam@fz-juelich.de>
Date: Wed, 20 Apr 2022 11:50:08 +0200
Subject: [PATCH] Adding image grid in validation loop: To illustrate and
compare pli, cyto, and generated cyto
---
code/model.py | 4 ++++
1 file changed, 4 insertions(+)
diff --git a/code/model.py b/code/model.py
index 1c010fd..32674a3 100644
--- a/code/model.py
+++ b/code/model.py
@@ -2,6 +2,7 @@ import torch.nn as nn
from torch.nn import functional as F
import torch
import pytorch_lightning as pl
+from torchvision.utils import make_grid
import segmentation_models_pytorch as smp
@@ -51,6 +52,9 @@ class TestModule(pl.LightningModule):
cyto_imag_generated = self.forward(batch['pli_image'])
loss = self.loss_f(cyto_imag_generated, batch['cyto_image'])
self.log("val_loss", loss)
+ if batch_idx == 0:
+ grid = make_grid([batch['pli_image'][0, 0], batch['cyto_image'][0, 0], cyto_imag_generated[0, 0]])
+ self.logger.experiment.add_image('Grid_images', grid, self.current_epoch, dataformats="HW")
def configure_optimizers(self):
optimizer = torch.optim.Adam(self.parameters(), lr=self.learning_rate)
--
GitLab