hardiksharma6555 commited on
Commit
eb245fb
·
verified ·
1 Parent(s): 764a99b

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +41 -32
app.py CHANGED
@@ -1,44 +1,53 @@
1
  import gradio as gr
2
- from huggingface_hub import hf_hub_download
3
  from ultralytics import YOLO
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()
 
1
  import gradio as gr
2
+ import torch
3
  from ultralytics import YOLO
4
+ from PIL import Image
 
5
 
6
+ def load_model():
 
 
 
 
 
 
 
7
  """
8
+ Load the YOLOv8 segmentation model onto GPU (if available)
9
+ with mixed‑precision enabled.
10
  """
11
+ device = "cuda" if torch.cuda.is_available() else "cpu" # GPU if available :contentReference[oaicite:2]{index=2}
12
+ model = YOLO('yolov8x-seg.pt').to(device) # Segmentation variant for finer masks :contentReference[oaicite:3]{index=3}
13
+ return model, device
14
+
15
+ model, device = load_model()
16
 
17
+ def count_persons(image: Image.Image) -> str:
18
+ """
19
+ Run inference on the input image, apply TTA, filter for class 0 (person),
20
+ and return the total count.
21
+ """
22
+ # Perform prediction with augmentation (TTA), limit detections, and only class 0
23
+ results = model.predict(
24
+ source=image,
25
+ conf=0.6, # Confidence threshold
26
+ imgsz=640, # Inference resolution
27
+ augment=True, # Test Time Augmentation :contentReference[oaicite:4]{index=4}
28
+ max_det=300, # Cap detections for crowded scenes
29
+ classes=[0] # Only detect persons (class 0) :contentReference[oaicite:5]{index=5}
30
+ )
31
 
32
+ # Sum counts across all results (usually one per image)
33
+ total = sum(len(r.boxes) for r in results)
34
+ return f"Persons detected: {total}"
35
 
36
+ # Build Gradio interface
37
+ demo = gr.Interface(
38
+ fn=count_persons,
39
  inputs=gr.Image(type="pil", label="Upload Image"),
40
+ outputs=gr.Text(label="Person Count"),
41
+ title="Advanced Person Counter with YOLOv8",
42
+ description=(
43
+ "Upload an image to count people using a state‑of‑the‑art "
44
+ "YOLOv8 segmentation model with Test‑Time Augmentation."
45
+ ),
46
+ examples=[ # optional: add example images if you like
47
+ # ["examples/crowd1.jpg"],
48
+ # ["examples/street_scene.jpg"],
49
+ ]
50
  )
51
 
52
  if __name__ == "__main__":
53
+ demo.launch() # Launch locally; add `share=True` for a public link