cvdetectors commited on
Commit
763f05b
·
verified ·
1 Parent(s): 9ca210a

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +54 -0
app.py ADDED
@@ -0,0 +1,54 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import torch
3
+ import torchvision
4
+ from torchvision import transforms
5
+ import numpy as np
6
+ from PIL import Image
7
+
8
+ # Use GPU if available
9
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
10
+ torch.backends.cudnn.benchmark = True
11
+ print(f"Using device: {device}")
12
+
13
+ # Load lightweight detection model
14
+ model = torchvision.models.detection.fasterrcnn_mobilenet_v3_large_fpn(pretrained=True)
15
+ model.to(device)
16
+ model.eval()
17
+
18
+ # Image transformation
19
+ transform = transforms.Compose([
20
+ transforms.ToTensor(),
21
+ ])
22
+
23
+ # Mixed precision (FP16) for CUDA
24
+ autocast = torch.cuda.amp.autocast if device.type == "cuda" else torch.cpu.amp.autocast
25
+
26
+ def count_persons(image):
27
+ # Convert image to tensor
28
+ image_rgb = np.array(image.convert("RGB"))
29
+ img_tensor = transform(image_rgb).to(device).unsqueeze(0)
30
+
31
+ # Inference
32
+ with torch.no_grad():
33
+ with autocast():
34
+ outputs = model(img_tensor)[0]
35
+
36
+ # Count persons (label 1 in COCO)
37
+ person_count = sum(
38
+ 1 for label, score in zip(outputs['labels'], outputs['scores'])
39
+ if label.item() == 1 and score.item() > 0.8
40
+ )
41
+
42
+ return f"Number of persons detected: {person_count}"
43
+
44
+ # Gradio interface for image upload
45
+ demo = gr.Interface(
46
+ fn=count_persons,
47
+ inputs=gr.Image(type="pil", label="Upload Image"),
48
+ outputs=gr.Text(label="Person Count"),
49
+ title="Person Counter in Image (Fast)",
50
+ description="Upload an image to count the number of people using a fast MobileNet-based detector. GPU supported."
51
+ )
52
+
53
+ if __name__ == "__main__":
54
+ demo.launch()