Update app.py

app.py

@@ -1,15 +1,57 @@
 import gradio as gr
 import monai
 import torch
+from monai.networks.nets import UNet
+from PIL import Image
+import albumentations as A
+from albumentations.pytorch import ToTensorV2
+import numpy as np
 
-def greet(name, intensity):
-    return "Hello "
+
+model = UNet(
+    spatial_dims=2,
+    in_channels=3,
+    out_channels=1,
+    channels=[16, 32, 64, 128, 256, 512],
+    strides=(2, 2, 2, 2, 2),
+    num_res_units=4,
+    dropout=0.15,
+)
+model.load_state_dict(torch.load("best_model.pth"))
+model.eval()
+
+def greet(image_path):
+    
+
+    image = Image.open(image_path).convert("RGB")
+    image = np.array(image) / 255.0
+    image = image.astype(np.float32)
+
+    inference_transforms = A.Compose([
+        A.Resize(height=512, width=512),
+        ToTensorV2(),
+    ])
+
+    image = inference_transforms(image=image)["image"]
+
+    image = image.unsqueeze(0)
+
+
+    with torch.no_grad():
+        mask_pred = model(image)
+        
+    return mask_pred[0]
+    
 
 demo = gr.Interface(
     fn=greet,
-    title="test-title",
-    inputs=["file"],
-    outputs=["text"],
+    title="Histapathology segmentation",
+    inputs=[
+        gr.File(label="Input image (512x512)")
+    ],
+    outputs=[
+        gr.File(label="Model Prediction")
+    ],
 )
 
 demo.launch()