Update app.py

app.py

@@ -1,98 +1,54 @@
-import os
-import tempfile
-import cv2
-
-from PIL import Image, ImageDraw
 import gradio as gr
+from huggingface_hub import hf_hub_download
 from ultralytics import YOLO
 from supervision import Detections
+from PIL import Image, ImageDraw
 
+# Download the YOLOv8 face detection model from Hugging Face Hub
+model_path = hf_hub_download(repo_id="arnabdhar/YOLOv8-Face-Detection", filename="model.pt")
 
-# Download and load the YOLOv8 face detection model
-def load_model():
-    model = YOLO("yolov8n-face.pt")  # Make sure the path or name of the model is correct
-    return model
-
-
-model = load_model()
-
+# Load the YOLOv8 face detection model
+model = YOLO(model_path)
 
 def detect_faces(image: Image.Image):
+    """
+    Detects faces in an input image using YOLOv8 face detection model and returns the annotated image
+    along with the number of faces detected.
+    """
+    # Run inference on the input image
     output = model(image)
+    
+    # Convert YOLO output to Detections format using the supervision library
     results = Detections.from_ultralytics(output[0])
-    boxes = results.xyxy
-
-    annotated = image.copy()
-    draw = ImageDraw.Draw(annotated)
-    for x1, y1, x2, y2 in boxes:
+    
+    # Extract bounding boxes; results.xyxy contains boxes in [x1, y1, x2, y2] format
+    boxes = results.xyxy  # This is assumed to be a list-like structure of bounding boxes
+    
+    # Create a copy of the input image for drawing
+    annotated_image = image.copy()
+    draw = ImageDraw.Draw(annotated_image)
+    
+    # Draw a red bounding box for each detected face
+    for box in boxes:
+        x1, y1, x2, y2 = box  # unpack the coordinates
         draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
-
-    return annotated, f"Number of faces detected: {len(boxes)}"
-
-
-def detect_faces_video(video_path: str):
-    """
-    Reads a video file, annotates faces on each frame, and writes out an annotated video.
-
-    Returns the new video path and a summary.
-    """
-    cap = cv2.VideoCapture(video_path)
-    fps = cap.get(cv2.CAP_PROP_FPS)
-    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-
-    # Prepare output
-    out_file = tempfile.mktemp(suffix=".mp4")
-    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
-    writer = cv2.VideoWriter(out_file, fourcc, fps, (width, height))
-
-    frame_count = 0
-    total_faces = 0
-
-    while True:
-        ret, frame = cap.read()
-        if not ret:
-            break
-
-        # Convert frame BGR -> RGB and to PIL Image for model
-        pil_img = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
-        output = model(pil_img)
-        results = Detections.from_ultralytics(output[0])
-        boxes = results.xyxy
-
-        # Draw boxes on original frame
-        for x1, y1, x2, y2 in boxes:
-            cv2.rectangle(frame, (int(x1), int(y1)), (int(x2), int(y2)), (0, 0, 255), 2)
-
-        writer.write(frame)
-        frame_count += 1
-        total_faces += len(boxes)
-
-    cap.release()
-    writer.release()
-
-    avg_per_frame = total_faces / frame_count if frame_count else 0
-    summary = (
-        f"Processed {frame_count} frames. "
-        f"Total faces detected: {total_faces}. "
-        f"Average per frame: {avg_per_frame:.2f}"
-    )
-    return out_file, summary
-
-
-# Build Gradio interface
-video_interface = gr.Interface(
-    fn=detect_faces_video,
-    inputs=gr.Video(label="Upload Video"),
-    outputs=[gr.Video(label="Annotated Video"), gr.Text(label="Summary")],
-    title="YOLOv8 Video Face Detector",
-    description="Detect and annotate faces in videos using a YOLOv8 model."
+    
+    # Count the number of faces detected
+    face_count = len(boxes)
+    
+    return annotated_image, f"Number of faces detected: {face_count}"
+
+# Create a Gradio interface for the face detection function
+demo = gr.Interface(
+    fn=detect_faces,
+    inputs=gr.Image(type="pil", label="Upload Image"),
+    outputs=[
+        gr.Image(type="pil", label="Annotated Image"),
+        gr.Text(label="Face Count")
+    ],
+    title="YOLOv8 Face Detector",
+    description="Upload an image to detect faces using a YOLOv8 face detection model. The detected faces will be highlighted with red bounding boxes."
 )
 
-
-def main():
-    video_interface.launch()
-
-
 if __name__ == "__main__":
-    main()
+    demo.launch()