Update mdr_pdf_parser.py

mdr_pdf_parser.py

@@ -2589,8 +2589,8 @@ class MDRExtractionEngine:
     def _get_yolo_model(self) -> YOLOv10 | None:
         """Loads the YOLOv10 layout detection model using hf_hub_download."""
         if self._yolo is None and YOLOv10 is not None:
-            repo_id = "juliozhao/DocLayout-YOLO-DocStructBench"
-            filename = "doclayout_yolo_docstructbench_imgsz1024.pt"
+            repo_id = "hantian/yolo-doclaynet"
+            filename = "yolov10b-doclaynet.pt"
             # Use a subdirectory within the main model dir for YOLO cache via HF Hub
             yolo_cache_dir = Path(self._model_dir) / "yolo_hf_cache"
             mdr_ensure_directory(str(yolo_cache_dir))  # Ensure cache dir exists
@@ -2684,37 +2684,193 @@ class MDRExtractionEngine:
         return MDRExtractionResult(rotation=optimizer.rotation, layouts=layouts, extracted_image=image,
                                    adjusted_image=optimizer.adjusted_image)
 
-    def _run_yolo_detection(self, img: Image, yolo: YOLOv10):
+        # In class MDRExtractionEngine:
+
+    def _run_yolo_detection(self, img: Image, yolo: Any):  # yolo can be doclayout_yolo.YOLOv10 or ultralytics.YOLO
         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:
+        # Standard predict call
+        # The conf threshold might need adjustment based on the new model's performance
+        # For DocLayNet, 'Text' is often a high-confidence class.
+        res_list = yolo.predict(source=img_rgb, imgsz=1024, conf=0.25,
+                                # Slightly higher conf for potentially better precision
+                                device=self._device, verbose=False)
+
+        if not res_list or not hasattr(res_list[0], 'boxes') or res_list[0].boxes is None:
+            print("  Engine: YOLO detection returned no results or no boxes.")
             return
 
-        plain_classes: set[MDRLayoutClass] = {
+        results = res_list[0]  # Get the first (and usually only) result object
+
+        # --- Determine Class Mapping ---
+        # This mapping needs to be verified against the actual model's output.
+        # The hantian/yolo-doclaynet model card or its files might specify this.
+        # Common DocLayNet class order (example, VERIFY THIS):
+        # 0: Caption, 1: Footnote, 2: Formula, 3: List-item, 4: Page-footer,
+        # 5: Page-header, 6: Picture, 7: Section-header, 8: Table, 9: Text, 10: Title
+
+        # Let's try to get names from the model directly if possible
+        model_class_names = {}
+        if hasattr(results, 'names') and isinstance(results.names, dict):
+            model_class_names = results.names  # results.names is usually {id: name}
+            print(f"  Engine: YOLO model class names: {model_class_names}")
+        else:
+            print(
+                "  Engine: Warning - Could not automatically get class names from YOLO model. Using predefined fallback mapping.")
+            # Fallback predefined mapping (MUST BE VERIFIED FOR hantian/yolo-doclaynet)
+            # This is a GUESS based on common DocLayNet order.
+            # You MUST verify this by inspecting the model's config or output.
+            _doclaynet_names_fallback = ['Caption', 'Footnote', 'Formula', 'List-item', 'Page-footer', 'Page-header',
+                                         'Picture', 'Section-header', 'Table', 'Text', 'Title']
+            model_class_names = {i: name for i, name in enumerate(_doclaynet_names_fallback)}
+
+        # Define your mapping from DocLayNet names (or indices if names are not available) to MDRLayoutClass
+        # This is crucial and specific to the new model's output classes.
+        doclaynet_to_mdr_map = {
+            model_class_names.get(k): v for k, v in {
+                # Map by string name if model_class_names is populated correctly
+                'Text': MDRLayoutClass.PLAIN_TEXT,
+                'Title': MDRLayoutClass.TITLE,
+                'Section-header': MDRLayoutClass.TITLE,  # Or a new MDRLayoutClass if needed
+                'List-item': MDRLayoutClass.PLAIN_TEXT,  # Treat list items as plain text
+                'Table': MDRLayoutClass.TABLE,
+                'Picture': MDRLayoutClass.FIGURE,
+                'Formula': MDRLayoutClass.ISOLATE_FORMULA,
+                'Caption': MDRLayoutClass.FIGURE_CAPTION,  # Or TABLE_CAPTION, needs context
+                'Footnote': MDRLayoutClass.TABLE_FOOTNOTE,  # Or a general footnote class
+                'Page-header': MDRLayoutClass.ABANDON,  # Often headers/footers are ignored
+                'Page-footer': MDRLayoutClass.ABANDON,
+            }.items() if k in model_class_names.values()  # Ensure key exists
+        }
+
+        # If mapping by string name failed (e.g. model_class_names was not populated as expected),
+        # try mapping by assumed index if you know the class ID for 'Text'.
+        # The hantian/yolo-doclaynet example uses `classes=[1]` for Text. This implies ID 1 is Text.
+        # This is risky if the order changes.
+        if 'Text' not in [name for name in model_class_names.values() if name in doclaynet_to_mdr_map]:
+            print(
+                "  Engine: Warning - 'Text' class not found in model_class_names via string mapping. Attempting index-based mapping for critical classes.")
+            # Example: If you know from model card that class ID 9 is 'Text' and 10 is 'Title' for hantian/yolo-doclaynet
+            # This is a COMMON order for DocLayNet, but VERIFY for hantian's model.
+            # From some sources, for DocLayNet, 'Text' is often ID 9, 'Title' is ID 10.
+            # The example `classes=[1]` from the HF page for hantian/yolo-doclaynet is confusing if 'Text' is ID 9.
+            # Let's assume the example `classes=[1]` meant "the class at index 1 in some list", not necessarily ID 1.
+            # We MUST get the correct ID for 'Text'.
+            # For now, let's try to find 'Text' and 'Title' by string in model_class_names and get their IDs.
+
+            text_id = None
+            title_id = None
+            table_id = None
+            figure_id = None
+            formula_id = None
+            caption_id = None  # Generic caption
+
+            for id_val, name_val in model_class_names.items():
+                if name_val == 'Text':
+                    text_id = id_val
+                elif name_val == 'Title':
+                    title_id = id_val
+                elif name_val == 'Table':
+                    table_id = id_val
+                elif name_val == 'Picture':
+                    figure_id = id_val
+                elif name_val == 'Formula':
+                    formula_id = id_val
+                elif name_val == 'Caption':
+                    caption_id = id_val
+                # Add other mappings as needed
+
+            temp_map_by_id = {}
+            if text_id is not None: temp_map_by_id[text_id] = MDRLayoutClass.PLAIN_TEXT
+            if title_id is not None: temp_map_by_id[title_id] = MDRLayoutClass.TITLE
+            if table_id is not None: temp_map_by_id[table_id] = MDRLayoutClass.TABLE
+            if figure_id is not None: temp_map_by_id[figure_id] = MDRLayoutClass.FIGURE
+            if formula_id is not None: temp_map_by_id[formula_id] = MDRLayoutClass.ISOLATE_FORMULA
+            if caption_id is not None: temp_map_by_id[
+                caption_id] = MDRLayoutClass.FIGURE_CAPTION  # Default, refine later
+
+            # Override doclaynet_to_mdr_map if direct ID mapping is more reliable
+            if temp_map_by_id:
+                print(f"  Engine: Using direct ID mapping for some classes: {temp_map_by_id}")
+                # This isn't quite right, the map should be from YOLO ID to MDR Class
+                # The previous doclaynet_to_mdr_map was from string name to MDR Class.
+                # We need a single, consistent map from YOLO's predicted class ID to MDRLayoutClass.
+
+                # Let's rebuild the map: yolo_class_id -> MDRLayoutClass
+                final_yolo_id_to_mdr_class_map = {}
+                if text_id is not None: final_yolo_id_to_mdr_class_map[text_id] = MDRLayoutClass.PLAIN_TEXT
+                if title_id is not None: final_yolo_id_to_mdr_class_map[title_id] = MDRLayoutClass.TITLE
+                # ... map others based on their found IDs ...
+                # For simplicity, let's assume the string-based map from above is preferred if names are available.
+                # The most important thing is to get the ID for 'Text'.
+                # If `model_class_names` is `{0: 'Caption', 1: 'Footnote', ..., 9: 'Text', 10: 'Title'}`
+                # then `doclaynet_to_mdr_map` should correctly map 'Text' to `MDRLayoutClass.PLAIN_TEXT`.
+
+        # Define which MDRLayoutClasses are considered "plain" for fragment merging later (if needed)
+        # This set should use your MDRLayoutClass enum members.
+        plain_mdr_classes: set[MDRLayoutClass] = {
             MDRLayoutClass.TITLE,
             MDRLayoutClass.PLAIN_TEXT,
-            MDRLayoutClass.ABANDON,
-            MDRLayoutClass.FIGURE_CAPTION,
-            MDRLayoutClass.TABLE_CAPTION,
-            MDRLayoutClass.TABLE_FOOTNOTE,
-            MDRLayoutClass.FORMULA_CAPTION,
+            # MDRLayoutClass.ABANDON, # ABANDON layouts usually shouldn't get general text fragments
+            MDRLayoutClass.FIGURE_CAPTION,  # Captions are text
+            MDRLayoutClass.TABLE_CAPTION,  # Captions are text
+            MDRLayoutClass.TABLE_FOOTNOTE,  # Footnotes are text
+            MDRLayoutClass.FORMULA_CAPTION,  # Captions are text
         }
+        print(f"  Engine: Mapping YOLO classes to MDR classes. Effective map used for generation:")
+
+        for cls_id_tensor, xyxy_tensor in zip(results.boxes.cls, results.boxes.xyxy):
+            yolo_cls_id = int(cls_id_tensor.item())  # Get integer class ID from tensor
+
+            # Get the string name for logging/mapping
+            yolo_cls_name = model_class_names.get(yolo_cls_id, f"UnknownID-{yolo_cls_id}")
+
+            # Map YOLO class name to your MDRLayoutClass
+            mdr_cls = None
+            if yolo_cls_name == 'Text':
+                mdr_cls = MDRLayoutClass.PLAIN_TEXT
+            elif yolo_cls_name == 'Title':
+                mdr_cls = MDRLayoutClass.TITLE
+            elif yolo_cls_name == 'Section-header':
+                mdr_cls = MDRLayoutClass.TITLE  # Or a specific header class
+            elif yolo_cls_name == 'List-item':
+                mdr_cls = MDRLayoutClass.PLAIN_TEXT
+            elif yolo_cls_name == 'Table':
+                mdr_cls = MDRLayoutClass.TABLE
+            elif yolo_cls_name == 'Picture':
+                mdr_cls = MDRLayoutClass.FIGURE
+            elif yolo_cls_name == 'Formula':
+                mdr_cls = MDRLayoutClass.ISOLATE_FORMULA
+            elif yolo_cls_name == 'Caption':
+                mdr_cls = MDRLayoutClass.FIGURE_CAPTION  # Default, could be table too
+            elif yolo_cls_name == 'Footnote':
+                mdr_cls = MDRLayoutClass.TABLE_FOOTNOTE  # Or general footnote
+            elif yolo_cls_name in ['Page-header', 'Page-footer']:
+                mdr_cls = MDRLayoutClass.ABANDON
+
+            if mdr_cls is None:
+                # print(f"  Engine: Skipping YOLO box with class '{yolo_cls_name}' (ID {yolo_cls_id}) as it's not mapped to an MDRLayoutClass.")
+                continue
+
+            # print(f"  Engine: Detected YOLO class '{yolo_cls_name}' (ID {yolo_cls_id}), mapped to MDR class '{mdr_cls.name}'")
 
-        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:
+            x1, y1, x2, y2 = map(float, xyxy_tensor)
+            rect = MDRRectangle(lt=(x1, y1), rt=(x2, y1), lb=(x1, y2), rb=(x2, y2))
+            if rect.area < 10:  # Filter tiny boxes
                 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=[])
+            if mdr_cls == MDRLayoutClass.TABLE:
+                yield MDRTableLayoutElement(rect=rect, fragments=[], parsed=None, cls=mdr_cls)  # Explicitly pass cls
+            elif mdr_cls == MDRLayoutClass.ISOLATE_FORMULA:
+                yield MDRFormulaLayoutElement(rect=rect, fragments=[], latex=None, cls=mdr_cls)  # Explicitly pass cls
+            elif mdr_cls == MDRLayoutClass.FIGURE:  # Figure is not in plain_mdr_classes for default fragment assignment
+                yield MDRPlainLayoutElement(cls=mdr_cls, rect=rect, fragments=[])
+            elif mdr_cls in plain_mdr_classes:  # For TITLE, PLAIN_TEXT, CAPTION, etc.
+                yield MDRPlainLayoutElement(cls=mdr_cls, rect=rect, fragments=[])
+            elif mdr_cls == MDRLayoutClass.ABANDON:  # ABANDON class if you want to track but not assign frags by default
+                yield MDRPlainLayoutElement(cls=mdr_cls, rect=rect, fragments=[])
+            # else:
+            # print(f"  Engine: MDR class '{mdr_cls.name}' not explicitly handled for yielding, but was mapped.")
 
     def _match_fragments_to_layouts(self, frags: list[MDROcrFragment], layouts: list[MDRLayoutElement]) -> list[
         MDRLayoutElement]: