segment_test_sv.py 2.3 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667
  1. import cv2
  2. import supervision as sv
  3. from ultralytics import YOLO
  4. cap = cv2.VideoCapture(0)
  5. #model = YOLO('best_magent_block.pt')
  6. model = YOLO('yolov8n-seg.pt')
  7. #model = YOLO('best0312.pt')
  8. byte_tracker = sv.ByteTrack()
  9. while True:
  10. ret, frame = cap.read()
  11. result = model(frame)[0]
  12. detections = sv.Detections.from_ultralytics(result)
  13. detections_tracker = byte_tracker.update_with_detections(detections=detections)
  14. #print(detections)
  15. if ret:
  16. polygon_annotator = sv.PolygonAnnotator()
  17. annotated_frame = polygon_annotator.annotate(
  18. scene=frame.copy(), detections=detections)
  19. mask_annotator = sv.MaskAnnotator()
  20. annotated_frame = mask_annotator.annotate(
  21. scene=annotated_frame,detections=detections
  22. )
  23. dot_annotator = sv.DotAnnotator()
  24. annotated_frame = dot_annotator.annotate(
  25. scene=annotated_frame, detections=detections_tracker)
  26. # bounding_box_annotator = sv.BoundingBoxAnnotator()
  27. # annotated_frame = bounding_box_annotator.annotate(
  28. # scene=annotated_frame, detections=detections)
  29. #labels = [detections.confidence]
  30. labels = [
  31. f"#{tracker_id} {result.names[class_id]} {confidence:.2f}"
  32. for class_id, tracker_id, confidence
  33. in zip(detections_tracker.class_id, detections_tracker.tracker_id, detections_tracker.confidence)
  34. ]
  35. # labels_con_id = [
  36. # f"{result.names[class_id]} {confidence:.2f}"
  37. # for class_id, confidence
  38. # in zip(detections.class_id, detections.confidence)
  39. # ]
  40. print(labels)
  41. label_annotator = sv.LabelAnnotator()
  42. annotated_frame = label_annotator.annotate(
  43. scene=annotated_frame,
  44. detections=detections_tracker,
  45. labels=labels
  46. )
  47. position = sv.Position.CENTER_OF_MASS
  48. xy = detections.get_anchors_coordinates(anchor=position)
  49. for detection_idx in range(len(detections)):
  50. print(xy)
  51. center = (int(xy[detection_idx,0]), int(xy[detection_idx,1]))
  52. cv2.circle(annotated_frame, center,radius=3, color=(255,255,0), thickness=-1)
  53. cv2.imshow('frame', annotated_frame)
  54. if cv2.waitKey(1) == ord('q'):
  55. break