tahsinraiyan000 commited on
Commit
d10c3f2
·
verified ·
1 Parent(s): 961117c

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +71 -0
app.py ADDED
@@ -0,0 +1,71 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import cv2
3
+ import numpy as np
4
+ from collections import defaultdict
5
+ from transformers import pipeline, AutoModelForObjectDetection, AutoFeatureExtractor
6
+
7
+ # Load Xenova's DETR model
8
+ model_name = "Xenova/detr-resnet-50"
9
+ feature_extractor = AutoFeatureExtractor.from_pretrained(model_name)
10
+ model = AutoModelForObjectDetection.from_pretrained(model_name)
11
+ detector = pipeline("object-detection", model=model, feature_extractor=feature_extractor)
12
+
13
+ # Global counter (persists across frames)
14
+ object_counter = defaultdict(int)
15
+
16
+ def process_video(video_path):
17
+ cap = cv2.VideoCapture(video_path)
18
+
19
+ while cap.isOpened():
20
+ ret, frame = cap.read()
21
+ if not ret:
22
+ break
23
+
24
+ # Convert frame to RGB (required by DETR)
25
+ rgb_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
26
+
27
+ # Detect objects
28
+ results = detector(rgb_frame)
29
+
30
+ # Draw boxes and update counter
31
+ for obj in results:
32
+ label = obj["label"]
33
+ score = obj["score"]
34
+ box = obj["box"]
35
+
36
+ # Update counter
37
+ object_counter[label] += 1
38
+
39
+ # Draw bounding box
40
+ xmin, ymin, xmax, ymax = int(box["xmin"]), int(box["ymin"]), int(box["xmax"]), int(box["ymax"])
41
+ cv2.rectangle(frame, (xmin, ymin), (xmax, ymax), (0, 255, 0), 2)
42
+ cv2.putText(frame, f"{label} ({score:.2f})", (xmin, ymin-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
43
+
44
+ # Overlay counter on the frame
45
+ counter_text = "\n".join([f"{k}: {v}" for k, v in object_counter.items()])
46
+ cv2.putText(frame, counter_text, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 0, 0), 2)
47
+
48
+ # Convert back to RGB for Gradio
49
+ output_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
50
+ yield output_frame
51
+
52
+ cap.release()
53
+
54
+ # Gradio UI
55
+ with gr.Blocks() as demo:
56
+ gr.Markdown("# 🎥 Video Object Detection (Xenova/detr-resnet-50)")
57
+ with gr.Row():
58
+ video_input = gr.Video(label="Upload Video")
59
+ video_output = gr.Image(label="Live Detection")
60
+ with gr.Row():
61
+ gr.Markdown("### Cumulative Object Counts")
62
+ counter_display = gr.Textbox(value="", label="Counter", interactive=False)
63
+
64
+ video_input.change(
65
+ fn=process_video,
66
+ inputs=video_input,
67
+ outputs=video_output,
68
+ show_progress=True,
69
+ )
70
+
71
+ demo.launch()