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
- detect.py +1 -1
- 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 |
-
|
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', '
|
37 |
-
VID_FORMATS = ['
|
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
|