Update app.py

app.py

@@ -2,83 +2,67 @@ from ultralytics import YOLO
 from PIL import Image
 import gradio as gr
 from huggingface_hub import snapshot_download
+import tempfile
 import os
 import cv2
-import tempfile
-import numpy as np
 
-# === Load model from Hugging Face ===
+# Load the YOLO model from Hugging Face
 def load_model(repo_id):
     download_dir = snapshot_download(repo_id)
-    print("Model downloaded to:", download_dir)
     model_path = os.path.join(download_dir, "best.pt")
     return YOLO(model_path)
 
-# === Prediction functions ===
-def predict_image(pil_image, conf):
-    result = detection_model.predict(pil_image, conf=conf, iou=0.6)
+# Process image input
+def predict_image(image, conf_threshold, iou_threshold):
+    result = detection_model.predict(image, conf=conf_threshold, iou=iou_threshold)
     img_bgr = result[0].plot()
-    output = Image.fromarray(result[0].plot())
-    return output
-
-def predict_video(video_file, conf):
-    cap = cv2.VideoCapture(video_file)
-    if not cap.isOpened():
-        raise IOError("Cannot open video file")
+    return Image.fromarray(img_bgr[..., ::-1])
 
+# Process video input
+def predict_video(video_path, conf_threshold, iou_threshold):
+    cap = cv2.VideoCapture(video_path)
     width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
     height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
     fps = cap.get(cv2.CAP_PROP_FPS)
-
-    temp_output = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False)
-    out = cv2.VideoWriter(temp_output.name, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
+    out_path = tempfile.NamedTemporaryFile(suffix=".mp4", delete=False).name
+    out_writer = cv2.VideoWriter(out_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
 
     while cap.isOpened():
         ret, frame = cap.read()
         if not ret:
             break
-
-        result = detection_model.predict(frame, conf=conf, iou=0.6)
+        result = detection_model.predict(frame, conf=conf_threshold, iou=iou_threshold)
         annotated = result[0].plot()
-        out.write(annotated)
+        out_writer.write(annotated)
 
     cap.release()
-    out.release()
-    return temp_output.name
+    out_writer.release()
+    return out_path
 
-# === Load model ===
+# Load model
 REPO_ID = "Cedri/battery_key_yolov8"
 detection_model = load_model(REPO_ID)
 
-# === Gradio UI ===
-def image_interface(image, conf_threshold):
-    return predict_image(image, conf=conf_threshold)
-
-def video_interface(video, conf_threshold):
-    return predict_video(video, conf=conf_threshold)
-
-image_tab = gr.Interface(
-    fn=image_interface,
-    inputs=[
-        gr.Image(label="Upload Image"),
-        gr.Slider(minimum=0.1, maximum=1.0, step=0.05, value=0.7, label="Confidence Threshold")
-    ],
-    outputs=gr.Image(label="Detected Image"),
-    title="Battery Key Detection (Image)"
-)
+# Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("## Battery Key Detection - Image & Video")
+    with gr.Tabs():
+        with gr.TabItem("Image"):
+            with gr.Row():
+                img_input = gr.Image(type="pil", label="Upload Image")
+                img_output = gr.Image(type="pil", label="Predicted Image")
+            conf_slider_img = gr.Slider(0.1, 1.0, 0.5, step=0.05, label="Confidence Threshold")
+            iou_slider_img = gr.Slider(0.1, 1.0, 0.6, step=0.05, label="IoU Threshold")
+            run_btn_img = gr.Button("Run Detection on Image")
+            run_btn_img.click(fn=predict_image, inputs=[img_input, conf_slider_img, iou_slider_img], outputs=img_output)
 
-video_tab = gr.Interface(
-    fn=video_interface,
-    inputs=[
-        gr.Video(label="Upload Video"),
-        gr.Slider(minimum=0.1, maximum=1.0, step=0.05, value=0.7, label="Confidence Threshold")
-    ],
-    outputs=gr.Video(label="Detected Video"),
-    title="Battery Key Detection (Video)"
-)
+        with gr.TabItem("Video"):
+            with gr.Row():
+                vid_input = gr.Video(label="Upload Video")
+                vid_output = gr.Video(label="Predicted Video")
+            conf_slider_vid = gr.Slider(0.1, 1.0, 0.5, step=0.05, label="Confidence Threshold")
+            iou_slider_vid = gr.Slider(0.1, 1.0, 0.6, step=0.05, label="IoU Threshold")
+            run_btn_vid = gr.Button("Run Detection on Video")
+            run_btn_vid.click(fn=predict_video, inputs=[vid_input, conf_slider_vid, iou_slider_vid], outputs=vid_output)
 
-# === Launch with tabs ===
-gr.TabbedInterface(
-    [image_tab, video_tab],
-    tab_names=["Image", "Video"]
-).launch()
+demo.launch()