Create app.py

app.py

@@ -0,0 +1,89 @@
+from ultralytics import YOLO
+from PIL import Image
+import gradio as gr
+from huggingface_hub import snapshot_download
+import os
+import cv2
+
+def load_model(repo_id, model_filename="best_int8_openvino_model"):
+    """
+    Loads a YOLO model from Hugging Face Hub.
+
+    Args:
+        repo_id: The ID of the Hugging Face Hub repository.
+        model_filename: The filename of the YOLO model within the repository.
+
+    Returns:
+        The loaded YOLO model.
+    """
+    download_dir = snapshot_download(repo_id)
+    model_path = os.path.join(download_dir, model_filename)
+    detection_model = YOLO(model_path, task='detect')
+    return detection_model
+
+def predict_image(pilimg, conf_thresh, iou_thresh):
+    """
+    Performs object detection on the input image.
+
+    Args:
+        pilimg: The input image as a PIL Image object.
+        conf_thresh: The confidence threshold for object detection.
+        iou_thresh: The IoU threshold for non-maximum suppression.
+
+    Returns:
+        The processed image with detected objects highlighted.
+    """
+    source = pilimg
+    result = detection_model(source, conf=conf_thresh, iou=iou_thresh)
+    img_bgr = result[0].plot()
+    out_pilimg = Image.fromarray(img_bgr[..., ::-1])  # Convert BGR to RGB
+    return out_pilimg
+
+def predict_video(video_path, conf_thresh, iou_thresh):
+    """
+    Performs object detection on a video.
+
+    Args:
+        video_path: Path to the video file.
+        conf_thresh: The confidence threshold for object detection.
+        iou_thresh: The IoU threshold for non-maximum suppression.
+
+    Returns:
+        A list of processed frames with detected objects highlighted.
+    """
+    cap = cv2.VideoCapture(video_path)
+    frame_list = []
+
+    while True:
+        ret, frame = cap.read()
+        if not ret:
+            break
+        pilimg = Image.fromarray(frame)
+        result = detection_model(pilimg, conf=conf_thresh, iou=iou_thresh)
+        img_bgr = result[0].plot()
+        frame_list.append(img_bgr)
+    
+    cap.release()
+    return frame_list
+
+REPO_ID = "qiqiyuan/glasses_and_mouth" 
+detection_model = load_model(REPO_ID)
+
+iface = gr.Interface(
+    fn=[predict_image, predict_video],
+    inputs=[
+        gr.Image(type="pil", label="Image"),
+        gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.5, label="Confidence Threshold"),
+        gr.Slider(minimum=0.0, maximum=1.0, step=0.05, value=0.5, label="IoU Threshold"),
+        gr.File(label="Video (optional)") 
+    ],
+    outputs=[
+        gr.Image(type="pil", label="Image Output"),
+        gr.Video(label="Video Output")
+    ],
+    title="Object Detection with YOLO",
+    description="Upload an image or video to detect glasses and mouth.",
+    examples=[["examples/image1.jpg"], ["examples/video1.mp4"]]  # Add example images/videos
+)
+
+iface.launch(share=True)