|
|
@ -38,7 +38,7 @@ from ultralytics.nn.autobackend import AutoBackend
|
|
|
|
from ultralytics.yolo.cfg import get_cfg
|
|
|
|
from ultralytics.yolo.cfg import get_cfg
|
|
|
|
from ultralytics.yolo.data import load_inference_source
|
|
|
|
from ultralytics.yolo.data import load_inference_source
|
|
|
|
from ultralytics.yolo.data.augment import LetterBox, classify_transforms
|
|
|
|
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.checks import check_imgsz, check_imshow
|
|
|
|
from ultralytics.yolo.utils.files import increment_path
|
|
|
|
from ultralytics.yolo.utils.files import increment_path
|
|
|
|
from ultralytics.yolo.utils.torch_utils import select_device, smart_inference_mode
|
|
|
|
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))
|
|
|
|
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
|
|
|
|
else: # stream
|
|
|
|
else: # stream
|
|
|
|
fps, w, h = 30, im0.shape[1], im0.shape[0]
|
|
|
|
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
|
|
|
|
suffix = '.mp4' if MACOS else '.avi' if WINDOWS else '.avi'
|
|
|
|
self.vid_writer[idx] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
|
|
|
|
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)
|
|
|
|
self.vid_writer[idx].write(im0)
|
|
|
|
|
|
|
|
|
|
|
|
def run_callbacks(self, event: str):
|
|
|
|
def run_callbacks(self, event: str):
|
|
|
|