Update app.py

app.py

@@ -1,44 +1,53 @@
 import gradio as gr
-from huggingface_hub import hf_hub_download
+import torch
 from ultralytics import YOLO
-from supervision import Detections
-from PIL import Image, ImageDraw
+from PIL import Image
 
-# Download & load the YOLOv8 face detection model
-model_path = hf_hub_download(
-    repo_id="arnabdhar/YOLOv8-Face-Detection",
-    filename="model.pt"
-)
-model = YOLO(model_path)
-
-def count_faces(image: Image.Image):
+def load_model():
     """
-    Detects faces in an image and returns the annotated image + face count.
+    Load the YOLOv8 segmentation model onto GPU (if available)
+    with mixed‑precision enabled.
     """
-    # 1. Run detection
-    results = model(image)[0]
-    dets = Detections.from_ultralytics(results)
+    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
+
+model, device = load_model()
 
-    # 2. Draw boxes
-    annotated = image.copy()
-    draw = ImageDraw.Draw(annotated)
-    for x1, y1, x2, y2 in dets.xyxy:
-        draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
+def count_persons(image: Image.Image) -> str:
+    """
+    Run inference on the input image, apply TTA, filter for class 0 (person),
+    and return the total count.
+    """
+    # 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}
+    )
 
-    # 3. Return
-    return annotated, f"Faces detected: {len(dets.xyxy)}"
+    # Sum counts across all results (usually one per image)
+    total = sum(len(r.boxes) for r in results)
+    return f"Persons detected: {total}"
 
-# Gradio interface
-image_app = gr.Interface(
-    fn=count_faces,
+# Build Gradio interface
+demo = gr.Interface(
+    fn=count_persons,
     inputs=gr.Image(type="pil", label="Upload Image"),
-    outputs=[
-        gr.Image(type="pil", label="Annotated Image"),
-        gr.Text(label="Face Count")
-    ],
-    title="Image Face Counter",
-    description="Detect and count faces in a single image using YOLOv8."
+    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"],
+    ]
 )
 
 if __name__ == "__main__":
-    image_app.launch()
+    demo.launch()  # Launch locally; add `share=True` for a public link