Update app.py

app.py

@@ -1,7 +1,6 @@
 import torch
 from PIL import Image
 import matplotlib.pyplot as plt
-import matplotlib.patches as patches
 import io
 from random import choice
 import gradio as gr
@@ -9,27 +8,22 @@ from yolo import YOLO
 
 yolo = YOLO()
 
-def predict_single_image(image, crop=False, count=True):
-    try:
-        r_image = yolo.detect_image(image, crop=crop, count=count)
-        return r_image
-    except Exception as e:
-        return str(e)
-
+# Colors for bounding boxes
 COLORS = ["#ff7f7f", "#ff7fbf", "#ff7fff", "#bf7fff",
-            "#7f7fff", "#7fbfff", "#7fffff", "#7fffbf",
-            "#7fff7f", "#bfff7f", "#ffff7f", "#ffbf7f"]
+          "#7f7fff", "#7fbfff", "#7fffff", "#7fffbf",
+          "#7fff7f", "#bfff7f", "#ffff7f", "#ffbf7f"]
 
+# Font dictionary for text annotations
 fdic = {
-    "family" : "DejaVu Serif",
-    "style" : "normal",
-    "size" : 18,
-    "color" : "yellow",
-    "weight" : "bold"
+    "family": "DejaVu Serif",
+    "style": "normal",
+    "size": 18,
+    "color": "yellow",
+    "weight": "bold"
 }
 
-
 def get_figure(in_pil_img, in_results):
+    """ Function to generate figure with bounding boxes and labels """
     plt.figure(figsize=(16, 10))
     plt.imshow(in_pil_img)
     ax = plt.gca()
@@ -37,32 +31,29 @@ def get_figure(in_pil_img, in_results):
     for score, label, box in zip(in_results["scores"], in_results["labels"], in_results["boxes"]):
         selected_color = choice(COLORS)
 
-        box_int = [i.item() for i in torch.round(box).to(torch.int32)]
-        x, y, w, h = box_int[0], box_int[1], box_int[2]-box_int[0], box_int[3]-box_int[1]
-        #x, y, w, h = torch.round(box[0]).item(), torch.round(box[1]).item(), torch.round(box[2]-box[0]).item(), torch.round(box[3]-box[1]).item()
+        box_int = [int(i.item()) for i in torch.round(box)]
+        x, y, w, h = box_int[0], box_int[1], box_int[2] - box_int[0], box_int[3] - box_int[1]
 
         ax.add_patch(plt.Rectangle((x, y), w, h, fill=False, color=selected_color, linewidth=3, alpha=0.8))
-        ax.text(x, y, 'MASFNet')
+        ax.text(x, y)
 
     plt.axis("off")
-
     return plt.gcf()
 
+def predict(image):
+    try:
+        r_image = yolo.detect_image(image)
+        return r_image
+    except Exception as e:
+        return str(e)
 
-
-with gr.Blocks(title="MASFNet Object Detection",
-                    css=".gradio-container {background:lightyellow;}"
-               ) as demo:
-    #sample_index = gr.State([])
-
-
-    with gr.Row():
-        input_image = gr.Image(label="Input image", type="pil")
-        output_image = gr.Image(label="Output image with predicted instances", type="pil")
-
-    gr.Examples(['img/1.png', 'img/2.png'], inputs=input_image)
-
-    send_btn = gr.Button("Predict")
-
-#demo.queue()
-demo.launch(debug=True)
+# Define the Gradio interface
+gr.Interface(
+    fn=predict,
+    inputs=gr.inputs.Image(label="Input image", type="pil"),
+    outputs=gr.outputs.Image(label="Output image with predicted instances", type="pil"),
+    examples=['img/1.png', 'img/2.png'],
+    title="MASFNet Object Detection",
+    theme="default",
+    debug=True
+).launch()