add IOU and Dice Score

app.py

@@ -9,13 +9,12 @@ import einops
 import torch
 from torch.utils.data import DataLoader
 from torchvision import transforms
-from torchmetrics.functional import dice
+from torchmetrics.functional import dice_score, jaccard_index
 
 from UNET_perso import UNET
 from src.medicalDataLoader import MedicalImageDataset
 from src.utils import getTargetSegmentation
 
-
 def disp_init(num):
     lst_patients = glob.glob('./Data/val/Img/*.png')
     patient = lst_patients[num]
@@ -25,7 +24,8 @@ def disp_init(num):
     plt.axis('off')
     return fig
 
-def disp_prediction(num, alpha=0.4):
+def disp_prediction(num, alpha=0.45):
+
     lst_patients = glob.glob('./Data/val/Img/*.png')
     patient = lst_patients[num]
 
@@ -43,7 +43,7 @@ def disp_prediction(num, alpha=0.4):
                                 equalize=False)
     val_loader = DataLoader(val_set,
                             batch_size=1,
-                            shuffle=True)
+                            shuffle=False)
     
     for _, (img, _, name) in enumerate(val_loader):
         if name[0] == patient:
@@ -60,7 +60,7 @@ def disp_prediction(num, alpha=0.4):
     
     fig = plt.figure()
     try:
-        plt.imshow(alpha*pred_mask + (1-alpha)*im)
+        plt.imshow((alpha*pred_mask + (1-alpha)*im)/2)
     except ValueError:
         print(type(pred_mask))
     plt.axis('off')
@@ -100,15 +100,14 @@ def compute_metrics(num):
                                 transform=transform,
                                 mask_transform=transform,
                                 equalize=False)
-    
     val_loader = DataLoader(val_set,
                             batch_size=1,
-                            shuffle=True)
+                            shuffle=False)
     
     for _, (img, label, name) in enumerate(val_loader):
         if name[0] == patient:
             im = img
-            label = label.reshape(256,256)
+            lab = label.reshape(256,256)
 
     pred = model(im)
 
@@ -116,30 +115,29 @@ def compute_metrics(num):
     pred = pred.reshape(4,256,256)
     pred = np.argmax(pred, axis=0)
     pred = torch.from_numpy(pred)
-    pred = torch.nn.functional.one_hot(pred)
+    pred = torch.nn.functional.one_hot(pred, num_classes=4)
     pred = einops.rearrange(pred, 'h w class -> class h w')
 
-    label = getTargetSegmentation(label)
-    label = torch.nn.functional.one_hot(label, num_classes=4)
-    label = einops.rearrange(label, 'h w class -> class h w')
-
-    print(pred.shape, label.shape)
-    #dsc = dice(pred, label, average='none', num_classes=4)
-
-    print(torch.mean(pred.float()), torch.max(pred.float()))
-    print(torch.mean(label.float()), torch.max(label.float()))
+    lab = getTargetSegmentation(lab)
+    lab = torch.nn.functional.one_hot(lab, num_classes=4)
+    lab = einops.rearrange(lab, 'h w class -> class h w')
+    
+    dsc1 = dice_score(pred[1], lab[1])
+    dsc2 = dice_score(pred[2], lab[2])
+    dsc3 = dice_score(pred[3], lab[3])
 
-    #print(dsc)
-    #dsc1 = dice_score(pred[1], label[1])
-    #dsc2 = dice_score(pred[2], label[2])
-    #dsc3 = dice_score(pred[3], label[3])
-    #print(dsc1, dsc2, dsc3)
+    iou1 = jaccard_index(pred[1], lab[1], num_classes=2)
+    iou2 = jaccard_index(pred[2], lab[2], num_classes=2)
+    iou3 = jaccard_index(pred[3], lab[3], num_classes=2)
 
     df = pd.DataFrame(columns=['class 1', 'class 2', 'class 3'])
-
+    df.loc[len(df)] = [round(dsc1.item(),2), round(dsc2.item(),2), round(dsc3.item(),2)]
+    df.loc[len(df)] = [round(iou1.item(),2), round(iou2.item(),2), round(iou3.item(),2)]
+    df = df.assign(metric=['Dice Score', 'IoU'])
+    df = df[['metric','class 1','class 2','class 3']]
+    
     return df
 
-
 with gr.Blocks() as demo:
     with gr.Row() as row1:
         slide = gr.Slider(minimum=0, maximum=89, value=0, step=1, label='inference with validation split data (choose patient)')