Update mdr_pdf_parser.py

mdr_pdf_parser.py

@@ -1443,7 +1443,7 @@ class MDROcrEngine:
   def _get_system(self) -> _MDR_TextSystem | None:
     if self._text_system is None:
         paths = {k: str(Path(self._model_dir)/Path(*p)) for k,p in _MDR_OCR_MODELS.items()}
-        self._onnx_params = _MDR_ONNXParams(use_gpu=(self._device=="cuda"), det_model_dir=paths["det"], cls_model_dir=paths["cls"], rec_model_dir=paths["rec"], rec_char_dict_path=paths["keys"])
+        self._onnx_params = _MDR_ONNXParams(use_gpu=(self._device=="cpu"), det_model_dir=paths["det"], cls_model_dir=paths["cls"], rec_model_dir=paths["rec"], rec_char_dict_path=paths["keys"])
         try: self._text_system = _MDR_TextSystem(self._onnx_params); print(f"MDR OCR System initialized.")
         except Exception as e: print(f"ERROR initializing MDR OCR System: {e}"); self._text_system = None
     return self._text_system
@@ -1688,7 +1688,7 @@ class MDRLatexExtractor:
       if not cp.exists():
           print(f"Warn: MDR LaTeX config not found {self._model_path}")
       try:
-          args = Munch({"config": str(cp), "checkpoint": str(wp), "device": self._device, "no_cuda": self._device == "cpu", "no_resize": False, "temperature": 0.0})
+          args = Munch({"config": str(cp), "checkpoint": str(wp), "device": self._device, "no_cuda": self._device == "cuda", "no_resize": False, "temperature": 0.0})
           self._model = LatexOCR(args)
           print(f"MDR LaTeX loaded on {self._device}.")
       except Exception as e:
@@ -1712,7 +1712,7 @@ class MDRTableParser:
   def __init__(self, device: Literal["cpu", "cuda"], model_path: str):
     self._model: Any | None = None; self._model_path = mdr_ensure_directory(model_path)
     self._device = device if torch.cuda.is_available() and device=="cuda" else "cpu"
-    self._disabled = self._device == "cpu"
+    self._disabled = self._device == "cuda"
     if self._disabled: print("Warning: MDR Table parsing requires CUDA. Disabled.")
 
   def parse_table_image(self, image: Image, format: MDRTableLayoutParsedFormat) -> str | None:
@@ -1950,10 +1950,21 @@ class MDRExtractionEngine:
         try:
             raw_layouts = list(self._run_yolo_detection(optimizer.image, yolo))
             print(f"  Engine: {len(raw_layouts)} raw layouts found.")
-        except Exception as e:
-            print(f"  Engine: YOLO error: {e}")
+        except Exception:
+            import traceback, sys
+            traceback.print_exc(file=sys.stderr)
     print("  Engine: Matching fragments...")
     layouts = self._match_fragments_to_layouts(frags, raw_layouts)
+    if not layouts and frags:
+        # treat the whole page as one plain-text layout
+        page_rect = MDRRectangle(
+            lt=(0, 0), rt=(optimizer.image.width, 0),
+            lb=(0, optimizer.image.height), rb=(optimizer.image.width, optimizer.image.height)
+        )
+        dummy = MDRPlainLayoutElement(
+            cls=MDRLayoutClass.PLAIN_TEXT, rect=page_rect, fragments=frags.copy()
+        )
+        layouts.append(dummy)
     print("  Engine: Removing overlaps...")
     layouts = mdr_remove_overlap_layouts(layouts)
     print(f"  Engine: {len(layouts)} layouts after overlap removal.")
@@ -1974,27 +1985,37 @@ class MDRExtractionEngine:
     print("  Engine: Analysis complete.")
     return MDRExtractionResult(rotation=optimizer.rotation, layouts=layouts, extracted_image=image, adjusted_image=optimizer.adjusted_image)
 
-  def _run_yolo_detection(self, img: Image, yolo: YOLOv10) -> Generator[MDRLayoutElement, None, None]:
-    img_rgb = img.convert("RGB")
-    res = yolo.predict(source=img_rgb, imgsz=1024, conf=0.2, device=self._device, verbose=False)
-    if not res or not hasattr(res[0], 'boxes') or res[0].boxes is None:
-        return
-    boxes = res[0].boxes
-    for cls_id_t, xyxy_t in zip(boxes.cls, boxes.xyxy):
-        cls_id = int(cls_id_t.item())
-        try:
-            cls = MDRLayoutClass(cls_id)
-        except ValueError:
-            continue
-        x1, y1, x2, y2 = [c.item() for c in xyxy_t]
-        rect = MDRRectangle(lt=(x1, y1), rt=(x2, y1), lb=(x1, y2), rb=(x2, y2))
-        if rect.is_valid and rect.area > 10:
-            if cls == MDRLayoutClass.TABLE:
-                yield MDRTableLayoutElement(cls=cls, rect=rect, fragments=[], parsed=None)
-            elif cls == MDRLayoutClass.ISOLATE_FORMULA:
-                yield MDRFormulaLayoutElement(cls=cls, rect=rect, fragments=[], latex=None)
-            elif cls in MDRPlainLayoutElement.__annotations__['cls'].__args__:
-                yield MDRPlainLayoutElement(cls=cls, rect=rect, fragments=[])
+  def _run_yolo_detection(self, img: Image, yolo: YOLOv10):
+      img_rgb = img.convert("RGB")
+      res = yolo.predict(source=img_rgb, imgsz=1024, conf=0.20,
+                         device=self._device, verbose=False)
+
+      if not res or not res[0].boxes:
+          return
+
+      plain_classes: set[MDRLayoutClass] = {
+          MDRLayoutClass.TITLE,
+          MDRLayoutClass.PLAIN_TEXT,
+          MDRLayoutClass.ABANDON,
+          MDRLayoutClass.FIGURE_CAPTION,
+          MDRLayoutClass.TABLE_CAPTION,
+          MDRLayoutClass.TABLE_FOOTNOTE,
+          MDRLayoutClass.FORMULA_CAPTION,
+      }
+
+      for cls_id_t, xyxy_t in zip(res[0].boxes.cls, res[0].boxes.xyxy):
+          cls = MDRLayoutClass(int(cls_id_t))
+          x1, y1, x2, y2 = map(float, xyxy_t)
+          rect = MDRRectangle((x1, y1), (x2, y1), (x1, y2), (x2, y2))
+          if rect.area < 10:
+              continue
+
+          if cls == MDRLayoutClass.TABLE:
+              yield MDRTableLayoutElement(rect=rect, fragments=[], parsed=None)
+          elif cls == MDRLayoutClass.ISOLATE_FORMULA:
+              yield MDRFormulaLayoutElement(rect=rect, fragments=[], latex=None)
+          elif cls in plain_classes:
+              yield MDRPlainLayoutElement(cls=cls, rect=rect, fragments=[])
 
   def _match_fragments_to_layouts(self, frags: list[MDROcrFragment], layouts: list[MDRLayoutElement]) -> list[MDRLayoutElement]:
     if not frags or not layouts:
@@ -2540,7 +2561,7 @@ if __name__ == '__main__':
     MDR_DEBUG_DIRECTORY = "./mdr_debug_output"
 
     # Specify device ('cuda' or 'cpu').
-    MDR_DEVICE = "cpu"
+    MDR_DEVICE = "cuda"
 
     # Specify desired table format
     MDR_TABLE_FORMAT = MDRExtractedTableFormat.MARKDOWN
@@ -2637,4 +2658,4 @@ if __name__ == '__main__':
         print(f"\nFATAL ERROR during processing: {e}")
         import traceback
         traceback.print_exc()
-        exit(1)
+        exit(1)