Spaces:
Build error
Build error
| from ultralytics import YOLO | |
| from PIL import Image | |
| import cv2 | |
| import ffmpeg | |
| ffmpeg.input('files/a.MOV').output('files/a.mp4').run() | |
| ffmpeg.input('input.mov').output('output.mp4').run() | |
| def draw_boxes(image, boxes): | |
| for box in boxes: | |
| x1, y1, x2, y2, name, prob = box | |
| cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 2) | |
| cv2.putText(image, f"{name} {prob:.2f}", (x1, y1-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0,255,0), 2) | |
| return image | |
| def detect_objects_on_image(buf): | |
| model = YOLO("yolov8n.pt") | |
| results = model.predict(buf) | |
| result = results[0] | |
| output = [] | |
| for box in result.boxes: | |
| x1, y1, x2, y2 = [ | |
| round(x) for x in box.xyxy[0].tolist() | |
| ] | |
| class_id = box.cls[0].item() | |
| prob = round(box.conf[0].item(), 2) | |
| output.append([ | |
| x1, y1, x2, y2, result.names[class_id], prob | |
| ]) | |
| return output | |
| # model = MaskRCNN("mask_rcnn_model.pth") | |
| # results = model.predict(img) | |
| # masks = results['masks'] | |
| # img = cv2.imread('a.png') | |
| # boxes=detect_objects_on_image(img) | |
| # img_with_boxes = draw_boxes(img, boxes) | |
| # cv2.imshow("test",img_with_boxes) | |
| # cv2.waitKey(0) | |
| model = YOLO("files/yolov8n.pt") | |
| video_path = "files/a.MOV" | |
| cap = cv2.VideoCapture(video_path) | |
| # Loop through the video frames | |
| while cap.isOpened(): | |
| # Read a frame from the video | |
| success, frame = cap.read() | |
| if success: | |
| # Run YOLOv8 tracking on the frame, persisting tracks between frames | |
| results = model.track(frame, persist=True) | |
| # Visualize the results on the frame | |
| annotated_frame = results[0].plot() | |
| # Display the annotated frame | |
| cv2.imshow("YOLOv8 Tracking", annotated_frame) | |
| # Break the loop if 'q' is pressed | |
| if cv2.waitKey(1) & 0xFF == ord("q"): | |
| break | |
| else: | |
| # Break the loop if the end of the video is reached | |
| break | |
| # Release the video capture object and close the display window | |
| cap.release() | |
| cv2.destroyAllWindows() |