Support FastSAM directory inference and plot (#4634)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>
This commit is contained in:
Gezhi Zhang 2023-08-30 20:19:36 +08:00 committed by GitHub
parent 620335de27
commit 8596ee241f
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 136 additions and 120 deletions

View file

@ -22,7 +22,7 @@ class FastSAMPredictor(DetectionPredictor):
max_det=self.args.max_det,
nc=len(self.model.names),
classes=self.args.classes)
full_box = torch.zeros_like(p[0][0])
full_box = torch.zeros(p[0].shape[1])
full_box[2], full_box[3], full_box[4], full_box[6:] = img.shape[3], img.shape[2], 1.0, 1.0
full_box = full_box.view(1, -1)
critical_iou_index = bbox_iou(full_box[0][:4], p[0][:, :4], iou_thres=0.9, image_shape=img.shape[2:])