Spaces:
Build error
Build error
Update app.py
Browse files
app.py
CHANGED
@@ -4,51 +4,41 @@ from ultralytics import YOLO
|
|
4 |
from supervision import Detections
|
5 |
from PIL import Image, ImageDraw
|
6 |
|
7 |
-
# Download the YOLOv8 face detection model
|
8 |
-
model_path = hf_hub_download(
|
9 |
-
|
10 |
-
|
|
|
11 |
model = YOLO(model_path)
|
12 |
|
13 |
-
def
|
14 |
"""
|
15 |
-
Detects faces in an
|
16 |
-
along with the number of faces detected.
|
17 |
"""
|
18 |
-
# Run
|
19 |
-
|
20 |
-
|
21 |
-
|
22 |
-
|
23 |
-
|
24 |
-
|
25 |
-
|
26 |
-
|
27 |
-
# Create a copy of the input image for drawing
|
28 |
-
annotated_image = image.copy()
|
29 |
-
draw = ImageDraw.Draw(annotated_image)
|
30 |
-
|
31 |
-
# Draw a red bounding box for each detected face
|
32 |
-
for box in boxes:
|
33 |
-
x1, y1, x2, y2 = box # unpack the coordinates
|
34 |
draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
|
35 |
-
|
36 |
-
# Count the number of faces detected
|
37 |
-
face_count = len(boxes)
|
38 |
-
|
39 |
-
return annotated_image, f"Number of faces detected: {face_count}"
|
40 |
|
41 |
-
#
|
42 |
-
|
43 |
-
|
|
|
|
|
|
|
44 |
inputs=gr.Image(type="pil", label="Upload Image"),
|
45 |
outputs=[
|
46 |
gr.Image(type="pil", label="Annotated Image"),
|
47 |
gr.Text(label="Face Count")
|
48 |
],
|
49 |
-
title="
|
50 |
-
description="
|
51 |
)
|
52 |
|
53 |
if __name__ == "__main__":
|
54 |
-
|
|
|
4 |
from supervision import Detections
|
5 |
from PIL import Image, ImageDraw
|
6 |
|
7 |
+
# Download & load the YOLOv8 face detection model
|
8 |
+
model_path = hf_hub_download(
|
9 |
+
repo_id="arnabdhar/YOLOv8-Face-Detection",
|
10 |
+
filename="model.pt"
|
11 |
+
)
|
12 |
model = YOLO(model_path)
|
13 |
|
14 |
+
def count_faces(image: Image.Image):
|
15 |
"""
|
16 |
+
Detects faces in an image and returns the annotated image + face count.
|
|
|
17 |
"""
|
18 |
+
# 1. Run detection
|
19 |
+
results = model(image)[0]
|
20 |
+
dets = Detections.from_ultralytics(results)
|
21 |
+
|
22 |
+
# 2. Draw boxes
|
23 |
+
annotated = image.copy()
|
24 |
+
draw = ImageDraw.Draw(annotated)
|
25 |
+
for x1, y1, x2, y2 in dets.xyxy:
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
26 |
draw.rectangle([x1, y1, x2, y2], outline="red", width=2)
|
|
|
|
|
|
|
|
|
|
|
27 |
|
28 |
+
# 3. Return
|
29 |
+
return annotated, f"Faces detected: {len(dets.xyxy)}"
|
30 |
+
|
31 |
+
# Gradio interface
|
32 |
+
image_app = gr.Interface(
|
33 |
+
fn=count_faces,
|
34 |
inputs=gr.Image(type="pil", label="Upload Image"),
|
35 |
outputs=[
|
36 |
gr.Image(type="pil", label="Annotated Image"),
|
37 |
gr.Text(label="Face Count")
|
38 |
],
|
39 |
+
title="Image Face Counter",
|
40 |
+
description="Detect and count faces in a single image using YOLOv8."
|
41 |
)
|
42 |
|
43 |
if __name__ == "__main__":
|
44 |
+
image_app.launch()
|