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

Return processed outputs from predictor (#161)

Browse Source 
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>
This commit is contained in:
Ayush Chaurasia 2023-01-10 00:10:44 +05:30 committed by GitHub
parent cb4801888e
commit 6e5638c128
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
7 changed files with 23 additions and 10 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

9
tests/test_engine.py
View File

@ -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!"

3
ultralytics/yolo/engine/model.py
View File

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

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

@ -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()

5
ultralytics/yolo/utils/plotting.py
View File

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

1
ultralytics/yolo/v8/classify/predict.py
View File

@ -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)}, "

2
ultralytics/yolo/v8/detect/predict.py
View File

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

7
ultralytics/yolo/v8/segment/predict.py
View File

@ -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()