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

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