Update app.py

app.py

@@ -1,18 +1,35 @@
 import glob
 import gradio as gr
 import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+import imageio.v3 as iio
+import einops
+
 import torch
 from torch.utils.data import DataLoader
 from torchvision import transforms
-import imageio.v3 as iio
+from torchmetrics.functional import dice, dice_score
 
 from UNET_perso import UNET
-import matplotlib.pyplot as plt
 from src.medicalDataLoader import MedicalImageDataset
+from src.utils import getTargetSegmentation
 
+import sys
+from IPython.core.ultratb import ColorTB
 
-def disp_prediction(num, alpha=0.4):
+sys.excepthook = ColorTB()
+
+def disp_init(num):
+    lst_patients = glob.glob('./Data/val/Img/*.png')
+    patient = lst_patients[num]
+    im = iio.imread(patient)
+    fig =  plt.figure()
+    plt.imshow(im)
+    plt.axis('off')
+    return fig
 
+def disp_prediction(num, alpha=0.4):
     lst_patients = glob.glob('./Data/val/Img/*.png')
     patient = lst_patients[num]
 
@@ -28,12 +45,11 @@ def disp_prediction(num, alpha=0.4):
                                 transform=transform,
                                 mask_transform=transform,
                                 equalize=False)
-    
     val_loader = DataLoader(val_set,
                             batch_size=1,
                             shuffle=True)
     
-    for _, (img, label, name) in enumerate(val_loader):
+    for _, (img, _, name) in enumerate(val_loader):
         if name[0] == patient:
             im = img
 
@@ -43,31 +59,100 @@ def disp_prediction(num, alpha=0.4):
     pred = pred.reshape(4,256,256)
     pred_mask = np.argmax(pred, axis=0)
 
-    fig = plt.figure()
     im = im.detach().numpy()
     im = im.reshape(256,256)
-    plt.imshow(alpha*pred_mask + (1-alpha)*im)
+    
+    fig = plt.figure()
+    try:
+        plt.imshow(alpha*pred_mask + (1-alpha)*im)
+    except ValueError:
+        print(type(pred_mask))
     plt.axis('off')
 
     return fig
 
-def disp_init(num):
-    lst_patients = glob.glob('./Data/val/Img/*.png')
-    patient = lst_patients[num]
-    print(patient)
-    im = iio.imread(patient)
-    fig =  plt.figure()
-    plt.imshow(im)
+def disp_label(num, alpha=0.4):
+    lst_GT = glob.glob('./Data/val/GT/*.png')
+    GT = lst_GT[num]
+
+    label = iio.imread(GT)
+    
+    fig = plt.figure()
+    try : 
+        plt.imshow(label)
+    except ValueError:
+        print(type(label))
     plt.axis('off')
+
     return fig
 
+
+def compute_metrics(num):
+
+    lst_patients = glob.glob('./Data/val/Img/*.png')
+    patient = lst_patients[num]
+
+    model = UNET(in_channels=1, out_channels=4).to('cpu')
+
+    filepath = 'UNET_perso'
+    model.load_state_dict(torch.load(filepath, map_location=torch.device('cpu')))
+    model.eval()
+
+    transform = transforms.Compose([transforms.ToTensor()])
+    val_set = MedicalImageDataset('val',
+                                './Data',
+                                transform=transform,
+                                mask_transform=transform,
+                                equalize=False)
+    
+    val_loader = DataLoader(val_set,
+                            batch_size=1,
+                            shuffle=True)
+    
+    for _, (img, label, name) in enumerate(val_loader):
+        if name[0] == patient:
+            im = img
+            label = label.reshape(256,256)
+
+    pred = model(im)
+
+    pred = pred.detach().numpy()
+    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 = 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()))
+
+    #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)
+
+    df = pd.DataFrame(columns=['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)')
     with gr.Row() as row2:
         slide.release(fn=disp_init, inputs=[slide], outputs=gr.Plot(label='initial image'))
         slide.release(fn=disp_prediction, inputs=[slide], outputs=gr.Plot(label='prediction'))
+        slide.release(fn=disp_label, inputs=[slide], outputs=gr.Plot(label='groundtruth'))
     with gr.Row() as row3:
-        gr.DataFrame()
+        slide.release(fn=compute_metrics, inputs=[slide], outputs=gr.DataFrame())
 
+demo.queue()
 demo.launch()