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 <glenn.jocher@ultralytics.com>

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                                |

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                                |

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(video_url, 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='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():

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

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',
+                 'save_frames', 'show_labels', 'show_conf', 'visualize', 'augment', 'agnostic_nms', 'retina_masks',
-                 'keras', 'optimize', 'int8', 'dynamic', 'simplify', 'nms', 'profile')
+                 'show_boxes', 'keras', 'optimize', 'int8', 'dynamic', 'simplify', 'nms', 'profile')
 
 
 def cfg2dict(cfg):

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

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(str(Path(save_path).with_suffix(suffix)),
-                self.vid_writer[idx] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*fourcc), fps, (w, h))
+                                                       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, []):