From 4096b261fcc3689dd656e779d362c9ca474441c7 Mon Sep 17 00:00:00 2001 From: Jason Sohn <75611662+tensorturtle@users.noreply.github.com> Date: Mon, 27 Nov 2023 18:36:24 +0900 Subject: [PATCH] `ultralytics 8.0.219` new `save_frames=False` predict arg (#6396) Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com> Co-authored-by: Glenn Jocher --- docs/en/modes/predict.md | 1 + docs/en/usage/cfg.md | 1 + tests/test_python.py | 7 ++++--- ultralytics/__init__.py | 2 +- ultralytics/cfg/__init__.py | 4 ++-- ultralytics/cfg/default.yaml | 1 + ultralytics/engine/predictor.py | 18 ++++++++++++++---- 7 files changed, 24 insertions(+), 10 deletions(-) diff --git a/docs/en/modes/predict.md b/docs/en/modes/predict.md index 4adf2e5a91..446d36edfb 100644 --- a/docs/en/modes/predict.md +++ b/docs/en/modes/predict.md @@ -364,6 +364,7 @@ Visualization arguments: |---------------|---------------|---------|-----------------------------------------------------------------| | `show` | `bool` | `False` | show predicted images and videos if environment allows | | `save` | `bool` | `False` | save predicted images and videos | +| `save_frames` | `bool` | `False` | save predicted individual video frames | | `save_txt` | `bool` | `False` | save results as `.txt` file | | `save_conf` | `bool` | `False` | save results with confidence scores | | `save_crop` | `bool` | `False` | save cropped images with results | diff --git a/docs/en/usage/cfg.md b/docs/en/usage/cfg.md index d416fc05ed..5bb0fadcec 100644 --- a/docs/en/usage/cfg.md +++ b/docs/en/usage/cfg.md @@ -153,6 +153,7 @@ Visualization arguments: |---------------|---------------|---------|-----------------------------------------------------------------| | `show` | `bool` | `False` | show predicted images and videos if environment allows | | `save` | `bool` | `False` | save predicted images and videos | +| `save_frames` | `bool` | `False` | save predicted individual video frames | | `save_txt` | `bool` | `False` | save results as `.txt` file | | `save_conf` | `bool` | `False` | save results with confidence scores | | `save_crop` | `bool` | `False` | save cropped images with results | diff --git a/tests/test_python.py b/tests/test_python.py index 4c740a9789..741ad5d08e 100644 --- a/tests/test_python.py +++ b/tests/test_python.py @@ -154,9 +154,10 @@ def test_track_stream(): """ import yaml + video_url = 'https://ultralytics.com/assets/decelera_portrait_min.mov' model = YOLO(MODEL) - model.track('https://ultralytics.com/assets/decelera_portrait_min.mov', imgsz=160, tracker='bytetrack.yaml') - model.track('https://ultralytics.com/assets/decelera_portrait_min.mov', imgsz=160, tracker='botsort.yaml') + model.track(video_url, imgsz=160, tracker='bytetrack.yaml') + model.track(video_url, imgsz=160, tracker='botsort.yaml', save_frames=True) # test frame saving also # Test Global Motion Compensation (GMC) methods for gmc in 'orb', 'sift', 'ecc': @@ -166,7 +167,7 @@ def test_track_stream(): data['gmc_method'] = gmc with open(tracker, 'w', encoding='utf-8') as f: yaml.safe_dump(data, f) - model.track('https://ultralytics.com/assets/decelera_portrait_min.mov', imgsz=160, tracker=tracker) + model.track(video_url, imgsz=160, tracker=tracker) def test_val(): diff --git a/ultralytics/__init__.py b/ultralytics/__init__.py index 53b6e640fb..710a7723bc 100644 --- a/ultralytics/__init__.py +++ b/ultralytics/__init__.py @@ -1,6 +1,6 @@ # Ultralytics YOLO 🚀, AGPL-3.0 license -__version__ = '8.0.218' +__version__ = '8.0.219' from ultralytics.models import RTDETR, SAM, YOLO from ultralytics.models.fastsam import FastSAM diff --git a/ultralytics/cfg/__init__.py b/ultralytics/cfg/__init__.py index aed0eac88f..652e8fc29c 100644 --- a/ultralytics/cfg/__init__.py +++ b/ultralytics/cfg/__init__.py @@ -71,8 +71,8 @@ CFG_INT_KEYS = ('epochs', 'patience', 'batch', 'workers', 'seed', 'close_mosaic' 'line_width', 'workspace', 'nbs', 'save_period') CFG_BOOL_KEYS = ('save', 'exist_ok', 'verbose', 'deterministic', 'single_cls', 'rect', 'cos_lr', 'overlap_mask', 'val', 'save_json', 'save_hybrid', 'half', 'dnn', 'plots', 'show', 'save_txt', 'save_conf', 'save_crop', - 'show_labels', 'show_conf', 'visualize', 'augment', 'agnostic_nms', 'retina_masks', 'show_boxes', - 'keras', 'optimize', 'int8', 'dynamic', 'simplify', 'nms', 'profile') + 'save_frames', 'show_labels', 'show_conf', 'visualize', 'augment', 'agnostic_nms', 'retina_masks', + 'show_boxes', 'keras', 'optimize', 'int8', 'dynamic', 'simplify', 'nms', 'profile') def cfg2dict(cfg): diff --git a/ultralytics/cfg/default.yaml b/ultralytics/cfg/default.yaml index 6549e4299d..c9df7ea11c 100644 --- a/ultralytics/cfg/default.yaml +++ b/ultralytics/cfg/default.yaml @@ -63,6 +63,7 @@ retina_masks: False # (bool) use high-resolution segmentation masks # Visualize settings --------------------------------------------------------------------------------------------------- show: False # (bool) show predicted images and videos if environment allows +save_frames: False # (bool) save predicted individual video frames save_txt: False # (bool) save results as .txt file save_conf: False # (bool) save results with confidence scores save_crop: False # (bool) save cropped images with results diff --git a/ultralytics/engine/predictor.py b/ultralytics/engine/predictor.py index c9bcca4d1f..3df7bdb5d7 100644 --- a/ultralytics/engine/predictor.py +++ b/ultralytics/engine/predictor.py @@ -98,7 +98,7 @@ class BasePredictor: self.imgsz = None self.device = None self.dataset = None - self.vid_path, self.vid_writer = None, None + self.vid_path, self.vid_writer, self.vid_frame = None, None, None self.plotted_img = None self.data_path = None self.source_type = None @@ -221,7 +221,9 @@ class BasePredictor: len(self.dataset) > 1000 or # images any(getattr(self.dataset, 'video_flag', [False]))): # videos LOGGER.warning(STREAM_WARNING) - self.vid_path, self.vid_writer = [None] * self.dataset.bs, [None] * self.dataset.bs + self.vid_path = [None] * self.dataset.bs + self.vid_writer = [None] * self.dataset.bs + self.vid_frame = [None] * self.dataset.bs @smart_inference_mode() def stream_inference(self, source=None, model=None, *args, **kwargs): @@ -341,8 +343,11 @@ class BasePredictor: if self.dataset.mode == 'image': cv2.imwrite(save_path, im0) else: # 'video' or 'stream' + frames_path = f'{save_path.split(".", 1)[0]}_frames/' if self.vid_path[idx] != save_path: # new video + Path(frames_path).mkdir(parents=True, exist_ok=True) self.vid_path[idx] = save_path + self.vid_frame[idx] = 0 if isinstance(self.vid_writer[idx], cv2.VideoWriter): self.vid_writer[idx].release() # release previous video writer if vid_cap: # video @@ -352,10 +357,15 @@ class BasePredictor: else: # stream fps, w, h = 30, im0.shape[1], im0.shape[0] suffix, fourcc = ('.mp4', 'avc1') if MACOS else ('.avi', 'WMV2') if WINDOWS else ('.avi', 'MJPG') - save_path = str(Path(save_path).with_suffix(suffix)) - self.vid_writer[idx] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h)) + self.vid_writer[idx] = cv2.VideoWriter(str(Path(save_path).with_suffix(suffix)), + cv2.VideoWriter_fourcc(*fourcc), fps, (w, h)) + # Write video self.vid_writer[idx].write(im0) + # Write frame + cv2.imwrite(f'{frames_path}{self.vid_frame[idx]}.jpg', im0) + self.vid_frame[idx] += 1 + def run_callbacks(self, event: str): """Runs all registered callbacks for a specific event.""" for callback in self.callbacks.get(event, []):