natural-pdf 0.1.4__py3-none-any.whl → 0.1.5__py3-none-any.whl

natural_pdf/analyzers/layout/tatr.py

@@ -1,14 +1,15 @@
 # layout_detector_tatr.py
-import logging
 import importlib.util
+import logging
 import os
 import tempfile
-from typing import List, Dict, Any, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
+
 from PIL import Image
 
 # Assuming base class and options are importable
 from .base import LayoutDetector
-from .layout_options import TATRLayoutOptions, BaseLayoutOptions
+from .layout_options import BaseLayoutOptions, TATRLayoutOptions
 
 logger = logging.getLogger(__name__)
 
@@ -26,9 +27,13 @@ if torch_spec and torchvision_spec and transformers_spec:
         from torchvision import transforms
         from transformers import AutoModelForObjectDetection
     except ImportError as e:
-        logger.warning(f"Could not import TATR dependencies (torch, torchvision, transformers): {e}")
+        logger.warning(
+            f"Could not import TATR dependencies (torch, torchvision, transformers): {e}"
+        )
 else:
-    logger.warning("torch, torchvision, or transformers not found. TableTransformerDetector will not be available.")
+    logger.warning(
+        "torch, torchvision, or transformers not found. TableTransformerDetector will not be available."
+    )
 
 
 class TableTransformerDetector(LayoutDetector):
@@ -36,26 +41,36 @@ class TableTransformerDetector(LayoutDetector):
 
     # Custom resize transform (keep as nested class or move outside)
     class MaxResize(object):
-        def __init__(self, max_size=800):
+        def __init__(self, max_size=2000):
             self.max_size = max_size
+
         def __call__(self, image):
             width, height = image.size
             current_max_size = max(width, height)
             scale = self.max_size / current_max_size
             # Use LANCZOS for resizing
-            resized_image = image.resize((int(round(scale*width)), int(round(scale*height))), Image.Resampling.LANCZOS)
+            resized_image = image.resize(
+                (int(round(scale * width)), int(round(scale * height))), Image.Resampling.LANCZOS
+            )
             return resized_image
 
     def __init__(self):
         super().__init__()
         self.supported_classes = {
-            'table', 'table row', 'table column', 'table column header', 'table projected row header', 'table spanning cell' # Add others if supported by models used
+            "table",
+            "table row",
+            "table column",
+            "table column header",
+            "table projected row header",
+            "table spanning cell",  # Add others if supported by models used
         }
         # Models are loaded via _get_model
 
     def is_available(self) -> bool:
         """Check if dependencies are installed."""
-        return torch is not None and transforms is not None and AutoModelForObjectDetection is not None
+        return (
+            torch is not None and transforms is not None and AutoModelForObjectDetection is not None
+        )
 
     def _get_cache_key(self, options: TATRLayoutOptions) -> str:
         """Generate cache key based on model IDs and device."""
@@ -63,26 +78,30 @@ class TableTransformerDetector(LayoutDetector):
             options = TATRLayoutOptions(device=options.device)
 
         device_key = str(options.device).lower()
-        det_model_key = options.detection_model.replace('/','_')
-        struct_model_key = options.structure_model.replace('/','_')
+        det_model_key = options.detection_model.replace("/", "_")
+        struct_model_key = options.structure_model.replace("/", "_")
         return f"{self.__class__.__name__}_{device_key}_{det_model_key}_{struct_model_key}"
 
     def _load_model_from_options(self, options: TATRLayoutOptions) -> Dict[str, Any]:
         """Load the TATR detection and structure models."""
         if not self.is_available():
-             raise RuntimeError("TATR dependencies (torch, torchvision, transformers) not installed.")
+            raise RuntimeError(
+                "TATR dependencies (torch, torchvision, transformers) not installed."
+            )
 
         device = options.device or ("cuda" if torch.cuda.is_available() else "cpu")
-        self.logger.info(f"Loading TATR models: Detection='{options.detection_model}', Structure='{options.structure_model}' onto device='{device}'")
+        self.logger.info(
+            f"Loading TATR models: Detection='{options.detection_model}', Structure='{options.structure_model}' onto device='{device}'"
+        )
         try:
             detection_model = AutoModelForObjectDetection.from_pretrained(
-                options.detection_model, revision="no_timm" # Important revision for some versions
+                options.detection_model, revision="no_timm"  # Important revision for some versions
             ).to(device)
             structure_model = AutoModelForObjectDetection.from_pretrained(
                 options.structure_model
             ).to(device)
             self.logger.info("TATR models loaded.")
-            return {'detection': detection_model, 'structure': structure_model}
+            return {"detection": detection_model, "structure": structure_model}
         except Exception as e:
             self.logger.error(f"Failed to load TATR models: {e}", exc_info=True)
             raise
@@ -97,19 +116,21 @@ class TableTransformerDetector(LayoutDetector):
     def rescale_bboxes(self, out_bbox, size):
         img_w, img_h = size
         boxes = self.box_cxcywh_to_xyxy(out_bbox)
-        boxes = boxes * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32).to(out_bbox.device) # Ensure tensor on correct device
+        boxes = boxes * torch.tensor([img_w, img_h, img_w, img_h], dtype=torch.float32).to(
+            out_bbox.device
+        )  # Ensure tensor on correct device
         return boxes
 
     def outputs_to_objects(self, outputs, img_size, id2label):
         logits = outputs.logits
         bboxes = outputs.pred_boxes
         # Use softmax activation function
-        prob = logits.softmax(-1)[0, :, :-1] # Exclude the "no object" class
+        prob = logits.softmax(-1)[0, :, :-1]  # Exclude the "no object" class
         scores, labels = prob.max(-1)
 
         # Convert to absolute coordinates
         img_w, img_h = img_size
-        boxes = self.rescale_bboxes(bboxes[0, ...], (img_w, img_h)) # Pass tuple size
+        boxes = self.rescale_bboxes(bboxes[0, ...], (img_w, img_h))  # Pass tuple size
 
         # Move results to CPU for list comprehension
         scores = scores.cpu().tolist()
@@ -118,49 +139,62 @@ class TableTransformerDetector(LayoutDetector):
 
         objects = []
         for score, label_idx, bbox in zip(scores, labels, boxes):
-            class_label = id2label.get(label_idx, 'unknown') # Use get with default
-            if class_label != 'no object' and class_label != 'unknown':
-                objects.append({
-                    'label': class_label,
-                    'score': float(score),
-                    'bbox': [round(float(c), 2) for c in bbox] # Round coordinates
-                })
+            class_label = id2label.get(label_idx, "unknown")  # Use get with default
+            if class_label != "no object" and class_label != "unknown":
+                objects.append(
+                    {
+                        "label": class_label,
+                        "score": float(score),
+                        "bbox": [round(float(c), 2) for c in bbox],  # Round coordinates
+                    }
+                )
         return objects
+
     # --- End Helper Methods ---
 
     def detect(self, image: Image.Image, options: BaseLayoutOptions) -> List[Dict[str, Any]]:
         """Detect tables and their structure in an image."""
         if not self.is_available():
-            raise RuntimeError("TATR dependencies (torch, torchvision, transformers) not installed.")
+            raise RuntimeError(
+                "TATR dependencies (torch, torchvision, transformers) not installed."
+            )
 
         if not isinstance(options, TATRLayoutOptions):
-             self.logger.warning("Received BaseLayoutOptions, expected TATRLayoutOptions. Using defaults.")
-             options = TATRLayoutOptions(
-                 confidence=options.confidence, classes=options.classes,
-                 exclude_classes=options.exclude_classes, device=options.device,
-                 extra_args=options.extra_args
-             )
+            self.logger.warning(
+                "Received BaseLayoutOptions, expected TATRLayoutOptions. Using defaults."
+            )
+            options = TATRLayoutOptions(
+                confidence=options.confidence,
+                classes=options.classes,
+                exclude_classes=options.exclude_classes,
+                device=options.device,
+                extra_args=options.extra_args,
+            )
 
         self.validate_classes(options.classes or [])
         if options.exclude_classes:
             self.validate_classes(options.exclude_classes)
 
         models = self._get_model(options)
-        detection_model = models['detection']
-        structure_model = models['structure']
+        detection_model = models["detection"]
+        structure_model = models["structure"]
         device = options.device or ("cuda" if torch.cuda.is_available() else "cpu")
 
         # Prepare transforms based on options
-        detection_transform = transforms.Compose([
-            self.MaxResize(options.max_detection_size),
-            transforms.ToTensor(),
-            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
-        ])
-        structure_transform = transforms.Compose([
-            self.MaxResize(options.max_structure_size),
-            transforms.ToTensor(),
-            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
-        ])
+        detection_transform = transforms.Compose(
+            [
+                self.MaxResize(options.max_detection_size),
+                transforms.ToTensor(),
+                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+            ]
+        )
+        structure_transform = transforms.Compose(
+            [
+                self.MaxResize(options.max_structure_size),
+                transforms.ToTensor(),
+                transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
+            ]
+        )
 
         # --- Detect Tables ---
         self.logger.debug("Running TATR table detection...")
@@ -169,38 +203,60 @@ class TableTransformerDetector(LayoutDetector):
             outputs = detection_model(pixel_values)
 
         id2label_det = detection_model.config.id2label
-        id2label_det[detection_model.config.num_labels] = "no object" # Add no object class
+        id2label_det[detection_model.config.num_labels] = "no object"  # Add no object class
         tables = self.outputs_to_objects(outputs, image.size, id2label_det)
-        tables = [t for t in tables if t['score'] >= options.confidence and t['label'] == 'table'] # Filter for tables
+        tables = [
+            t for t in tables if t["score"] >= options.confidence and t["label"] == "table"
+        ]  # Filter for tables
         self.logger.debug(f"Detected {len(tables)} table regions.")
 
         all_detections = []
 
         # Add table detections if requested
-        normalized_classes_req = {self._normalize_class_name(c) for c in options.classes} if options.classes else None
-        normalized_classes_excl = {self._normalize_class_name(c) for c in options.exclude_classes} if options.exclude_classes else set()
-
-        if normalized_classes_req is None or 'table' in normalized_classes_req:
-             if 'table' not in normalized_classes_excl:
-                 for table in tables:
-                     all_detections.append({
-                         'bbox': tuple(table['bbox']),
-                         'class': 'table',
-                         'confidence': float(table['score']),
-                         'normalized_class': 'table',
-                         'source': 'layout',
-                         'model': 'tatr'
-                     })
+        normalized_classes_req = (
+            {self._normalize_class_name(c) for c in options.classes} if options.classes else None
+        )
+        normalized_classes_excl = (
+            {self._normalize_class_name(c) for c in options.exclude_classes}
+            if options.exclude_classes
+            else set()
+        )
+
+        if normalized_classes_req is None or "table" in normalized_classes_req:
+            if "table" not in normalized_classes_excl:
+                for table in tables:
+                    all_detections.append(
+                        {
+                            "bbox": tuple(table["bbox"]),
+                            "class": "table",
+                            "confidence": float(table["score"]),
+                            "normalized_class": "table",
+                            "source": "layout",
+                            "model": "tatr",
+                        }
+                    )
 
         # --- Process Structure ---
-        structure_class_names = {'table row', 'table column', 'table column header', 'table projected row header', 'table spanning cell'}
-        normalized_structure_classes = {self._normalize_class_name(c) for c in structure_class_names}
+        structure_class_names = {
+            "table row",
+            "table column",
+            "table column header",
+            "table projected row header",
+            "table spanning cell",
+        }
+        normalized_structure_classes = {
+            self._normalize_class_name(c) for c in structure_class_names
+        }
 
         needed_structure = False
-        if normalized_classes_req is None: # If no specific classes requested
-             needed_structure = any(norm_cls not in normalized_classes_excl for norm_cls in normalized_structure_classes)
-        else: # Specific classes requested
-             needed_structure = any(norm_cls in normalized_classes_req for norm_cls in normalized_structure_classes)
+        if normalized_classes_req is None:  # If no specific classes requested
+            needed_structure = any(
+                norm_cls not in normalized_classes_excl for norm_cls in normalized_structure_classes
+            )
+        else:  # Specific classes requested
+            needed_structure = any(
+                norm_cls in normalized_classes_req for norm_cls in normalized_structure_classes
+            )
 
         if needed_structure and tables:
             self.logger.debug("Running TATR structure recognition...")
@@ -208,44 +264,53 @@ class TableTransformerDetector(LayoutDetector):
             id2label_struct[structure_model.config.num_labels] = "no object"
 
             for table in tables:
-                x_min, y_min, x_max, y_max = map(int, table['bbox'])
+                x_min, y_min, x_max, y_max = map(int, table["bbox"])
                 # Ensure coordinates are within image bounds
                 x_min, y_min = max(0, x_min), max(0, y_min)
                 x_max, y_max = min(image.width, x_max), min(image.height, y_max)
-                if x_max <= x_min or y_max <= y_min: continue # Skip invalid crop
+                if x_max <= x_min or y_max <= y_min:
+                    continue  # Skip invalid crop
 
                 cropped_table = image.crop((x_min, y_min, x_max, y_max))
-                if cropped_table.width == 0 or cropped_table.height == 0: continue # Skip empty crop
+                if cropped_table.width == 0 or cropped_table.height == 0:
+                    continue  # Skip empty crop
 
                 pixel_values_struct = structure_transform(cropped_table).unsqueeze(0).to(device)
                 with torch.no_grad():
                     outputs_struct = structure_model(pixel_values_struct)
 
-                structure_elements = self.outputs_to_objects(outputs_struct, cropped_table.size, id2label_struct)
-                structure_elements = [e for e in structure_elements if e['score'] >= options.confidence]
+                structure_elements = self.outputs_to_objects(
+                    outputs_struct, cropped_table.size, id2label_struct
+                )
+                structure_elements = [
+                    e for e in structure_elements if e["score"] >= options.confidence
+                ]
 
                 for element in structure_elements:
-                    element_class_orig = element['label']
+                    element_class_orig = element["label"]
                     normalized_class = self._normalize_class_name(element_class_orig)
 
                     # Apply class filtering
-                    if normalized_classes_req and normalized_class not in normalized_classes_req: continue
-                    if normalized_class in normalized_classes_excl: continue
+                    if normalized_classes_req and normalized_class not in normalized_classes_req:
+                        continue
+                    if normalized_class in normalized_classes_excl:
+                        continue
 
                     # Adjust coordinates
-                    ex0, ey0, ex1, ey1 = element['bbox']
+                    ex0, ey0, ex1, ey1 = element["bbox"]
                     adj_bbox = (ex0 + x_min, ey0 + y_min, ex1 + x_min, ey1 + y_min)
 
-                    all_detections.append({
-                        'bbox': adj_bbox,
-                        'class': element_class_orig,
-                        'confidence': float(element['score']),
-                        'normalized_class': normalized_class,
-                        'source': 'layout',
-                        'model': 'tatr'
-                    })
+                    all_detections.append(
+                        {
+                            "bbox": adj_bbox,
+                            "class": element_class_orig,
+                            "confidence": float(element["score"]),
+                            "normalized_class": normalized_class,
+                            "source": "layout",
+                            "model": "tatr",
+                        }
+                    )
             self.logger.debug(f"Added {len(all_detections) - len(tables)} structure elements.")
 
         self.logger.info(f"TATR detected {len(all_detections)} layout elements matching criteria.")
         return all_detections
-

natural_pdf/analyzers/layout/yolo.py

@@ -1,24 +1,38 @@
 # layout_detector_yolo.py
-import logging
 import importlib.util
+import logging
 import os
 import tempfile
-from typing import List, Dict, Any, Optional
+from typing import Any, Dict, List, Optional
+
 from PIL import Image
 
 # Assuming base class and options are importable
 try:
     from .base import LayoutDetector
-    from .layout_options import YOLOLayoutOptions, BaseLayoutOptions
+    from .layout_options import BaseLayoutOptions, YOLOLayoutOptions
 except ImportError:
     # Placeholders if run standalone or imports fail
-    class BaseLayoutOptions: pass
-    class YOLOLayoutOptions(BaseLayoutOptions): pass
+    class BaseLayoutOptions:
+        pass
+
+    class YOLOLayoutOptions(BaseLayoutOptions):
+        pass
+
     class LayoutDetector:
-         def __init__(self): self.logger=logging.getLogger(); self.supported_classes=set()
-         def _get_model(self, options): raise NotImplementedError
-         def _normalize_class_name(self, n): return n
-         def validate_classes(self, c): pass
+        def __init__(self):
+            self.logger = logging.getLogger()
+            self.supported_classes = set()
+
+        def _get_model(self, options):
+            raise NotImplementedError
+
+        def _normalize_class_name(self, n):
+            return n
+
+        def validate_classes(self, c):
+            pass
+
     logging.basicConfig()
 
 logger = logging.getLogger(__name__)
@@ -36,7 +50,9 @@ if yolo_spec and hf_spec:
     except ImportError as e:
         logger.warning(f"Could not import YOLO dependencies: {e}")
 else:
-    logger.warning("doclayout_yolo or huggingface_hub not found. YOLODocLayoutDetector will not be available.")
+    logger.warning(
+        "doclayout_yolo or huggingface_hub not found. YOLODocLayoutDetector will not be available."
+    )
 
 
 class YOLODocLayoutDetector(LayoutDetector):
@@ -45,9 +61,16 @@ class YOLODocLayoutDetector(LayoutDetector):
     def __init__(self):
         super().__init__()
         self.supported_classes = {
-            'title', 'plain text', 'abandon', 'figure', 'figure_caption',
-            'table', 'table_caption', 'table_footnote', 'isolate_formula',
-            'formula_caption'
+            "title",
+            "plain text",
+            "abandon",
+            "figure",
+            "figure_caption",
+            "table",
+            "table_caption",
+            "table_footnote",
+            "isolate_formula",
+            "formula_caption",
         }
 
     def is_available(self) -> bool:
@@ -58,8 +81,8 @@ class YOLODocLayoutDetector(LayoutDetector):
         """Generate cache key based on model repo/file and device."""
         # Ensure options is the correct type
         if not isinstance(options, YOLOLayoutOptions):
-             # This shouldn't happen if called correctly, but handle defensively
-             options = YOLOLayoutOptions(device=options.device) # Use base device
+            # This shouldn't happen if called correctly, but handle defensively
+            options = YOLOLayoutOptions(device=options.device)  # Use base device
 
         device_key = str(options.device).lower()
         model_key = f"{options.model_repo.replace('/','_')}_{options.model_file}"
@@ -68,7 +91,7 @@ class YOLODocLayoutDetector(LayoutDetector):
     def _load_model_from_options(self, options: YOLOLayoutOptions) -> Any:
         """Load the YOLOv10 model based on options."""
         if not self.is_available():
-             raise RuntimeError("YOLO dependencies (doclayout_yolo, huggingface_hub) not installed.")
+            raise RuntimeError("YOLO dependencies (doclayout_yolo, huggingface_hub) not installed.")
         self.logger.info(f"Loading YOLO model: {options.model_repo}/{options.model_file}")
         try:
             model_path = hf_hub_download(repo_id=options.model_repo, filename=options.model_file)
@@ -86,12 +109,16 @@ class YOLODocLayoutDetector(LayoutDetector):
 
         # Ensure options are the correct type, falling back to defaults if base type passed
         if not isinstance(options, YOLOLayoutOptions):
-             self.logger.warning("Received BaseLayoutOptions, expected YOLOLayoutOptions. Using defaults.")
-             options = YOLOLayoutOptions(
-                 confidence=options.confidence, classes=options.classes,
-                 exclude_classes=options.exclude_classes, device=options.device,
-                 extra_args=options.extra_args
-             )
+            self.logger.warning(
+                "Received BaseLayoutOptions, expected YOLOLayoutOptions. Using defaults."
+            )
+            options = YOLOLayoutOptions(
+                confidence=options.confidence,
+                classes=options.classes,
+                exclude_classes=options.exclude_classes,
+                device=options.device,
+                extra_args=options.extra_args,
+            )
 
         # Validate classes before proceeding
         self.validate_classes(options.classes or [])
@@ -108,58 +135,71 @@ class YOLODocLayoutDetector(LayoutDetector):
             temp_image_path = os.path.join(temp_dir, "temp_layout_image.png")
             try:
                 self.logger.debug(f"Saving temporary image for YOLO detector to: {temp_image_path}")
-                image.convert("RGB").save(temp_image_path) # Ensure RGB
+                image.convert("RGB").save(temp_image_path)  # Ensure RGB
 
                 # Run model prediction
-                self.logger.debug(f"Running YOLO prediction (imgsz={options.image_size}, conf={options.confidence}, device={options.device})...")
+                self.logger.debug(
+                    f"Running YOLO prediction (imgsz={options.image_size}, conf={options.confidence}, device={options.device})..."
+                )
                 results = model.predict(
                     temp_image_path,
                     imgsz=options.image_size,
                     conf=options.confidence,
-                    device=options.device or 'cpu' # Default to cpu if None
+                    device=options.device or "cpu",  # Default to cpu if None
                     # Add other predict args from options.extra_args if needed
                     # **options.extra_args
                 )
                 self.logger.debug(f"YOLO prediction returned {len(results)} result objects.")
 
                 # Process results into standardized format
-                img_width, img_height = image.size # Get original image size for context if needed
+                img_width, img_height = image.size  # Get original image size for context if needed
                 for result in results:
-                    if result.boxes is None: continue
+                    if result.boxes is None:
+                        continue
                     boxes = result.boxes.xyxy
                     labels = result.boxes.cls
                     scores = result.boxes.conf
-                    class_names = result.names # Dictionary mapping index to name
+                    class_names = result.names  # Dictionary mapping index to name
 
                     for box, label_idx_tensor, score_tensor in zip(boxes, labels, scores):
                         x_min, y_min, x_max, y_max = map(float, box.tolist())
-                        label_idx = int(label_idx_tensor.item()) # Get int index
-                        score = float(score_tensor.item()) # Get float score
+                        label_idx = int(label_idx_tensor.item())  # Get int index
+                        score = float(score_tensor.item())  # Get float score
 
                         if label_idx not in class_names:
-                             self.logger.warning(f"Label index {label_idx} not found in model names dict. Skipping.")
-                             continue
+                            self.logger.warning(
+                                f"Label index {label_idx} not found in model names dict. Skipping."
+                            )
+                            continue
                         label_name = class_names[label_idx]
                         normalized_class = self._normalize_class_name(label_name)
 
                         # Apply class filtering (using normalized names)
-                        if options.classes and normalized_class not in [self._normalize_class_name(c) for c in options.classes]:
+                        if options.classes and normalized_class not in [
+                            self._normalize_class_name(c) for c in options.classes
+                        ]:
                             continue
-                        if options.exclude_classes and normalized_class in [self._normalize_class_name(c) for c in options.exclude_classes]:
+                        if options.exclude_classes and normalized_class in [
+                            self._normalize_class_name(c) for c in options.exclude_classes
+                        ]:
                             continue
 
-                        detections.append({
-                            'bbox': (x_min, y_min, x_max, y_max),
-                            'class': label_name,
-                            'confidence': score,
-                            'normalized_class': normalized_class,
-                            'source': 'layout',
-                            'model': 'yolo'
-                        })
-                self.logger.info(f"YOLO detected {len(detections)} layout elements matching criteria.")
+                        detections.append(
+                            {
+                                "bbox": (x_min, y_min, x_max, y_max),
+                                "class": label_name,
+                                "confidence": score,
+                                "normalized_class": normalized_class,
+                                "source": "layout",
+                                "model": "yolo",
+                            }
+                        )
+                self.logger.info(
+                    f"YOLO detected {len(detections)} layout elements matching criteria."
+                )
 
             except Exception as e:
                 self.logger.error(f"Error during YOLO detection: {e}", exc_info=True)
-                raise # Re-raise the exception
+                raise  # Re-raise the exception
 
         return detections