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):