cvdetectors commited on
Commit
764a99b
·
verified ·
1 Parent(s): 61dafbf

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +24 -34
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 from Hugging Face Hub
8
- model_path = hf_hub_download(repo_id="arnabdhar/YOLOv8-Face-Detection", filename="model.pt")
9
-
10
- # Load the YOLOv8 face detection model
 
11
  model = YOLO(model_path)
12
 
13
- def detect_faces(image: Image.Image):
14
  """
15
- Detects faces in an input image using YOLOv8 face detection model and returns the annotated image
16
- along with the number of faces detected.
17
  """
18
- # Run inference on the input image
19
- output = model(image)
20
-
21
- # Convert YOLO output to Detections format using the supervision library
22
- results = Detections.from_ultralytics(output[0])
23
-
24
- # Extract bounding boxes; results.xyxy contains boxes in [x1, y1, x2, y2] format
25
- boxes = results.xyxy # This is assumed to be a list-like structure of bounding boxes
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
- # Create a Gradio interface for the face detection function
42
- demo = gr.Interface(
43
- fn=detect_faces,
 
 
 
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="YOLOv8 Face Detector",
50
- description="Upload an image to detect faces using a YOLOv8 face detection model. The detected faces will be highlighted with red bounding boxes."
51
  )
52
 
53
  if __name__ == "__main__":
54
- demo.launch()
 
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()