Update app.py

app.py

@@ -1,23 +1,63 @@
 import gradio as gr
 import torch
-from PIL import Image
+import yolov7
 
-# Load the model using torch.hub (you must have the model locally)
-model_name = "skin_burn_2022_8_21 (1).pt"
-model =torch.hub.load("WongKinYiu/yolov7", 'custom',model_name)
 
-def predict(img):
-    # Run inference with YOLO model
-    results = model(img)
-    # Convert the result to an image
-    result_img = Image.fromarray(results.render())
-    return result_img
+# Images
+torch.hub.download_url_to_file('https://github.com/Michael-OvO/Burn-Detection-Classification/blob/main/inference/images/1st_degree_1.jpg', '1st_degree_1.jpg')
+torch.hub.download_url_to_file('https://github.com/Michael-OvO/Burn-Detection-Classification/blob/main/inference/images/3rd_degree_1.jpg', '3rd_degree_1.jpg')
+    
+def yolov7_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,
+):
+    """
+    YOLOv7 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
+    """
 
-# Define a Gradio interface
-iface = gr.Interface(
-    fn=predict,  # the function to wrap
-    inputs=gr.inputs.Image(),  # input type
-    outputs=gr.outputs.Image()  # output type
-)
+    model = torch.hub.load('WongKinYiu/yolov7', 'custom', path='path/to/best.pt', source='local', device="cpu", hf_model=True, trace=False)
+    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(
+        choices=[
+            "skin_burn",
+            "kadirnar/yolov7-v0.1",
+        ],
+        default="skin_burn",
+        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"),
+]
 
-iface.launch()
+outputs = gr.outputs.Image(type="filepath", label="Output Image")
+title = "Yolov7: Trainable bag-of-freebies sets new state-of-the-art for real-time object detectors"
+
+demo_app = gr.Interface(
+    fn=yolov7_inference,
+    inputs=inputs,
+    outputs=outputs,
+    title=title,
+    examples=examples,
+    cache_examples=True,
+    theme='huggingface',
+)
+demo_app.launch(debug=True, enable_queue=True)