Update app.py

app.py

@@ -1,5 +1,6 @@
 # Import necessary libraries and load the model
 import gradio as gr
+from gradio.layouts import Column, Row
 import torch
 import numpy as np
 from PIL import Image
@@ -17,6 +18,14 @@ model.load_state_dict(torch.load('unet_model.pth', map_location=device))
 model = model.to(device)
 model.eval()
 
+
+# Define your examples
+examples = [
+    ["examples/image_1.jpg", "examples/gt_1.png", "examples/scribble_1.png"],
+    # ["path_to_image2", "path_to_gt_mask2", "path_to_scribble_mask2"],
+    # ["path_to_image3", "path_to_gt_mask3", "path_to_scribble_mask3"]
+]
+
 # Define the segmentation function
 def segment(img):
     img = Image.fromarray(img.astype('uint8'), 'RGB')
@@ -35,7 +44,7 @@ def segment(img):
     output = torch.argmax(torch.softmax(output, dim=1), dim=1).squeeze().cpu().numpy()
     
     # Resize the mask back to the original image size
-    output = Image.fromarray(output.astype('uint8')).resize(original_size, resample=Image.BILINEAR)
+    output = Image.fromarray(output.astype('uint8')).resize(original_size, resample=Image.NEAREST)
     
     # Convert the PIL Image back to a numpy array
     output = np.array(output)
@@ -44,10 +53,18 @@ def segment(img):
     
     return binary_mask
 
+
+gr.Button('hhhhh')
 # Create a Gradio interface
-iface = gr.Interface(fn=segment, inputs="image", outputs="image", title="Segmentation Model",
-    description="Segment objects in an image.",
-    allow_flagging=False)
+demo = gr.Interface(fn=segment, 
+                    inputs="image",
+                    outputs="image",
+                    title="<p>S<sup>2</sup>ME: Spatial-Spectral Mutual Teaching and Ensemble Learning</p><p>for Scribble-supervised Polyp Segmentation</p>",
+                    description="<p>MICCAI 2023, the 26th International Conference on Medical Image Computing and Computer Assisted Intervention</p><p>An Wang, Mengya Xu, Yang Zhang, Mobarakol Islam, and Hongliang Ren</p>",
+                    allow_flagging=False,
+                    examples=examples)  # Add your examples here
+
 
 # Launch the interface
-iface.launch()
+demo.launch()
+