Initial Setup

app.py

@@ -1,36 +1,64 @@
 from ultralytics import YOLO
 from PIL import Image
 import gradio as gr
-from huggingface_hub import snapshot_download
 import os
 
-model_path = "/Users/markk/Downloads/best_int8_openvino_model"
+# Define the model path
+model_path = "C:\Users\Ichiya Rei\235008L\best.onnx"
 
-def load_model(repo_id):
-    download_dir = snapshot_download(repo_id)
-    print(download_dir)
-    path  = os.path.join(download_dir, "best_int8_openvino_model")
-    print(path)
-    detection_model = YOLO(path, task='detect')
+
+# Function to load the YOLO model
+def load_model(model_path):
+    print(f"Loading model from: {model_path}")
+    detection_model = YOLO(model_path, task='detect')
     return detection_model
 
+# Load the YOLO model
+detection_model = load_model(model_path)
+
+# Function to perform prediction on an image
+def predict_image(input_img):
+    try:
+        result = detection_model.predict(input_img, conf=0.5, iou=0.6)
+        img_bgr = result[0].plot()
+        out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # Convert BGR to RGB for PIL image
+        return out_pilimg
+    except Exception as e:
+        return f"Error processing image: {str(e)}"
 
-def predict(pilimg):
+# Function to perform prediction on a video
+def predict_video(video_path):
+    try:
+        result_video = detection_model.predict(video_path, conf=0.5, iou=0.6, save=True)
+        
+        # Find the processed video output path
+        processed_video_path = str(result_video[0].save_dir / result_video[0].path.name)
+        return processed_video_path
+    except Exception as e:
+        return f"Error processing video: {str(e)}"
 
-    source = pilimg
-    # x = np.asarray(pilimg)
-    # print(x.shape)
-    result = detection_model.predict(source, conf=0.5, iou=0.6)
-    img_bgr = result[0].plot()
-    out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # RGB-order PIL image
-    
-    return out_pilimg
+# Define Gradio interface for image and video separately
+image_interface = gr.Interface(
+    fn=predict_image, 
+    inputs=gr.Image(type="pil"), 
+    outputs=gr.Image(type="pil"),
+    title="Object Detection - Image",
+    description="Upload an image and the model will detect objects."
+)
 
+video_interface = gr.Interface(
+    fn=predict_video, 
+    inputs=gr.Video(), 
+    outputs=gr.File(label="Processed Video"),
+    title="Object Detection - Video",
+    description="Upload a video and the model will detect objects."
+)
 
-REPO_ID = "khengkok/balloon_yolov8"
-detection_model = load_model(REPO_ID)
+# Combine both into a tabbed interface
+app = gr.TabbedInterface(
+    [image_interface, video_interface], 
+    ["Image Detection", "Video Detection"]
+)
 
-gr.Interface(fn=predict,
-             inputs=gr.Image(type="pil"),
-             outputs=gr.Image(type="pil")
-             ).launch(share=True)
+# Launch the app
+app.launch(share=True)

best.onnx

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:389195fb8fd39f37d3f7ec8cebf0cc63a0e6a225dcac0e80350c77efd8bbfea9
+size 103608880

requirements.txt

@@ -1,2 +1,4 @@
-ultralytics
-huggingface_hub
+ultralytics==8.3.63
+gradio==5.12.0
+pillow==11.1.0
+onnxruntime==1.20.1