Chiebot-Mirror
/
ultralytics
mirror of https://github.com/ultralytics/ultralytics.git
9
0
Fork 1
Code Issues Packages Projects Releases Wiki Activity

Simplify postprocessing methods (#4497)

Browse Source
pull/4502/head
Glenn Jocher 1 year ago 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
Whitespace
Show all changes Ignore whitespace when comparing lines Ignore changes in amount of whitespace Ignore changes in whitespace at EOL
Split View
Diff Options
Show Stats Download Patch File Download Diff File
  1. 19
      ultralytics/models/fastsam/predict.py
  2. 10
      ultralytics/models/nas/predict.py
  3. 10
      ultralytics/models/rtdetr/predict.py
  4. 8
      ultralytics/models/sam/predict.py
  5. 9
      ultralytics/models/yolo/classify/predict.py
  6. 12
      ultralytics/models/yolo/detect/predict.py
  7. 17
      ultralytics/models/yolo/pose/predict.py
  8. 21
      ultralytics/models/yolo/segment/predict.py

19
ultralytics/models/fastsam/predict.py
Unescape 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
Unescape 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
Unescape 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
Unescape 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
Unescape 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
Unescape 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
Unescape 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
Unescape 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

Write Preview
Loading…
Cancel
Save