fix: 🐞 add if condition for check existince of 'pred_masks' for detectron2

Signed-off-by: Onuralp SEZER <thunderbirdtr@gmail.com>
This commit is contained in:
Onuralp SEZER 2024-06-25 15:59:49 +03:00
parent 3df33f0f9d
commit f66444bc91
No known key found for this signature in database
GPG Key ID: CF0835DFDF14CA38
1 changed files with 3 additions and 1 deletions

View File

@ -540,7 +540,9 @@ class Detections:
return cls(
xyxy=detectron2_results["instances"].pred_boxes.tensor.cpu().numpy(),
confidence=detectron2_results["instances"].scores.cpu().numpy(),
mask=detectron2_results["instances"].pred_masks.detach().cpu().numpy(),
mask=detectron2_results["instances"].pred_masks.cpu().numpy()
if hasattr(detectron2_results["instances"], "pred_masks")
else None,
class_id=detectron2_results["instances"]
.pred_classes.cpu()
.numpy()