fix: fix probability data type (#391)

This PR fixes the data type for probabilities from int to float.
Unstructured-IO · Oct 11, 2024 · 4431fe5 · 4431fe5
1 parent e411862
commit 4431fe5
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 2 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,3 +1,7 @@
+## 0.7.40
+
+* fix: store probabilities with `float` data type instead of `int`
+
 ## 0.7.39
 
 * fix: Correctly assign mutable default value to variable in `LayoutElements` class

diff --git a/unstructured_inference/__version__.py b/unstructured_inference/__version__.py
@@ -1 +1 @@
-__version__ = "0.7.39"  # pragma: no cover
+__version__ = "0.7.40"  # pragma: no cover
diff --git a/unstructured_inference/models/yolox.py b/unstructured_inference/models/yolox.py
@@ -137,7 +137,7 @@ def image_processing(
 
         return LayoutElements(
             element_coords=sorted_dets[:, :4],
-            element_probs=sorted_dets[:, 4].astype(int),
+            element_probs=sorted_dets[:, 4].astype(float),
             element_class_ids=sorted_dets[:, 5].astype(int),
             element_class_id_map=self.layout_classes,
             source=Source.YOLOX,