fix[layoutmodel]: Fix type hints in LayoutPredictor::predict()

Signed-off-by: Nikos Livathinos <[email protected]>
DS4SD · Jul 18, 2024 · d6b8f0d · d6b8f0d
1 parent cbfcd4e
commit d6b8f0d
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/docling_ibm_models/layoutmodel/layout_predictor.py b/docling_ibm_models/layoutmodel/layout_predictor.py
@@ -118,7 +118,7 @@ def info(self) -> dict:
         }
         return info
 
-    def predict(self, orig_img: Union[Image, np.array]) -> Iterable[dict]:
+    def predict(self, orig_img: Union[Image.Image, np.ndarray]) -> Iterable[dict]:
         r"""
         Predict bounding boxes for a given image.
         The origin (0, 0) is the top-left corner and the predicted bbox coords are provided as: