diff --git a/CHANGELOG.md b/CHANGELOG.md index c353c99b..827aee7f 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/test_unstructured_inference/models/test_supergradients.py b/test_unstructured_inference/models/test_supergradients.py index eecd1051..d56df7d0 100644 --- a/test_unstructured_inference/models/test_supergradients.py +++ b/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), ) diff --git a/unstructured_inference/__version__.py b/unstructured_inference/__version__.py index d4341911..11245184 100644 --- a/unstructured_inference/__version__.py +++ b/unstructured_inference/__version__.py @@ -1 +1 @@ -__version__ = "0.7.15" # pragma: no cover +__version__ = "0.7.16" # pragma: no cover diff --git a/unstructured_inference/models/super_gradients.py b/unstructured_inference/models/super_gradients.py index 0719d7c3..6d9a25fa 100644 --- a/unstructured_inference/models/super_gradients.py +++ b/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))