Update app.py

app.py

@@ -20,11 +20,7 @@ model = UNet(
 model.load_state_dict(torch.load("best_model.pth", map_location=torch.device('cpu')))
 model.eval()
 
-def greet(image):
-    
-
-    # image = Image.open(image_path).convert("RGB")
-    # image = np.array(image) / 255.0
+def process_image(image):
     image = image / 255.0
     image = image.astype(np.float32)
 
@@ -32,39 +28,34 @@ def greet(image):
         A.Resize(height=512, width=512),
         ToTensorV2(),
     ])
-    print(image.shape)
+    
     image = inference_transforms(image=image)["image"]
-    print(image.shape)
     image = image.unsqueeze(0)
 
-
     with torch.no_grad():
         mask_pred = torch.sigmoid(model(image))
 
-    print(image.shape)
-    print(mask_pred.shape)
-    print(mask_pred[0, 0, :, :].shape)
     return mask_pred[0, 0, :, :].numpy()
     
 
 demo = gr.Interface(
-    fn=greet,
+    fn=process_image,
     title="Histapathology segmentation",
     inputs=[
         gr.Image(
             label="Input image",
             image_mode="RGB",
-            # height=400,
+            height=400,
             type="numpy",
-            # width=400,
+            width=400,
         )
     ],
     outputs=[
         gr.Image(
             label="Model Prediction",
             image_mode="L",
-            # height=400,
-            # width=400,
+            height=400,
+            width=400,
         )
     ],
     # examples=[