Spaces:
Runtime error
Runtime error
| 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() | |