diff --git a/README.md b/README.md index 5743e45..3495b7c 100644 --- a/README.md +++ b/README.md @@ -120,32 +120,20 @@ python main.py data/sample_videos/test.mp4 ```python import asone -from asone import utils -from asone import ASOne -import cv2 - -video_path = 'data/sample_videos/test.mp4' -detector = ASOne(detector=asone.YOLOV7_PYTORCH, use_cuda=True) # Set use_cuda to False for cpu +from asone import ASOne, VideoReader +from tqdm import tqdm filter_classes = ['car'] # Set to None to detect all classes -cap = cv2.VideoCapture(video_path) - -while True: - _, frame = cap.read() - if not _: - break +detector = ASOne(detector=asone.YOLOV7_PYTORCH, use_cuda=True) # Set use_cuda to False for cpu +cap = VideoReader('data/sample_videos/test.mp4') +for frame in tqdm(cap): dets, img_info = detector.detect(frame, filter_classes=filter_classes) + frame = ASOne.draw_boxes(frame, dets) - bbox_xyxy = dets[:, :4] - scores = dets[:, 4] - class_ids = dets[:, 5] - - frame = utils.draw_boxes(frame, bbox_xyxy, class_ids=class_ids) - + # Show Video cv2.imshow('result', frame) - if cv2.waitKey(25) & 0xFF == ord('q'): break ```