xdobro23/YOLOv8-16bit
1
0
Fork 0
Code Issues Pull Requests 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
7 changed files with 23 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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