obb.py 812 B

1234567891011121314151617181920212223242526272829303132
  1. import cv2
  2. from ultralytics import YOLO
  3. # Load a model
  4. model = YOLO('yolov8n-obb.pt') # load an official model
  5. # Open the input video
  6. cap = cv2.VideoCapture('obb.mp4')
  7. if not cap.isOpened():
  8. raise Exception("Error: Could not open video.")
  9. while cap.isOpened():
  10. # Read a frame from the input video
  11. success, frame = cap.read()
  12. if success:
  13. results = model.predict(source=frame, verbose=False)
  14. for result in results:
  15. boxs = result.obb.xyxyxyxy.cpu().numpy().astype(int)
  16. for box in boxs:
  17. color = (0, 255, 0)
  18. image = cv2.polylines(frame, [box], True, color, 1)
  19. cv2.imshow("Obb example", frame)
  20. key = cv2.waitKey(1)
  21. if key & 0xff == ord('q'):
  22. break
  23. cap.release()
  24. cv2.destroyAllWindows()