Spaces:
Running
Running
Update app.py
Browse files
app.py
CHANGED
@@ -44,15 +44,7 @@ colors = np.random.uniform(0, 255, size=(len(model.names), 3))
|
|
44 |
# Performance tracking
|
45 |
total_inference_time = 0
|
46 |
inference_count = 0
|
47 |
-
|
48 |
-
for _ in range(30): # Initialize with reasonable values
|
49 |
-
fps_queue.put(30.0)
|
50 |
-
|
51 |
-
# Threading variables
|
52 |
-
processing_lock = threading.Lock()
|
53 |
-
stop_event = threading.Event()
|
54 |
-
frame_queue = Queue(maxsize=2) # Small queue to avoid lag
|
55 |
-
result_queue = Queue(maxsize=2)
|
56 |
|
57 |
def detect_objects(image):
|
58 |
"""Process a single image for object detection"""
|
@@ -107,135 +99,71 @@ def detect_objects(image):
|
|
107 |
|
108 |
return output_image
|
109 |
|
110 |
-
def
|
111 |
-
"""
|
112 |
-
|
113 |
-
try:
|
114 |
-
if not frame_queue.empty():
|
115 |
-
frame = frame_queue.get()
|
116 |
-
|
117 |
-
# Skip if there's a processing lock (from image upload)
|
118 |
-
if processing_lock.locked():
|
119 |
-
result_queue.put(frame) # Return unprocessed frame
|
120 |
-
continue
|
121 |
-
|
122 |
-
# Process the frame
|
123 |
-
start_time = time.time()
|
124 |
-
with torch.no_grad(): # Ensure no gradients for inference
|
125 |
-
input_size = 384 # Smaller size for real-time processing
|
126 |
-
results = model(frame['image'], size=input_size)
|
127 |
-
|
128 |
-
# Calculate FPS
|
129 |
-
inference_time = time.time() - start_time
|
130 |
-
current_fps = 1 / inference_time if inference_time > 0 else 30
|
131 |
-
|
132 |
-
# Update rolling FPS average
|
133 |
-
if not fps_queue.full():
|
134 |
-
fps_queue.put(current_fps)
|
135 |
-
else:
|
136 |
-
try:
|
137 |
-
fps_queue.get_nowait()
|
138 |
-
fps_queue.put(current_fps)
|
139 |
-
except:
|
140 |
-
pass
|
141 |
-
|
142 |
-
fps_values = list(fps_queue.queue)
|
143 |
-
avg_fps = sum(fps_values) / len(fps_values) if fps_values else 30.0
|
144 |
-
|
145 |
-
# Draw detections
|
146 |
-
output = frame['image'].copy()
|
147 |
-
detections = results.pred[0].cpu().numpy()
|
148 |
-
|
149 |
-
for *xyxy, conf, cls in detections:
|
150 |
-
x1, y1, x2, y2 = map(int, xyxy)
|
151 |
-
class_id = int(cls)
|
152 |
-
color = colors[class_id].tolist()
|
153 |
-
|
154 |
-
# Draw rectangle and label
|
155 |
-
cv2.rectangle(output, (x1, y1), (x2, y2), color, 2, lineType=cv2.LINE_AA)
|
156 |
-
|
157 |
-
label = f"{model.names[class_id]} {conf:.2f}"
|
158 |
-
font_scale, font_thickness = 0.6, 1 # Smaller for real-time
|
159 |
-
(w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness)
|
160 |
-
|
161 |
-
cv2.rectangle(output, (x1, y1 - h - 5), (x1 + w + 5, y1), color, -1)
|
162 |
-
cv2.putText(output, label, (x1 + 3, y1 - 3),
|
163 |
-
cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), font_thickness, lineType=cv2.LINE_AA)
|
164 |
-
|
165 |
-
# Add FPS counter
|
166 |
-
cv2.rectangle(output, (10, 10), (210, 80), (0, 0, 0), -1)
|
167 |
-
cv2.putText(output, f"FPS: {current_fps:.1f}", (20, 40),
|
168 |
-
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2, lineType=cv2.LINE_AA)
|
169 |
-
cv2.putText(output, f"Avg FPS: {avg_fps:.1f}", (20, 70),
|
170 |
-
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2, lineType=cv2.LINE_AA)
|
171 |
-
|
172 |
-
# Put the processed frame in the result queue
|
173 |
-
if not result_queue.full():
|
174 |
-
result_queue.put({'image': output, 'fps': current_fps})
|
175 |
-
else:
|
176 |
-
time.sleep(0.001) # Small sleep to prevent CPU spinning
|
177 |
-
except Exception as e:
|
178 |
-
print(f"Error in frame processing thread: {e}")
|
179 |
-
time.sleep(0.1) # Pause briefly on error
|
180 |
-
|
181 |
-
def webcam_feed():
|
182 |
-
"""Generator function for webcam feed"""
|
183 |
-
# Start the processing thread if not already running
|
184 |
-
if not any(thread.name == "frame_processor" for thread in threading.enumerate()):
|
185 |
-
stop_event.clear()
|
186 |
-
processor = threading.Thread(target=process_frame_thread, name="frame_processor", daemon=True)
|
187 |
-
processor.start()
|
188 |
|
189 |
-
|
190 |
-
|
191 |
-
if not cap.isOpened():
|
192 |
-
print("Warning: Unable to open webcam! Using dummy frames instead.")
|
193 |
-
# Create a dummy frame with a message
|
194 |
-
dummy_frame = np.zeros((480, 640, 3), dtype=np.uint8)
|
195 |
-
cv2.putText(dummy_frame, "Webcam not available", (100, 240),
|
196 |
-
cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2)
|
197 |
-
while True:
|
198 |
-
yield dummy_frame
|
199 |
-
time.sleep(0.033) # ~30 FPS
|
200 |
|
201 |
-
|
202 |
-
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
|
203 |
-
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
|
204 |
-
cap.set(cv2.CAP_PROP_FPS, 30) # Request 30 FPS from camera if supported
|
205 |
|
206 |
-
|
207 |
-
|
208 |
-
|
209 |
-
|
210 |
-
|
211 |
-
|
212 |
-
|
213 |
-
|
214 |
-
|
215 |
-
|
216 |
-
|
217 |
-
|
218 |
-
|
219 |
-
|
220 |
-
|
221 |
-
|
222 |
-
|
223 |
-
|
224 |
-
|
225 |
-
|
226 |
-
|
227 |
-
|
228 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
229 |
|
230 |
def process_uploaded_image(image):
|
231 |
-
"""Process an uploaded image
|
232 |
-
|
233 |
-
return detect_objects(image)
|
234 |
|
235 |
# Setup Gradio interface
|
236 |
example_images = ["spring_street_after.jpg", "pexels-hikaique-109919.jpg"]
|
237 |
os.makedirs("examples", exist_ok=True)
|
238 |
|
|
|
239 |
with gr.Blocks(title="YOLOv5 Object Detection - Real-time & Image Upload") as demo:
|
240 |
gr.Markdown("""
|
241 |
# YOLOv5 Object Detection
|
@@ -248,7 +176,17 @@ with gr.Blocks(title="YOLOv5 Object Detection - Real-time & Image Upload") as de
|
|
248 |
### Real-time Object Detection
|
249 |
Using your webcam for continuous object detection at 30+ FPS.
|
250 |
""")
|
251 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
252 |
|
253 |
with gr.TabItem("Image Upload"):
|
254 |
gr.Markdown("""
|
@@ -275,16 +213,5 @@ with gr.Blocks(title="YOLOv5 Object Detection - Real-time & Image Upload") as de
|
|
275 |
# Set up event handlers
|
276 |
submit_button.click(fn=process_uploaded_image, inputs=input_image, outputs=output_image)
|
277 |
clear_button.click(lambda: (None, None), None, [input_image, output_image])
|
278 |
-
|
279 |
-
# Start webcam feed
|
280 |
-
demo.load(fn=lambda: None, inputs=None, outputs=webcam_output)
|
281 |
-
webcam_output.update(webcam_feed)
|
282 |
-
|
283 |
-
# Cleanup function to stop threads when app closes
|
284 |
-
def cleanup():
|
285 |
-
stop_event.set()
|
286 |
-
print("Cleaning up threads...")
|
287 |
|
288 |
-
# Register cleanup handler
|
289 |
-
demo.close = cleanup
|
290 |
demo.launch(share=False)
|
|
|
44 |
# Performance tracking
|
45 |
total_inference_time = 0
|
46 |
inference_count = 0
|
47 |
+
last_fps_values = [] # Store recent FPS values
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
48 |
|
49 |
def detect_objects(image):
|
50 |
"""Process a single image for object detection"""
|
|
|
99 |
|
100 |
return output_image
|
101 |
|
102 |
+
def process_webcam_frame(frame):
|
103 |
+
"""Process a single frame from webcam"""
|
104 |
+
global last_fps_values
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
105 |
|
106 |
+
if frame is None:
|
107 |
+
return None
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
108 |
|
109 |
+
start_time = time.time()
|
|
|
|
|
|
|
110 |
|
111 |
+
# Use a smaller size for real-time
|
112 |
+
input_size = 384
|
113 |
+
|
114 |
+
# Process the frame
|
115 |
+
with torch.no_grad():
|
116 |
+
results = model(frame, size=input_size)
|
117 |
+
|
118 |
+
# Calculate FPS
|
119 |
+
inference_time = time.time() - start_time
|
120 |
+
current_fps = 1 / inference_time if inference_time > 0 else 30
|
121 |
+
|
122 |
+
# Update FPS history (keep last 30 values)
|
123 |
+
last_fps_values.append(current_fps)
|
124 |
+
if len(last_fps_values) > 30:
|
125 |
+
last_fps_values.pop(0)
|
126 |
+
avg_fps = sum(last_fps_values) / len(last_fps_values)
|
127 |
+
|
128 |
+
# Create output image
|
129 |
+
output = frame.copy()
|
130 |
+
|
131 |
+
# Draw detections
|
132 |
+
detections = results.pred[0].cpu().numpy()
|
133 |
+
for *xyxy, conf, cls in detections:
|
134 |
+
x1, y1, x2, y2 = map(int, xyxy)
|
135 |
+
class_id = int(cls)
|
136 |
+
color = colors[class_id].tolist()
|
137 |
+
|
138 |
+
# Draw rectangle and label
|
139 |
+
cv2.rectangle(output, (x1, y1), (x2, y2), color, 2, lineType=cv2.LINE_AA)
|
140 |
+
|
141 |
+
label = f"{model.names[class_id]} {conf:.2f}"
|
142 |
+
font_scale, font_thickness = 0.6, 1
|
143 |
+
(w, h), _ = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, font_scale, font_thickness)
|
144 |
+
|
145 |
+
cv2.rectangle(output, (x1, y1 - h - 5), (x1 + w + 5, y1), color, -1)
|
146 |
+
cv2.putText(output, label, (x1 + 3, y1 - 3),
|
147 |
+
cv2.FONT_HERSHEY_SIMPLEX, font_scale, (255, 255, 255), font_thickness, lineType=cv2.LINE_AA)
|
148 |
+
|
149 |
+
# Add FPS counter
|
150 |
+
cv2.rectangle(output, (10, 10), (210, 80), (0, 0, 0), -1)
|
151 |
+
cv2.putText(output, f"FPS: {current_fps:.1f}", (20, 40),
|
152 |
+
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2, lineType=cv2.LINE_AA)
|
153 |
+
cv2.putText(output, f"Avg FPS: {avg_fps:.1f}", (20, 70),
|
154 |
+
cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2, lineType=cv2.LINE_AA)
|
155 |
+
|
156 |
+
return output
|
157 |
|
158 |
def process_uploaded_image(image):
|
159 |
+
"""Process an uploaded image"""
|
160 |
+
return detect_objects(image)
|
|
|
161 |
|
162 |
# Setup Gradio interface
|
163 |
example_images = ["spring_street_after.jpg", "pexels-hikaique-109919.jpg"]
|
164 |
os.makedirs("examples", exist_ok=True)
|
165 |
|
166 |
+
# Simplified interface with proper webcam handling
|
167 |
with gr.Blocks(title="YOLOv5 Object Detection - Real-time & Image Upload") as demo:
|
168 |
gr.Markdown("""
|
169 |
# YOLOv5 Object Detection
|
|
|
176 |
### Real-time Object Detection
|
177 |
Using your webcam for continuous object detection at 30+ FPS.
|
178 |
""")
|
179 |
+
# Use Gradio's webcam component with processing function
|
180 |
+
webcam = gr.Webcam(label="Webcam Input")
|
181 |
+
webcam_output = gr.Image(label="Real-time Detection")
|
182 |
+
detect_button = gr.Button("Detect Objects")
|
183 |
+
|
184 |
+
# Connect webcam to processor
|
185 |
+
detect_button.click(
|
186 |
+
fn=process_webcam_frame,
|
187 |
+
inputs=webcam,
|
188 |
+
outputs=webcam_output
|
189 |
+
)
|
190 |
|
191 |
with gr.TabItem("Image Upload"):
|
192 |
gr.Markdown("""
|
|
|
213 |
# Set up event handlers
|
214 |
submit_button.click(fn=process_uploaded_image, inputs=input_image, outputs=output_image)
|
215 |
clear_button.click(lambda: (None, None), None, [input_image, output_image])
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
216 |
|
|
|
|
|
217 |
demo.launch(share=False)
|