Return processed outputs from predictor (#161)

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>
Co-authored-by: Laughing-q <1185102784@qq.com>

tests/test_engine.py

@@ -24,7 +24,8 @@ def test_detect():
 
     # predictor
     pred = detect.DetectionPredictor(overrides={"imgsz": [640, 640]})
-    pred(source=SOURCE, model=trained_model)
+    p = pred(source=SOURCE, model="yolov8n.pt")
+    assert len(p) == 2, "predictor test failed"
 
     overrides["resume"] = trainer.last
     trainer = detect.DetectionTrainer(overrides=overrides)
@@ -54,7 +55,8 @@ def test_segment():
 
     # predictor
     pred = segment.SegmentationPredictor(overrides={"imgsz": [640, 640]})
-    pred(source=SOURCE, model=trained_model)
+    p = pred(source=SOURCE, model="yolov8n-seg.pt")
+    assert len(p) == 2, "predictor test failed"
 
     # test resume
     overrides["resume"] = trainer.last
@@ -91,4 +93,5 @@ def test_classify():
 
     # predictor
     pred = classify.ClassificationPredictor(overrides={"imgsz": [640, 640]})
-    pred(source=SOURCE, model=trained_model)
+    p = pred(source=SOURCE, model=trained_model)
+    assert len(p) == 2, "Predictor test failed!"

ultralytics/yolo/engine/model.py

@@ -121,11 +121,12 @@ class YOLO:
         overrides["conf"] = 0.25
         overrides.update(kwargs)
         overrides["mode"] = "predict"
+        overrides["save"] = kwargs.get("save", False)  # not save files by default
         predictor = self.PredictorClass(overrides=overrides)
 
         predictor.args.imgsz = check_imgsz(predictor.args.imgsz, min_dim=2)  # check image size
         predictor.setup(model=self.model, source=source)
-        predictor()
+        return predictor()
 
     @smart_inference_mode()
     def val(self, data=None, **kwargs):

ultralytics/yolo/engine/predictor.py

@@ -76,6 +76,7 @@ class BasePredictor:
         project = self.args.project or Path(SETTINGS['runs_dir']) / self.args.task
         name = self.args.name or f"{self.args.mode}"
         self.save_dir = increment_path(Path(project) / name, exist_ok=self.args.exist_ok)
+        if self.args.save:
             (self.save_dir / 'labels' if self.args.save_txt else self.save_dir).mkdir(parents=True, exist_ok=True)
         if self.args.conf is None:
             self.args.conf = 0.25  # default conf=0.25
@@ -149,7 +150,9 @@ class BasePredictor:
     def __call__(self, source=None, model=None):
         self.run_callbacks("on_predict_start")
         model = self.model if self.done_setup else self.setup(source, model)
+        model.eval()
         self.seen, self.windows, self.dt = 0, [], (ops.Profile(), ops.Profile(), ops.Profile())
+        self.all_outputs = []
         for batch in self.dataset:
             self.run_callbacks("on_predict_batch_start")
             path, im, im0s, vid_cap, s = batch
@@ -194,6 +197,7 @@ class BasePredictor:
             LOGGER.info(f"Results saved to {colorstr('bold', self.save_dir)}{s}")
 
         self.run_callbacks("on_predict_end")
+        return self.all_outputs
 
     def show(self, p):
         im0 = self.annotator.result()

ultralytics/yolo/utils/plotting.py

@@ -108,8 +108,9 @@ class Annotator:
         im_gpu = im_gpu.flip(dims=[0])  # flip channel
         im_gpu = im_gpu.permute(1, 2, 0).contiguous()  # shape(h,w,3)
         im_gpu = im_gpu * inv_alph_masks[-1] + mcs
-        im_mask = (im_gpu * 255).byte().cpu().numpy()
+        im_mask = (im_gpu * 255)
-        self.im[:] = im_mask if retina_masks else scale_image(im_gpu.shape, im_mask, self.im.shape)
+        im_mask_np = im_mask.byte().cpu().numpy()
+        self.im[:] = im_mask_np if retina_masks else scale_image(im_gpu.shape, im_mask_np, self.im.shape)
         if self.pil:
             # convert im back to PIL and update draw
             self.fromarray(self.im)

ultralytics/yolo/v8/classify/predict.py

@@ -37,6 +37,7 @@ class ClassificationPredictor(BasePredictor):
         self.annotator = self.get_annotator(im0)
 
         prob = preds[idx]
+        self.all_outputs.append(prob)
         # Print results
         top5i = prob.argsort(0, descending=True)[:5].tolist()  # top 5 indices
         log_string += f"{', '.join(f'{self.model.names[j]} {prob[j]:.2f}' for j in top5i)}, "

ultralytics/yolo/v8/detect/predict.py

@@ -51,12 +51,12 @@ class DetectionPredictor(BasePredictor):
         self.annotator = self.get_annotator(im0)
 
         det = preds[idx]
+        self.all_outputs.append(det)
         if len(det) == 0:
             return log_string
         for c in det[:, 5].unique():
             n = (det[:, 5] == c).sum()  # detections per class
             log_string += f"{n} {self.model.names[int(c)]}{'s' * (n > 1)}, "
-
         # write
         gn = torch.tensor(im0.shape)[[1, 0, 1, 0]]  # normalization gain whwh
         for *xyxy, conf, cls in reversed(det):

ultralytics/yolo/v8/segment/predict.py

@@ -58,7 +58,7 @@ class SegmentationPredictor(DetectionPredictor):
         mask = masks[idx]
         if self.args.save_txt:
             segments = [
-                ops.scale_segments(im0.shape if self.arg.retina_masks else im.shape[2:], x, im0.shape, normalize=True)
+                ops.scale_segments(im0.shape if self.args.retina_masks else im.shape[2:], x, im0.shape, normalize=True)
                 for x in reversed(ops.masks2segments(mask))]
 
         # Print results
@@ -73,6 +73,9 @@ class SegmentationPredictor(DetectionPredictor):
             im_gpu=torch.as_tensor(im0, dtype=torch.float16).to(self.device).permute(2, 0, 1).flip(0).contiguous() /
             255 if self.args.retina_masks else im[idx])
 
+        det = reversed(det[:, :6])
+        self.all_outputs.append([det, mask])
+
         # Write results
         for j, (*xyxy, conf, cls) in enumerate(reversed(det[:, :6])):
             if self.args.save_txt:  # Write to file
@@ -96,7 +99,7 @@ class SegmentationPredictor(DetectionPredictor):
 
 @hydra.main(version_base=None, config_path=str(DEFAULT_CONFIG.parent), config_name=DEFAULT_CONFIG.name)
 def predict(cfg):
-    cfg.model = cfg.model or "n.pt"
+    cfg.model = cfg.model or "yolov8n-seg.pt"
     cfg.imgsz = check_imgsz(cfg.imgsz, min_dim=2)  # check image size
     predictor = SegmentationPredictor(cfg)
     predictor()