From 8b50b2a89789d479b44a7c6048b4b7016cb8d737 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 25 Jun 2023 23:11:00 +0200 Subject: [PATCH] Fix video writing (#3379) --- ultralytics/yolo/engine/predictor.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/ultralytics/yolo/engine/predictor.py b/ultralytics/yolo/engine/predictor.py index a397db0..2ea7a27 100644 --- a/ultralytics/yolo/engine/predictor.py +++ b/ultralytics/yolo/engine/predictor.py @@ -38,7 +38,7 @@ from ultralytics.nn.autobackend import AutoBackend from ultralytics.yolo.cfg import get_cfg from ultralytics.yolo.data import load_inference_source from ultralytics.yolo.data.augment import LetterBox, classify_transforms -from ultralytics.yolo.utils import DEFAULT_CFG, LOGGER, SETTINGS, callbacks, colorstr, ops +from ultralytics.yolo.utils import DEFAULT_CFG, LOGGER, MACOS, SETTINGS, WINDOWS, callbacks, colorstr, ops from ultralytics.yolo.utils.checks import check_imgsz, check_imshow from ultralytics.yolo.utils.files import increment_path from ultralytics.yolo.utils.torch_utils import select_device, smart_inference_mode @@ -341,8 +341,10 @@ class BasePredictor: h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) else: # stream fps, w, h = 30, im0.shape[1], im0.shape[0] - save_path = str(Path(save_path).with_suffix('.mp4')) # force *.mp4 suffix on results videos - self.vid_writer[idx] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)) + suffix = '.mp4' if MACOS else '.avi' if WINDOWS else '.avi' + fourcc = 'avc1' if MACOS else 'WMV2' if WINDOWS else '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].write(im0) def run_callbacks(self, event: str):