glenn-jocher commited on
Commit
482af47
1 Parent(s): 8efe977

Add `detect.py` GIF video inference (#6410)

Browse files

* Add detect.py GIF video inference

* Cleanup

Files changed (2) hide show
  1. detect.py +1 -1
  2. utils/datasets.py +2 -2
detect.py CHANGED
@@ -199,7 +199,7 @@ def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s)
199
  h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
200
  else: # stream
201
  fps, w, h = 30, im0.shape[1], im0.shape[0]
202
- save_path += '.mp4'
203
  vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
204
  vid_writer[i].write(im0)
205
 
 
199
  h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
200
  else: # stream
201
  fps, w, h = 30, im0.shape[1], im0.shape[0]
202
+ save_path = str(Path(save_path).with_suffix('.mp4')) # force *.mp4 suffix on results videos
203
  vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
204
  vid_writer[i].write(im0)
205
 
utils/datasets.py CHANGED
@@ -33,8 +33,8 @@ from utils.torch_utils import torch_distributed_zero_first
33
 
34
  # Parameters
35
  HELP_URL = 'https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data'
36
- IMG_FORMATS = ['bmp', 'jpg', 'jpeg', 'png', 'tif', 'tiff', 'dng', 'webp', 'mpo'] # acceptable image suffixes
37
- VID_FORMATS = ['mov', 'avi', 'mp4', 'mpg', 'mpeg', 'm4v', 'wmv', 'mkv'] # acceptable video suffixes
38
  DEVICE_COUNT = max(torch.cuda.device_count(), 1)
39
 
40
  # Get orientation exif tag
 
33
 
34
  # Parameters
35
  HELP_URL = 'https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data'
36
+ IMG_FORMATS = ['bmp', 'dng', 'jpeg', 'jpg', 'mpo', 'png', 'tif', 'tiff', 'webp'] # include image suffixes
37
+ VID_FORMATS = ['avi', 'gif', 'm4v', 'mkv', 'mov', 'mp4', 'mpeg', 'mpg', 'wmv'] # include video suffixes
38
  DEVICE_COUNT = max(torch.cuda.device_count(), 1)
39
 
40
  # Get orientation exif tag