pratikshahp commited on
Commit
25e2600
·
verified ·
1 Parent(s): 09ac133

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +44 -0
app.py ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ from transformers import DetrImageProcessor, DetrForObjectDetection
2
+ from PIL import Image, ImageDraw
3
+ import torch
4
+ import gradio as gr
5
+
6
+ # Load model and processor
7
+ processor = DetrImageProcessor.from_pretrained("facebook/detr-resnet-50")
8
+ model = DetrForObjectDetection.from_pretrained("facebook/detr-resnet-50")
9
+
10
+ FACE_CLASS_INDEX = 1 # COCO class ID for 'person'
11
+
12
+ def detect_faces(img: Image.Image):
13
+ # Make a copy to draw on
14
+ img_draw = img.copy()
15
+ draw = ImageDraw.Draw(img_draw)
16
+
17
+ # Preprocess and predict
18
+ inputs = processor(images=img, return_tensors="pt")
19
+ outputs = model(**inputs)
20
+
21
+ # Get results
22
+ target_sizes = torch.tensor([img.size[::-1]])
23
+ results = processor.post_process_object_detection(outputs, target_sizes=target_sizes, threshold=0.8)[0]
24
+
25
+ count = 0
26
+ for score, label, box in zip(results["scores"], results["labels"], results["boxes"]):
27
+ if label.item() == FACE_CLASS_INDEX:
28
+ count += 1
29
+ box = [round(i, 2) for i in box.tolist()]
30
+ draw.rectangle(box, outline="lime", width=3)
31
+ draw.text((box[0], box[1] - 10), f"{score:.2f}", fill="lime")
32
+
33
+ return img_draw, f"Total Persons Detected: {count}"
34
+
35
+ # Gradio Interface
36
+ iface = gr.Interface(
37
+ fn=detect_faces,
38
+ inputs=gr.Image(type="pil"),
39
+ outputs=[gr.Image(type="pil"), gr.Text()],
40
+ title="Person Detection with DETR",
41
+ description="Uses DETR model to detect people (class 1 - COCO dataset). Note: not specialized for face detection."
42
+ )
43
+
44
+ iface.launch()