From 6e5638c128e3914d5161a5ef4cc77c305e593b35 Mon Sep 17 00:00:00 2001 From: Ayush Chaurasia Date: Tue, 10 Jan 2023 00:10:44 +0530 Subject: [PATCH] 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 Co-authored-by: Laughing-q <1185102784@qq.com> --- tests/test_engine.py | 9 ++++++--- ultralytics/yolo/engine/model.py | 3 ++- ultralytics/yolo/engine/predictor.py | 6 +++++- ultralytics/yolo/utils/plotting.py | 5 +++-- ultralytics/yolo/v8/classify/predict.py | 1 + ultralytics/yolo/v8/detect/predict.py | 2 +- ultralytics/yolo/v8/segment/predict.py | 7 +++++-- 7 files changed, 23 insertions(+), 10 deletions(-) diff --git a/tests/test_engine.py b/tests/test_engine.py index 0e74dc2..d0bd83f 100644 --- a/tests/test_engine.py +++ b/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!" diff --git a/ultralytics/yolo/engine/model.py b/ultralytics/yolo/engine/model.py index 5a57022..cc727ad 100644 --- a/ultralytics/yolo/engine/model.py +++ b/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): diff --git a/ultralytics/yolo/engine/predictor.py b/ultralytics/yolo/engine/predictor.py index d654257..868a493 100644 --- a/ultralytics/yolo/engine/predictor.py +++ b/ultralytics/yolo/engine/predictor.py @@ -76,7 +76,8 @@ 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) - (self.save_dir / 'labels' if self.args.save_txt else self.save_dir).mkdir(parents=True, exist_ok=True) + 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 self.done_setup = False @@ -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() diff --git a/ultralytics/yolo/utils/plotting.py b/ultralytics/yolo/utils/plotting.py index 89128c9..e978806 100644 --- a/ultralytics/yolo/utils/plotting.py +++ b/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() - self.im[:] = im_mask if retina_masks else scale_image(im_gpu.shape, im_mask, self.im.shape) + im_mask = (im_gpu * 255) + 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) diff --git a/ultralytics/yolo/v8/classify/predict.py b/ultralytics/yolo/v8/classify/predict.py index 47501cc..04d617b 100644 --- a/ultralytics/yolo/v8/classify/predict.py +++ b/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)}, " diff --git a/ultralytics/yolo/v8/detect/predict.py b/ultralytics/yolo/v8/detect/predict.py index 9736dca..f7544b8 100644 --- a/ultralytics/yolo/v8/detect/predict.py +++ b/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): diff --git a/ultralytics/yolo/v8/segment/predict.py b/ultralytics/yolo/v8/segment/predict.py index 1705dba..810cd07 100644 --- a/ultralytics/yolo/v8/segment/predict.py +++ b/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()