Spaces:
Runtime error
Runtime error
File size: 5,167 Bytes
ef9b1b9 | 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 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 | from ultralytics import YOLO
import subprocess
import threading
import os
import time
import cv2
from PIL import Image
import threading
class ObjectTracker:
def __init__(self, video_source=None, model_path='best.onnx'):
# Initialize the YOLO model and video capture
self.model = YOLO(model_path, task="detect")
if video_source is not None:
self.cap = cv2.VideoCapture(video_source)
# Object classes and tracking information
self.class_names = ["head", "helmet", "motorbike"]
self.passing_ids = set()
self.previous_centers = {}
self.threads = []
self.start = True
#
def process_object(self, buffer):
"""Run subprocess to process object asynchronously."""
process = subprocess.Popen(['python', 'recognition.py'], stdin=subprocess.PIPE)
process.communicate(input=buffer.tobytes())
print(2)
def detect_and_track(self):
"""Main loop for detecting and tracking objects."""
while True:
ret, frame = self.cap.read()
if not ret:
break # Stop if no frame is returned
frame = cv2.resize(frame, (600, 600))
motorbike_detect = list(self.model.track(frame, persist=True, conf=0.65, classes=[2]))
annotated_frame = motorbike_detect[0].plot()
self.handle_detections(motorbike_detect, frame)
cv2.imshow('Webcam', annotated_frame)
if cv2.waitKey(1) == ord('q'):
break
self.cap.release()
cv2.destroyAllWindows()
def detect_and_track_2(self, video):
self.cap = cv2.VideoCapture(video)
self.start = True
"""Main loop for detecting and tracking objects."""
while self.start:
ret, frame = self.cap.read()
if not ret:
break # Stop if no frame is returned
frame = cv2.resize(frame, (600, 600))
motorbike_detect = list(self.model.track(frame, persist=True, conf=0.65, classes=[2]))
annotated_frame = motorbike_detect[0].plot()
self.handle_detections(motorbike_detect, frame)
yield annotated_frame
if cv2.waitKey(1) == ord('q'):
break
self.cap.release()
cv2.destroyAllWindows()
def handle_detections(self, detections, frame):
"""Process detected objects and track their movement."""
for result in detections[0].boxes.data.tolist():
if len(result) == 7: # Ensure detection has all necessary attributes
x1, y1, x2, y2, obj_id, conf, label = result
if obj_id not in self.previous_centers:
self.previous_centers[obj_id] = 0
center_x = int((x1 + x2) / 2)
print(f"id: {obj_id} {self.previous_centers[obj_id]} {center_x}")
if self.crossed_line(obj_id, center_x):
self.save_and_process_object(frame, x1, y1, x2, y2, obj_id)
self.previous_centers[obj_id] = center_x
def crossed_line(self, obj_id, center_x):
"""Check if the object has crossed the predefined line."""
return (
self.previous_centers[obj_id] >= 150 and
center_x <= 150 and
obj_id not in self.passing_ids
)
def save_and_process_object(self, frame, x1, y1, x2, y2, obj_id):
"""Save the object image and process it in a new thread."""
self.passing_ids.add(obj_id)
print(f"id:{obj_id} crossed the line")
img_crop = frame[int(y1 - 150):int(y2), int(x1):int(x2 + 20), :]
# cv2.imwrite(f'img/object_crossed{obj_id}.jpg', img_crop)
_, buffer = cv2.imencode('.jpg',img_crop)
with open(f'img/object_crossed{obj_id}.jpg', "wb") as f:
f.write(buffer)
print(1)
thread = threading.Thread(target=self.process_object, args=(buffer,))
thread.start()
self.threads.append(thread)
print(3)
def test_camera():
cap = cv2.VideoCapture('http://100.86.155.229:8080/video')
if not cap.isOpened():
print("Error: Could not open the webcam.")
else:
print("Press 's' to save a frame, or 'q' to quit.")
# Loop to continuously capture frames
while True:
ret, frame = cap.read() # Capture a frame
if not ret:
print("Failed to grab frame.")
break
# Display the frame in a window
cv2.imshow("Webcam Stream", frame)
# Check for key press events
key = cv2.waitKey(1) & 0xFF # Mask for cross-platform compatibility
if key == ord('s'): # Press 's' to save a frame
cv2.imwrite("saved_frame.jpg", frame) # Save the frame as a JPEG
print("Frame saved as 'saved_frame.jpg'.")
elif key == ord('q'): # Press 'q' to quit
print("Exiting...")
break
# Release the webcam and close the window
cap.release()
cv2.destroyAllWindows()
if __name__ == "__main__":
tracker = ObjectTracker()
tracker.detect_and_track()
#test_camera()
|