fix(back): Fix converting BBox coords to float

pixano · Dec 8, 2023 · b1e936a · b1e936a
1 parent 351dced
commit b1e936a
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 6 deletions.
diff --git a/pixano_inference/pytorch_models/maskrcnnv2.py b/pixano_inference/pytorch_models/maskrcnnv2.py
@@ -128,8 +128,8 @@ def preannotate(
                             "item_id": batch["id"][x].as_py(),
                             "view_id": view,
                             "bbox": BBox.from_xyxy(
-                                [float(coord) for coord in output["boxes"][i]],
-                                confidence=float(output["scores"][i]),
+                                [coord.item() for coord in output["boxes"][i]],
+                                confidence=output["scores"][i].item(),
                             )
                             .normalize(h, w)
                             .to_dict(),

diff --git a/pixano_inference/pytorch_models/yolov5.py b/pixano_inference/pytorch_models/yolov5.py
@@ -98,8 +98,8 @@ def preannotate(
                             "item_id": batch["id"][x].as_py(),
                             "view_id": view,
                             "bbox": BBox.from_xyxy(
-                                [float(coord) for coord in pred[0:4]],
-                                confidence=float(pred[4]),
+                                [coord.item() for coord in pred[0:4]],
+                                confidence=pred[4].item(),
                             )
                             .normalize(h, w)
                             .to_dict(),

diff --git a/pixano_inference/segment_anything/sam.py b/pixano_inference/segment_anything/sam.py
@@ -114,8 +114,8 @@ def preannotate(
                             "item_id": batch["id"][x].as_py(),
                             "view_id": view,
                             "bbox": BBox.from_xywh(
-                                [float(coord) for coord in output[i]["bbox"]],
-                                confidence=float(output[i]["predicted_iou"]),
+                                [coord.item() for coord in output[i]["bbox"]],
+                                confidence=output[i]["predicted_iou"].item(),
                             )
                             .normalize(h, w)
                             .to_dict(),