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

Replace enumerate with zip in models/yolo (#14599)

Browse source
This commit is contained in:
Kayzwer 2024-07-23 01:52:43 +08:00 • committed by GitHub
parent 22a44d82c5
commit db82d1c6ae
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 4 additions and 12 deletions
Download patch file Download diff file Expand all files Collapse all files

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

@ -54,8 +54,6 @@ class ClassificationPredictor(BasePredictor):
orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)
results = []
for i, pred in enumerate(preds):
orig_img = orig_imgs[i]
img_path = self.batch[0][i]
for pred, orig_img, img_path in zip(preds, orig_imgs, self.batch[0]):
results.append(Results(orig_img, path=img_path, names=self.model.names, probs=pred))
return results

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

@ -35,9 +35,7 @@ class DetectionPredictor(BasePredictor):
orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)
results = []
for i, pred in enumerate(preds):
orig_img = orig_imgs[i]
for pred, orig_img, img_path in zip(preds, orig_imgs, self.batch[0]):
pred[:, :4] = ops.scale_boxes(img.shape[2:], pred[:, :4], orig_img.shape)
img_path = self.batch[0][i]
results.append(Results(orig_img, path=img_path, names=self.model.names, boxes=pred))
return results

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

@ -46,12 +46,10 @@ class PosePredictor(DetectionPredictor):
orig_imgs = ops.convert_torch2numpy_batch(orig_imgs)
results = []
for i, pred in enumerate(preds):
orig_img = orig_imgs[i]
for pred, orig_img, img_path in zip(preds, orig_imgs, self.batch[0]):
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, orig_img.shape)
img_path = self.batch[0][i]
results.append(
Results(orig_img, path=img_path, names=self.model.names, boxes=pred[:, :6], keypoints=pred_kpts)
)

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

@ -42,9 +42,7 @@ class SegmentationPredictor(DetectionPredictor):
results = []
proto = preds[1][-1] if isinstance(preds[1], tuple) else preds[1] # tuple if PyTorch model or array if exported
for i, pred in enumerate(p):
orig_img = orig_imgs[i]
img_path = self.batch[0][i]
for i, (pred, orig_img, img_path) in enumerate(zip(p, orig_imgs, self.batch[0])):
if not len(pred): # save empty boxes
masks = None
elif self.args.retina_masks: