fix bug which allows label map from initialize params to be used from… (#293)

… reading in a json file.

Author: pravin-unstructured

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.7.16
+
+* bug: Allow supplied ONNX models to use label_map dictionary from json file
+
 ## 0.7.15
 
 * enhancement: Enable env variables for model definition

test_unstructured_inference/models/test_supergradients.py

@@ -11,27 +11,27 @@ def test_supergradients_model():
     model.initialize(
         model_path=model_path,
         label_map={
-            0: "Picture",
-            1: "Caption",
-            2: "Text",
-            3: "Formula",
-            4: "Page number",
-            5: "Address",
-            6: "Footer",
-            7: "Subheadline",
-            8: "Chart",
-            9: "Metadata",
-            10: "Title",
-            11: "Misc",
-            12: "Header",
-            13: "Table",
-            14: "Headline",
-            15: "List-item",
-            16: "List",
-            17: "Author",
-            18: "Value",
-            19: "Link",
-            20: "Field-Name",
+            "0": "Picture",
+            "1": "Caption",
+            "2": "Text",
+            "3": "Formula",
+            "4": "Page number",
+            "5": "Address",
+            "6": "Footer",
+            "7": "Subheadline",
+            "8": "Chart",
+            "9": "Metadata",
+            "10": "Title",
+            "11": "Misc",
+            "12": "Header",
+            "13": "Table",
+            "14": "Headline",
+            "15": "List-item",
+            "16": "List",
+            "17": "Author",
+            "18": "Value",
+            "19": "Link",
+            "20": "Field-Name",
         },
         input_shape=(1024, 1024),
     )

unstructured_inference/__version__.py

@@ -1 +1 @@
-__version__ = "0.7.15"  # pragma: no cover
+__version__ = "0.7.16"  # pragma: no cover

unstructured_inference/models/super_gradients.py

@@ -69,15 +69,15 @@ def image_processing(
                 class_id = pred_classes[image_index, i]
                 prob = pred_scores[image_index, i]
                 x1, y1, x2, y2 = pred_boxes[image_index, i]
-                detected_class = self.layout_classes[int(class_id)]
+                detected_class = self.layout_classes[str(class_id)]
                 region = LayoutElement.from_coords(
-                    x1,
-                    y1,
-                    x2,
-                    y2,
+                    float(x1),
+                    float(y1),
+                    float(x2),
+                    float(y2),
                     text=None,
                     type=detected_class,
-                    prob=prob,
+                    prob=float(prob),
                     source=Source.SUPER_GRADIENTS,
                 )
                 regions.append(cast(LayoutElement, region))