Spaces:
Runtime error
Runtime error
File size: 2,877 Bytes
1fae9dc |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 |
import streamlit as st
import torch
import cv2
import numpy as np
from PIL import Image
from io import BytesIO
st.title("Real-time Object Detection with YOLOv5")
# Load the YOLOv5 model
model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)
# Function to draw the bounding box while selecting ROI
def draw_roi_box(event, x, y, flags, param):
global roi_start, roi_end, selecting_roi, frame_with_box
if event == cv2.EVENT_LBUTTONDOWN:
selecting_roi = True
roi_start = (x, y)
roi_end = (x, y)
elif event == cv2.EVENT_MOUSEMOVE:
if selecting_roi:
roi_end = (x, y)
elif event == cv2.EVENT_LBUTTONUP:
selecting_roi = False
roi_end = (x, y)
cv2.rectangle(frame_with_box, roi_start, roi_end, (0, 255, 0), 2)
selecting_roi = False
roi_start = (0, 0)
roi_end = (0, 0)
frame_with_box = None
# Set the mouse callback function for ROI selection
cv2.namedWindow("ROI Selector")
cv2.setMouseCallback("ROI Selector", draw_roi_box)
st.text("Select ROI by clicking and dragging on the 'ROI Selector' window.")
# Display the camera input widget
img_file_buffer = st.camera_input("Take a picture")
if img_file_buffer is not None:
# Convert the file-like object to a PIL image
img = Image.open(img_file_buffer)
img_array = np.array(img) # Convert PIL Image to numpy array
st.text("Press 'q' to exit the video stream")
while True:
ret, frame = cap.read()
if not ret:
st.error("Failed to capture image")
break
frame_with_box = frame.copy()
if not selecting_roi:
cv2.rectangle(frame_with_box, roi_start, roi_end, (0, 255, 0), 2)
cv2.imshow("ROI Selector", frame_with_box)
key = cv2.waitKey(1) & 0xFF
if key == 13: # Enter key
break
cv2.destroyWindow("ROI Selector")
# Crop the selected ROI
roi = frame[roi_start[1]:roi_end[1], roi_start[0]:roi_end[0]]
while True:
ret, frame = cap.read()
if not ret:
break
# Resize frame for consistent processing
frame = cv2.resize(frame, (1020, 500))
# Perform object detection on the selected ROI
results = model(roi)
# Draw bounding boxes and labels
for index, row in results.pandas().xyxy[0].iterrows():
x1, y1, x2, y2 = int(row['xmin']), int(row['ymin']), int(row['xmax']), int(row['ymax'])
label = row['name']
cv2.rectangle(roi, (x1, y1), (x2, y2), (0, 255, 0), 2)
cv2.putText(roi, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36, 255, 12), 2)
# Display the resulting frame with detections
st.image(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
|