From 9c44560ec4bfd5f6632a748add7fdc1ce5c6cdfa Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 18 Jun 2023 02:34:31 +0200 Subject: [PATCH] Return `save_dir` with prediction results (#3240) --- ultralytics/yolo/engine/predictor.py | 7 +++++-- ultralytics/yolo/engine/results.py | 1 + 2 files changed, 6 insertions(+), 2 deletions(-) diff --git a/ultralytics/yolo/engine/predictor.py b/ultralytics/yolo/engine/predictor.py index 5a37878..b71785c 100644 --- a/ultralytics/yolo/engine/predictor.py +++ b/ultralytics/yolo/engine/predictor.py @@ -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 diff --git a/ultralytics/yolo/engine/results.py b/ultralytics/yolo/engine/results.py index 3d4ff0c..68e0de2 100644 --- a/ultralytics/yolo/engine/results.py +++ b/ultralytics/yolo/engine/results.py @@ -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):