Update model.py

model.py

@@ -1,7 +1,8 @@
+## MODEL.py
 import pytorch_lightning as pl
 import segmentation_models_pytorch as smp
 import torch
-
+import torchmetrics
 
 class ModelRoiLeish(pl.LightningModule):
 
@@ -63,6 +64,7 @@ class ModelRoiLeish(pl.LightningModule):
         # first convert mask values to probabilities, then 
         # apply thresholding
         prob_mask = logits_mask.sigmoid()
+        iou_score = torchmetrics.functional.jaccard_index(prob_mask, mask.long())
         pred_mask = (prob_mask > 0.5).float()
 
         # We will compute IoU metric by two ways
@@ -79,6 +81,7 @@ class ModelRoiLeish(pl.LightningModule):
             f"{stage}_fp": fp.to(torch.float32).mean(),
             f"{stage}_fn": fn.to(torch.float32).mean(),
             f"{stage}_tn": tn.to(torch.float32).mean(),
+            f"{stage}_jaccard": iou_score.to(torch.float32).mean()
         }
 
         self.log_dict(loss_metrics, prog_bar=True)
@@ -135,4 +138,5 @@ class ModelRoiLeish(pl.LightningModule):
         return self.shared_epoch_end(outputs, "test")
 
     def configure_optimizers(self):
-        return torch.optim.AdamW(self.parameters(), lr=self.lr)
+        return torch.optim.AdamW(self.parameters(), lr=self.lr)
+