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

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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
ultralytics/models/fastsam/predict.py
View File

@ -15,7 +15,6 @@ class FastSAMPredictor(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,
@ -32,22 +31,20 @@ class FastSAMPredictor(DetectionPredictor):
full_box[0][6:] = p[0][critical_iou_index][:, 6:]
p[0][critical_iou_index] = full_box
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

10
ultralytics/models/nas/predict.py
View File

@ -24,11 +24,11 @@ class NASPredictor(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

10
ultralytics/models/rtdetr/predict.py
View File

@ -28,6 +28,7 @@ class RTDETRPredictor(BasePredictor):
nd = preds[0].shape[-1]
bboxes, scores = preds[0].split((4, nd - 4), dim=-1)
results = []
is_list = isinstance(orig_imgs, list) # input images are a list, not a torch.Tensor
for i, bbox in enumerate(bboxes): # (300, 4)
bbox = ops.xywh2xyxy(bbox)
score, cls = scores[i].max(-1, keepdim=True) # (300, 1)
@ -35,14 +36,13 @@ class RTDETRPredictor(BasePredictor):
if self.args.classes is not None:
idx = (cls == torch.tensor(self.args.classes, device=cls.device)).any(1) & idx
pred = torch.cat([bbox, score, cls], dim=-1)[idx] # filter
orig_img = orig_imgs[i] if isinstance(orig_imgs, list) else orig_imgs
orig_img = orig_imgs[i] if is_list else orig_imgs
oh, ow = orig_img.shape[:2]
if not isinstance(orig_imgs, torch.Tensor):
if is_list:
pred[..., [0, 2]] *= ow
pred[..., [1, 3]] *= oh
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
def pre_transform(self, im):

8
ultralytics/models/sam/predict.py
View File

@ -318,8 +318,9 @@ class Predictor(BasePredictor):
pred_bboxes = preds[2] if self.segment_all else None
names = dict(enumerate(str(i) for i in range(len(pred_masks))))
results = []
is_list = isinstance(orig_imgs, list) # input images are a list, not a torch.Tensor
for i, masks in enumerate([pred_masks]):
orig_img = orig_imgs[i] if isinstance(orig_imgs, list) else orig_imgs
orig_img = orig_imgs[i] if is_list else orig_imgs
if pred_bboxes is not None:
pred_bboxes = ops.scale_boxes(img.shape[2:], pred_bboxes.float(), orig_img.shape, padding=False)
cls = torch.arange(len(pred_masks), dtype=torch.int32, device=pred_masks.device)
@ -327,9 +328,8 @@ class Predictor(BasePredictor):
masks = ops.scale_masks(masks[None].float(), orig_img.shape[:2], padding=False)[0]
masks = masks > self.model.mask_threshold # to bool
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=names, masks=masks, boxes=pred_bboxes))
img_path = self.batch[0][i]
results.append(Results(orig_img, path=img_path, names=names, masks=masks, boxes=pred_bboxes))
# Reset segment-all mode.
self.segment_all = False
return results

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