Update app.py

app.py

@@ -1,44 +1,57 @@
-import glob
-
 import gradio as gr
+import torch
 import yolov5
-from PIL import Image
-# import torch
-from huggingface_hub import hf_hub_download
 
-# Add new models here
-model_names = [
-    "TW_plate_number",
+# Images
+torch.hub.download_url_to_file('https://github.com/ultralytics/yolov5/raw/master/data/images/zidane.jpg', 'zidane.jpg')
+torch.hub.download_url_to_file('https://raw.githubusercontent.com/WongKinYiu/yolov7/main/inference/images/image3.jpg', 'image3.jpg')
+
+def yolov5_inference(
+    image: gr.inputs.Image = None,
+    model_path: gr.inputs.Dropdown = None,
+    image_size: gr.inputs.Slider = 640,
+    conf_threshold: gr.inputs.Slider = 0.25,
+    iou_threshold: gr.inputs.Slider = 0.45,
+):
+    """
+    YOLOv5 inference function
+    Args:
+        image: Input image
+        model_path: Path to the model
+        image_size: Image size
+        conf_threshold: Confidence threshold
+        iou_threshold: IOU threshold
+    Returns:
+        Rendered image
+    """
+    model = yolov5.load(model_path, device="cpu")
+    model.conf = conf_threshold
+    model.iou = iou_threshold
+    results = model([image], size=image_size)
+    return results.render()[0]
+        
+
+inputs = [
+    gr.inputs.Image(type="pil", label="Input Image"),
+    gr.inputs.Dropdown(["plate.pt", "yolov5m.pt", "yolov5l.pt", "yolov5x.pt"], label="Model"),
+    gr.inputs.Slider(minimum=320, maximum=1280, default=640, step=32, label="Image Size"),
+    gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.25, step=0.05, label="Confidence Threshold"),
+    gr.inputs.Slider(minimum=0.0, maximum=1.0, default=0.45, step=0.05, label="IOU Threshold"),
 ]
 
-def load_model(model_name):
-    model_path = hf_hub_download(repo_id=model_name, filename="plate.pt")
-    return model_path
-
-def yolo_inference(im, model_path, threshold=0.70):
-    model_loaded = load_model(model_path)
-    model = yolov5.load(model_loaded)
-    model.conf = threshold
-    results = model(im)  # inference
-    numpy_image = results.render()[0]
-    output_image = Image.fromarray(numpy_image)
-    return output_image
-
-title = "TW_plate_number"
-
-inputs = [gr.Image(type="pil", label="document image"),
-          gr.inputs.Dropdown(model_names, label="Model", default=model_names[0]),
-          gr.Slider(maximum=1, step=0.01, value=0.70)]
-
-
-examples = [[str(file),model_names[0], 0.70] for file in glob.glob("./sample/*.jpg")]
-
-demo=gr.Interface(fn=yolo_inference,
-                  inputs=inputs,
-                  outputs=gr.Image(type="pil", label="annotated document").style(height=800),
-                  title=title,
-                  theme="huggingface",
-                  examples=examples)
-
-if __name__ == "__main__":
-    demo.launch(debug=True)
+outputs = gr.outputs.Image(type="filepath", label="Output Image")
+title = "YOLOv5"
+description = "YOLOv5 is a family of object detection models pretrained on COCO dataset. This model is a pip implementation of the original YOLOv5 model."
+
+examples = [['zidane.jpg', 'yolov5s.pt', 640, 0.25, 0.45], ['image3.jpg', 'yolov5s.pt', 640, 0.25, 0.45]]
+demo_app = gr.Interface(
+    fn=yolov5_inference,
+    inputs=inputs,
+    outputs=outputs,
+    title=title,
+    examples=examples,
+    cache_examples=True,
+    live=True,
+    theme='huggingface',
+)
+demo_app.launch(debug=True, enable_queue=True)