fix(LayoutPredictor): Ensure that the predicted bboxes are minmaxed i…

…nside the image boundaries

Update the LayoutPredictor unit test

Signed-off-by: Nikos Livathinos <[email protected]>

.gitignore

@@ -17,6 +17,9 @@ runs/*
 .DS_Store
 viz/
 
+# VSCode
+.vscode
+
 # VIM
 *.swp
 *.swo

docling_ibm_models/layoutmodel/layout_predictor.py

@@ -153,11 +153,15 @@ def predict(self, orig_img: Union[Image.Image, np.ndarray]) -> Iterable[dict]:
 
             # Check against threshold
             if score > self._threshold:
+                l = min(w, max(0, box[0]))
+                t = min(h, max(0, box[1]))
+                r = min(w, max(0, box[2]))
+                b = min(h, max(0, box[3]))
                 yield {
-                    "l": box[0],
-                    "t": box[1],
-                    "r": box[2],
-                    "b": box[3],
+                    "l": l,
+                    "t": t,
+                    "r": r,
+                    "b": b,
                     "label": label,
                     "confidence": score,
                 }

tests/test_layout_predictor.py

@@ -78,13 +78,23 @@ def test_layoutpredictor(init: dict):
     # Predict on the test image
     for img_fn in init["test_imgs"]:
         with Image.open(img_fn) as img:
+            w, h = img.size
             # Load images as PIL objects
             for i, pred in enumerate(lpredictor.predict(img)):
                 print("PIL pred: {}".format(pred))
+                assert pred["l"] >= 0 and pred["l"] <= w
+                assert pred["t"] >= 0 and pred["t"] <= h
+                assert pred["r"] >= 0 and pred["r"] <= w
+                assert pred["b"] >= 0 and pred["b"] <= h
+
             assert i + 1 == init["pred_bboxes"]
 
             # Load images as numpy arrays
             np_arr = np.asarray(img)
             for i, pred in enumerate(lpredictor.predict(np_arr)):
                 print("numpy pred: {}".format(pred))
+                assert pred["l"] >= 0 and pred["l"] <= w
+                assert pred["t"] >= 0 and pred["t"] <= h
+                assert pred["r"] >= 0 and pred["r"] <= w
+                assert pred["b"] >= 0 and pred["b"] <= h
             assert i + 1 == init["pred_bboxes"]