xclele/yolov10
1
0
Fork 0
mirror of https://github.com/THU-MIG/yolov10.git synced 2025-05-23 13:34:23 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

ultralytics 8.0.219 new save_frames=False predict arg (#6396)

Browse Source 
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>
This commit is contained in:
Jason Sohn 2023-11-27 18:36:24 +09:00 committed by GitHub
parent fdcf0dd4fd
commit 4096b261fc
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
7 changed files with 24 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
docs/en/modes/predict.md
View File

@ -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 |

1
docs/en/usage/cfg.md
View File

@ -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 |

7
tests/test_python.py
View File

@ -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('https://ultralytics.com/assets/decelera_portrait_min.mov', imgsz=160, tracker='botsort.yaml')
model.track(video_url, imgsz=160, tracker='bytetrack.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():

2
ultralytics/__init__.py
View File

@ -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

4
ultralytics/cfg/__init__.py
View File

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

1
ultralytics/cfg/default.yaml
View File

@ -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

18
ultralytics/engine/predictor.py
View File

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