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

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