Update mdr_pdf_parser.py

mdr_pdf_parser.py

@@ -2656,19 +2656,19 @@ class MDRExtractionEngine:
         self._layout_reader = MDRLayoutReader(model_path=os.path.join(self._model_dir, "layoutreader"))
         print(f"MDR Extraction Engine initialized on device: {self._device}")
 
-    # --- MODIFIED _get_yolo_model METHOD for HF ---
     # In class MDRExtractionEngine:
-
-    def _get_yolo_model(self) -> Any | None:  # Return type can be ultralytics.YOLO
-        """Loads the YOLOv10b-DocLayNet layout detection model."""
+    def _get_yolo_model(self) -> Any | None:  # Return type will be ultralytics.YOLO
+        """Loads the YOLOv10b-DocLayNet layout detection model using ultralytics.YOLO."""
         if self._yolo is None:
+            # Using hantian/yolo-doclaynet (or ppaanngggg if that's the one you have the .pt for)
+            # Ensure these match the model you intend to use
             repo_id = "hantian/yolo-doclaynet"
-            filename = "yolov10b-doclaynet.pt"
+            filename = "yolov10b-doclaynet.pt"  # Or the exact .pt filename from the repo
 
             yolo_cache_dir = Path(self._model_dir) / "yolo_hf_cache_doclaynet"
             mdr_ensure_directory(str(yolo_cache_dir))
 
-            print(f"Attempting to load YOLO model '{filename}' from repo '{repo_id}'...")
+            print(f"Attempting to load YOLO model '{filename}' from repo '{repo_id}' using ultralytics.YOLO...")
             print(f"Hugging Face Hub cache directory for YOLO: {yolo_cache_dir}")
 
             try:
@@ -2677,50 +2677,26 @@ class MDRExtractionEngine:
                     filename=filename,
                     cache_dir=yolo_cache_dir,
                     local_files_only=False,
-                    force_download=False,
+                    force_download=False,  # Set to True if you suspect a corrupted download
                 )
                 print(f"YOLO model file path: {yolo_model_filepath}")
 
-                # --- START: MODIFIED LOADING ---
-                # Attempt to load directly using ultralytics.YOLO
-                try:
-                    from ultralytics import YOLO as UltralyticsYOLO
-                    self._yolo = UltralyticsYOLO(yolo_model_filepath)
-                    print("MDR YOLOv10b-DocLayNet model loaded successfully using ultralytics.YOLO.")
-                except ImportError:
-                    print("ERROR: ultralytics library not found. Cannot load YOLOv10b-DocLayNet.")
-                    print("Please install it: pip install ultralytics")
-                    self._yolo = None
-                    return None  # Critical failure
-                except Exception as e_ultra:  # Catch other ultralytics loading errors
-                    print(f"ERROR: Failed to load YOLO model with ultralytics.YOLO: {e_ultra}")
-                    # If direct ultralytics fails, and your YOLOv10 wrapper exists, you could try it as a fallback,
-                    # but it's likely to fail if ultralytics.YOLO failed due to model structure.
-                    if YOLOv10 is not None:
-                        print("Attempting fallback to doclayout_yolo.YOLOv10 wrapper...")
-                        try:
-                            self._yolo = YOLOv10(yolo_model_filepath)
-                            print("MDR YOLOv10b-DocLayNet model loaded with doclayout_yolo.YOLOv10 wrapper (fallback).")
-                        except Exception as e_wrapper:
-                            print(f"ERROR: Fallback to doclayout_yolo.YOLOv10 also failed: {e_wrapper}")
-                            self._yolo = None
-                            return None
-                    else:
-                        self._yolo = None
-                        return None
-                # --- END: MODIFIED LOADING ---
+                from ultralytics import YOLO as UltralyticsYOLO  # Import here
+                self._yolo = UltralyticsYOLO(yolo_model_filepath)  # This is the line that fails with SCDown
+                print("MDR YOLOv10b-DocLayNet model loaded successfully using ultralytics.YOLO.")
 
-            except HfHubHTTPError as e:
-                print(f"ERROR: Failed to download/access YOLO model '{filename}' via Hugging Face Hub: {e}")
+            except ImportError:
+                print("ERROR: ultralytics library not found. Cannot load YOLOv10b-DocLayNet.")
+                print("Please ensure it's installed: pip install ultralytics (matching version if possible)")
                 self._yolo = None
-            except FileNotFoundError as e:  # This might be redundant if hf_hub_download raises its own error
-                print(f"ERROR: YOLO model file '{filename}' not found after download attempt: {e}")
+            except HfHubHTTPError as e:
+                print(f"ERROR: Failed to download YOLO model '{filename}' via Hugging Face Hub: {e}")
                 self._yolo = None
-            except Exception as e:  # General catch-all for unexpected issues during hf_hub_download or path ops
-                print(f"ERROR: An unexpected issue occurred related to YOLO model file handling for {filename}: {e}")
+            except Exception as e:  # Catch other model loading errors (like the SCDown error)
+                print(f"ERROR: Failed to load YOLO model '{yolo_model_filepath}' with ultralytics.YOLO: {e}")
                 import traceback
                 traceback.print_exc()
-                self._yolo = None
+                self._yolo = None  # Ensure self._yolo is None on failure
 
         return self._yolo
 
@@ -2779,153 +2755,60 @@ class MDRExtractionEngine:
 
         # In class MDRExtractionEngine:
 
-    def _run_yolo_detection(self, img: Image, yolo: Any):  # yolo can be doclayout_yolo.YOLOv10 or ultralytics.YOLO
+    # In class MDRExtractionEngine
+    def _run_yolo_detection(self, img: Image, yolo: Any):  # yolo is an ultralytics.YOLO instance
         img_rgb = img.convert("RGB")
 
-        # 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
+        res_list = yolo.predict(source=img_rgb, imgsz=1024, conf=0.25,  # Adjust conf as needed
                                 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.")
+            print("  Engine: YOLO detection (ultralytics) returned no results or no boxes.")
             return
 
-        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
+        results = res_list[0]
 
-        # 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}")
+            model_class_names = results.names
+            print(f"  Engine: YOLO model class names from ultralytics: {model_class_names}")
         else:
+            # This fallback is a major source of potential error if results.names isn't populated.
+            # It's better to fail or have a very explicit warning if names aren't found.
             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']
+                "  Engine: CRITICAL WARNING - Could not get class names from YOLO model. Layout mapping will likely be incorrect.")
+            # Forcing a known DocLayNet order as a last resort (HIGHLY UNRELIABLE without verification)
+            _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)}
+            print(f"  Engine: Using FALLBACK class names (VERIFY!): {model_class_names}")
 
-        # 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, # 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
+            MDRLayoutClass.TITLE, MDRLayoutClass.PLAIN_TEXT,
+            MDRLayoutClass.FIGURE_CAPTION, MDRLayoutClass.TABLE_CAPTION,
+            MDRLayoutClass.TABLE_FOOTNOTE, MDRLayoutClass.FORMULA_CAPTION,
         }
-        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
+        if results.boxes.cls is None or results.boxes.xyxy is None:
+            print("  Engine: YOLO results.boxes.cls or .xyxy is None.")
+            return
+
+        print(f"  Engine: Processing {len(results.boxes.cls)} detected YOLO boxes...")
+        for i in range(len(results.boxes.cls)):
+            yolo_cls_id = int(results.boxes.cls[i].item())
+            xyxy_tensor = results.boxes.xyxy[i]
 
-            # 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
+            # --- THIS MAPPING IS BASED ON STANDARD DOCLAYNET ---
+            # --- VERIFY IT AGAINST `model_class_names` PRINTED ABOVE ---
             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
+                mdr_cls = MDRLayoutClass.TITLE
             elif yolo_cls_name == 'List-item':
                 mdr_cls = MDRLayoutClass.PLAIN_TEXT
             elif yolo_cls_name == 'Table':
@@ -2935,36 +2818,32 @@ class MDRExtractionEngine:
             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
+                mdr_cls = MDRLayoutClass.FIGURE_CAPTION  # Needs context to be TABLE_CAPTION
             elif yolo_cls_name == 'Footnote':
-                mdr_cls = MDRLayoutClass.TABLE_FOOTNOTE  # Or general footnote
+                mdr_cls = MDRLayoutClass.TABLE_FOOTNOTE  # Needs context
             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.")
+                # print(f"    Skipping YOLO box: class '{yolo_cls_name}' (ID {yolo_cls_id}) - not mapped.")
                 continue
 
-            # print(f"  Engine: Detected YOLO class '{yolo_cls_name}' (ID {yolo_cls_id}), mapped to MDR class '{mdr_cls.name}'")
+            # print(f"    Detected: {yolo_cls_name} (ID {yolo_cls_id}) -> {mdr_cls.name}")
 
             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 rect.area < 10: continue
 
             if mdr_cls == MDRLayoutClass.TABLE:
-                yield MDRTableLayoutElement(rect=rect, fragments=[], parsed=None, cls=mdr_cls)  # Explicitly pass cls
+                yield MDRTableLayoutElement(rect=rect, fragments=[], parsed=None, cls=mdr_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 MDRFormulaLayoutElement(rect=rect, fragments=[], latex=None, cls=mdr_cls)
+            elif mdr_cls == MDRLayoutClass.FIGURE:
                 yield MDRPlainLayoutElement(cls=mdr_cls, rect=rect, fragments=[])
-            elif mdr_cls in plain_mdr_classes:  # For TITLE, PLAIN_TEXT, CAPTION, etc.
+            elif mdr_cls in plain_mdr_classes:
                 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
+            elif mdr_cls == MDRLayoutClass.ABANDON:
                 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]:
         if not frags or not layouts: