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

Return save_dir with prediction results (#3240)

Browse Source
This commit is contained in:
Glenn Jocher 2023-06-18 02:34:31 +02:00 committed by GitHub
parent 15c90bd404
commit 9c44560ec4
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 6 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

7
ultralytics/yolo/engine/predictor.py
View File

@ -102,6 +102,8 @@ class BasePredictor:
self.data_path = None
self.source_type = None
self.batch = None
self.results = None
self.transforms = None
self.callbacks = _callbacks or callbacks.get_default_callbacks()
callbacks.add_integration_callbacks(self)
@ -260,12 +262,13 @@ class BasePredictor:
if self.args.verbose or self.args.save or self.args.save_txt or self.args.show:
s += self.write_results(i, self.results, (p, im, im0))
if self.args.save or self.args.save_txt:
self.results[i].save_dir = self.save_dir.__str__()
if self.args.show and self.plotted_img is not None:
self.show(p)
if self.args.save and self.plotted_img is not None:
self.save_preds(vid_cap, i, str(self.save_dir / p.name))
self.run_callbacks('on_predict_batch_end')
yield from self.results

1
ultralytics/yolo/engine/results.py
View File

@ -101,6 +101,7 @@ class Results(SimpleClass):
self.speed = {'preprocess': None, 'inference': None, 'postprocess': None} # milliseconds per image
self.names = names
self.path = path
self.save_dir = None
self._keys = ('boxes', 'masks', 'probs', 'keypoints')
def __getitem__(self, idx):