HongChuang/ultralytics-ascend
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions 3

Simplify postprocessing methods (#4497)

Browse source
This commit is contained in:
Glenn Jocher 2023-08-22 19:01:33 +02:00 • committed by GitHub
parent 6da8f7f51e
commit b890e1c937
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
8 changed files with 45 additions and 61 deletions
Download patch file Download diff file Expand all files Collapse all files

9
ultralytics/models/yolo/classify/predict.py
View file

@ -39,10 +39,9 @@ class ClassificationPredictor(BasePredictor):
def postprocess(self, preds, img, orig_imgs):
"""Post-processes predictions to return Results objects."""
results = []
is_list = isinstance(orig_imgs, list) # input images are a list, not a torch.Tensor
for i, pred in enumerate(preds):
orig_img = orig_imgs[i] if isinstance(orig_imgs, list) else orig_imgs
path = self.batch[0]
img_path = path[i] if isinstance(path, list) else path
results.append(Results(orig_img=orig_img, path=img_path, names=self.model.names, probs=pred))
orig_img = orig_imgs[i] if is_list else orig_imgs
img_path = self.batch[0][i]
results.append(Results(orig_img, path=img_path, names=self.model.names, probs=pred))
return results

12
ultralytics/models/yolo/detect/predict.py
View file

@ -1,7 +1,5 @@
# Ultralytics YOLO 🚀, AGPL-3.0 license
import torch
from ultralytics.engine.predictor import BasePredictor
from ultralytics.engine.results import Results
from ultralytics.utils import ops
@ -32,11 +30,11 @@ class DetectionPredictor(BasePredictor):
classes=self.args.classes)
results = []
is_list = isinstance(orig_imgs, list) # input images are a list, not a torch.Tensor
for i, pred in enumerate(preds):
orig_img = orig_imgs[i] if isinstance(orig_imgs, list) else orig_imgs
if not isinstance(orig_imgs, torch.Tensor):
orig_img = orig_imgs[i] if is_list else orig_imgs
if is_list:
pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
path = self.batch[0]
img_path = path[i] if isinstance(path, list) else path
results.append(Results(orig_img=orig_img, path=img_path, names=self.model.names, boxes=pred))
img_path = self.batch[0][i]
results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred))
return results

17
ultralytics/models/yolo/pose/predict.py
View file

@ -38,18 +38,13 @@ class PosePredictor(DetectionPredictor):
nc=len(self.model.names))
results = []
is_list = isinstance(orig_imgs, list) # input images are a list, not a torch.Tensor
for i, pred in enumerate(preds):
orig_img = orig_imgs[i] if isinstance(orig_imgs, list) else orig_imgs
shape = orig_img.shape
pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], shape).round()
orig_img = orig_imgs[i] if is_list else orig_imgs
pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape).round()
pred_kpts = pred[:, 6:].view(len(pred), *self.model.kpt_shape) if len(pred) else pred[:, 6:]
pred_kpts = ops.scale_coords(img.shape[2:], pred_kpts, shape)
path = self.batch[0]
img_path = path[i] if isinstance(path, list) else path
pred_kpts = ops.scale_coords(img.shape[2:], pred_kpts, orig_img.shape)
img_path = self.batch[0][i]
results.append(
Results(orig_img=orig_img,
path=img_path,
names=self.model.names,
boxes=pred[:, :6],
keypoints=pred_kpts))
Results(orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6], keypoints=pred_kpts))
return results

21
ultralytics/models/yolo/segment/predict.py
View file

@ -1,7 +1,5 @@
# Ultralytics YOLO 🚀, AGPL-3.0 license
import torch
from ultralytics.engine.results import Results
from ultralytics.models.yolo.detect.predict import DetectionPredictor
from ultralytics.utils import DEFAULT_CFG, ops
@ -27,7 +25,6 @@ class SegmentationPredictor(DetectionPredictor):
self.args.task = 'segment'
def postprocess(self, preds, img, orig_imgs):
"""TODO: filter by classes."""
p = ops.non_max_suppression(preds[0],
self.args.conf,
self.args.iou,
@ -36,22 +33,20 @@ class SegmentationPredictor(DetectionPredictor):
nc=len(self.model.names),
classes=self.args.classes)
results = []
is_list = isinstance(orig_imgs, list) # input images are a list, not a torch.Tensor
proto = preds[1][-1] if len(preds[1]) == 3 else preds[1] # second output is len 3 if pt, but only 1 if exported
for i, pred in enumerate(p):
orig_img = orig_imgs[i] if isinstance(orig_imgs, list) else orig_imgs
path = self.batch[0]
img_path = path[i] if isinstance(path, list) else path
orig_img = orig_imgs[i] if is_list else orig_imgs
img_path = self.batch[0][i]
if not len(pred): # save empty boxes
results.append(Results(orig_img=orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6]))
continue
if self.args.retina_masks:
if not isinstance(orig_imgs, torch.Tensor):
masks = None
elif self.args.retina_masks:
if is_list:
pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
masks = ops.process_mask_native(proto[i], pred[:, 6:], pred[:, :4], orig_img.shape[:2]) # HWC
else:
masks = ops.process_mask(proto[i], pred[:, 6:], pred[:, :4], img.shape[2:], upsample=True) # HWC
if not isinstance(orig_imgs, torch.Tensor):
if is_list:
pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
results.append(
Results(orig_img=orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6], masks=masks))
results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6], masks=masks))
return results