Update app.py

app.py

@@ -1,8 +1,28 @@
 import gradio as gr
+import torch
+from UNET_perso import UNET
+import matplotlib.pyplot as plt
+from src.medicalDataLoader import *
+
 
 def greet(name):
     return "Hello " + name + "!"
 
+import torchvision
+from torchvision.io import read_image
+  
+pic = read_image('Data/val/Img/patient001_01_1.png')
+print(pic.shape)
+
+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()
+
+plt.imshow(pic)
+plt.show()
+
 demo = gr.Interface(fn=greet, inputs="text", outputs="text")
     
 demo.launch()