Update app.py

app.py

@@ -1,55 +1,54 @@
-# app.py
-
 import gradio as gr
-import torch
+from huggingface_hub import hf_hub_download
 from ultralytics import YOLO
-from PIL import Image
+from supervision import Detections
+from PIL import Image, ImageDraw
 
-def load_model():
-    """
-    Load the YOLOv8 segmentation model onto GPU (if available)
-    with mixed‑precision enabled.
-    """
-    device = "cuda" if torch.cuda.is_available() else "cpu"        # GPU if available :contentReference[oaicite:2]{index=2}
-    model = YOLO('yolov8x-seg.pt').to(device)                     # Segmentation variant for finer masks :contentReference[oaicite:3]{index=3}
-    return model, device
+# Download the YOLOv8 face detection model from Hugging Face Hub
+model_path = hf_hub_download(repo_id="arnabdhar/YOLOv8-Face-Detection", filename="model.pt")
 
-model, device = load_model()
+# Load the YOLOv8 face detection model
+model = YOLO(model_path)
 
-def count_persons(image: Image.Image) -> str:
+def detect_faces(image: Image.Image):
     """
-    Run inference on the input image, apply TTA, filter for class 0 (person),
-    and return the total count.
+    Detects faces in an input image using YOLOv8 face detection model and returns the annotated image
+    along with the number of faces detected.
     """
-    # Perform prediction with augmentation (TTA), limit detections, and only class 0
-    results = model.predict(
-        source=image,
-        conf=0.6,          # Confidence threshold
-        imgsz=640,         # Inference resolution
-        augment=True,      # Test Time Augmentation :contentReference[oaicite:4]{index=4}
-        max_det=300,       # Cap detections for crowded scenes
-        classes=[0]        # Only detect persons (class 0) :contentReference[oaicite:5]{index=5}
-    )
-
-    # Sum counts across all results (usually one per image)
-    total = sum(len(r.boxes) for r in results)
-    return f"Persons detected: {total}"
+    # 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])
+    
+    # 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)
+    
+    # Count the number of faces detected
+    face_count = len(boxes)
+    
+    return annotated_image, f"Number of faces detected: {face_count}"
 
-# Build Gradio interface
+# Create a Gradio interface for the face detection function
 demo = gr.Interface(
-    fn=count_persons,
+    fn=detect_faces,
     inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs=gr.Text(label="Person Count"),
-    title="Advanced Person Counter with YOLOv8",
-    description=(
-        "Upload an image to count people using a state‑of‑the‑art "
-        "YOLOv8 segmentation model with Test‑Time Augmentation."
-    ),
-    examples=[  # optional: add example images if you like
-        # ["examples/crowd1.jpg"],
-        # ["examples/street_scene.jpg"],
-    ]
+    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."
 )
 
 if __name__ == "__main__":
-    demo.launch()  # Launch locally; add `share=True` for a public link
+    demo.launch()