debug: Add prints to find out where predictions fail

Signed-off-by: Christoph Auer <cau@zurich.ibm.com>

Author: cau-git

docling/models/stages/ocr/nemotron_ocr_model.py

@@ -25,7 +25,7 @@
 
 
 class NemotronOcrPrediction(TypedDict):
-    """Exact prediction schema returned by `nemotron_ocr` 1.0.1."""
+    """Exact prediction schema returned by `nemotron_ocr`."""
 
     text: str
     confidence: float
@@ -119,7 +119,7 @@ def _prediction_to_cell(
         image_height: int,
         scale: int,
     ) -> TextCell:
-        # `nemotron_ocr` 1.0.1 returns normalized `left/right` and an inverted
+        # `nemotron_ocr` returns normalized `left/right` and an inverted
         # pair `lower/upper`, where `lower` is the top Y and `upper` is the
         # bottom Y in image coordinates.
         left = (prediction["left"] * image_width) / scale + ocr_rect.l
@@ -160,8 +160,14 @@ def __call__(
                 with TimeRecorder(conv_res, "ocr"):
                     ocr_rects = self.get_ocr_rects(page)
 
+                    print(
+                        "[nemotron-debug] "
+                        f"page={page.page_no} rect_count={len(ocr_rects)} "
+                        f"rects={[rect.as_tuple() for rect in ocr_rects]}"
+                    )
+
                     all_ocr_cells = []
-                    for ocr_rect in ocr_rects:
+                    for crop_index, ocr_rect in enumerate(ocr_rects):
                         if ocr_rect.area() == 0:
                             continue
 
@@ -171,15 +177,37 @@ def __call__(
                         image_width, image_height = high_res_image.size
                         image_array = numpy.array(high_res_image)
 
-                        raw_predictions = cast(
-                            Sequence[NemotronOcrPrediction],
-                            self.reader(
-                                image_array,
-                                merge_level=self.options.merge_level,
-                                visualize=False,
-                            ),
+                        print(
+                            "[nemotron-debug] "
+                            f"page={page.page_no} crop={crop_index} "
+                            f"rect={ocr_rect.as_tuple()} "
+                            f"size={image_width}x{image_height} "
+                            f"shape={image_array.shape} "
+                            f"dtype={image_array.dtype} "
+                            f"contiguous={image_array.flags.c_contiguous}"
                         )
 
+                        try:
+                            raw_predictions = cast(
+                                Sequence[NemotronOcrPrediction],
+                                self.reader(
+                                    image_array,
+                                    merge_level=self.options.merge_level,
+                                    visualize=False,
+                                ),
+                            )
+                        except Exception:
+                            print(
+                                "[nemotron-debug] "
+                                f"FAILED page={page.page_no} crop={crop_index} "
+                                f"rect={ocr_rect.as_tuple()} "
+                                f"size={image_width}x{image_height} "
+                                f"shape={image_array.shape} "
+                                f"dtype={image_array.dtype} "
+                                f"contiguous={image_array.flags.c_contiguous}"
+                            )
+                            raise
+
                         del high_res_image
                         del image_array