Update app.py

app.py

@@ -1,8 +1,11 @@
-import gradio as gr
+import torch
 from PIL import Image
-import os
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
+import io
+from random import choice
+import gradio as gr
 from yolo import YOLO
-from tqdm import tqdm
 
 yolo = YOLO()
 
@@ -13,19 +16,53 @@ def predict_single_image(image, crop=False, count=True):
     except Exception as e:
         return str(e)
 
-demo = gr.Interface(
-    fn=predict_single_image,
-    inputs=[
-        gr.Image(value="img/1.png", type="pil", label="Input Images"),
-    ],
-	
-    outputs=gr.Image(label="Detection Images", 
-                        type="pil",
-                        show_download_button=True,
-                        width=400, 
-                        height=600
-                        ),
-)
-
-if __name__ == "__main__":
-    demo.launch()
+COLORS = ["#ff7f7f", "#ff7fbf", "#ff7fff", "#bf7fff",
+            "#7f7fff", "#7fbfff", "#7fffff", "#7fffbf",
+            "#7fff7f", "#bfff7f", "#ffff7f", "#ffbf7f"]
+
+fdic = {
+    "family" : "DejaVu Serif",
+    "style" : "normal",
+    "size" : 18,
+    "color" : "yellow",
+    "weight" : "bold"
+}
+
+
+def get_figure(in_pil_img, in_results):
+    plt.figure(figsize=(16, 10))
+    plt.imshow(in_pil_img)
+    ax = plt.gca()
+
+    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()
+
+        ax.add_patch(plt.Rectangle((x, y), w, h, fill=False, color=selected_color, linewidth=3, alpha=0.8))
+        ax.text(x, y, 'MASFNet')
+
+    plt.axis("off")
+
+    return plt.gcf()
+
+
+
+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)