natural-pdf 0.1.8__py3-none-any.whl → 0.1.9__py3-none-any.whl

natural_pdf/classification/mixin.py

@@ -3,21 +3,23 @@ from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple, Un
 
 # Assuming PIL is installed as it's needed for vision
 try:
-     from PIL import Image
+    from PIL import Image
 except ImportError:
-     Image = None # type: ignore
+    Image = None  # type: ignore
 
 # Import result classes
-from .results import ClassificationResult # Assuming results.py is in the same dir
+from .results import ClassificationResult  # Assuming results.py is in the same dir
 
 if TYPE_CHECKING:
     # Avoid runtime import cycle
     from natural_pdf.core.page import Page
     from natural_pdf.elements.region import Region
+
     from .manager import ClassificationManager
 
 logger = logging.getLogger(__name__)
 
+
 class ClassificationMixin:
     """
     Mixin class providing classification capabilities to Page and Region objects.
@@ -38,18 +40,18 @@ class ClassificationMixin:
     # Host class needs 'analyses' attribute initialized as Dict[str, Any]
     # analyses: Dict[str, Any]
 
-    # --- End Abstract --- # 
+    # --- End Abstract --- #
 
     def classify(
         self,
         categories: List[str],
-        model: Optional[str] = None, # Default handled by manager
-        using: Optional[str] = None, # Renamed parameter
+        model: Optional[str] = None,  # Default handled by manager
+        using: Optional[str] = None,  # Renamed parameter
         min_confidence: float = 0.0,
-        analysis_key: str = 'classification', # Default key
+        analysis_key: str = "classification",  # Default key
         multi_label: bool = False,
-        **kwargs
-    ) -> "ClassificationMixin": # Return self for chaining
+        **kwargs,
+    ) -> "ClassificationMixin":  # Return self for chaining
         """
         Classifies this item (Page or Region) using the configured manager.
 
@@ -71,22 +73,30 @@ class ClassificationMixin:
             Self for method chaining.
         """
         # Ensure analyses dict exists
-        if not hasattr(self, 'analyses') or self.analyses is None:
-             logger.warning("'analyses' attribute not found or is None. Initializing as empty dict.")
-             self.analyses = {}
+        if not hasattr(self, "analyses") or self.analyses is None:
+            logger.warning("'analyses' attribute not found or is None. Initializing as empty dict.")
+            self.analyses = {}
 
         try:
             manager = self._get_classification_manager()
-            
+
             # Determine the effective model ID and engine type
             effective_model_id = model
-            inferred_using = manager.infer_using(model if model else manager.DEFAULT_TEXT_MODEL, using)
+            inferred_using = manager.infer_using(
+                model if model else manager.DEFAULT_TEXT_MODEL, using
+            )
 
             # If model was not provided, use the manager's default for the inferred engine type
             if effective_model_id is None:
-                effective_model_id = manager.DEFAULT_TEXT_MODEL if inferred_using == 'text' else manager.DEFAULT_VISION_MODEL
-                logger.debug(f"No model provided, using default for mode '{inferred_using}': '{effective_model_id}'")
-            
+                effective_model_id = (
+                    manager.DEFAULT_TEXT_MODEL
+                    if inferred_using == "text"
+                    else manager.DEFAULT_VISION_MODEL
+                )
+                logger.debug(
+                    f"No model provided, using default for mode '{inferred_using}': '{effective_model_id}'"
+                )
+
             # Get content based on the *final* determined engine type
             content = self._get_classification_content(model_type=inferred_using, **kwargs)
 
@@ -94,11 +104,11 @@ class ClassificationMixin:
             result_obj: ClassificationResult = manager.classify_item(
                 item_content=content,
                 categories=categories,
-                model_id=effective_model_id, # Pass the resolved model ID
-                using=inferred_using, # Pass renamed argument
+                model_id=effective_model_id,  # Pass the resolved model ID
+                using=inferred_using,  # Pass renamed argument
                 min_confidence=min_confidence,
                 multi_label=multi_label,
-                **kwargs
+                **kwargs,
             )
 
             # Store the structured result object under the specified key
@@ -106,8 +116,8 @@ class ClassificationMixin:
             logger.debug(f"Stored classification result under key '{analysis_key}': {result_obj}")
 
         except NotImplementedError as nie:
-             logger.error(f"Classification cannot proceed: {nie}")
-             raise
+            logger.error(f"Classification cannot proceed: {nie}")
+            raise
         except Exception as e:
             logger.error(f"Classification failed: {e}", exc_info=True)
             # Optionally re-raise or just log and return self
@@ -118,32 +128,36 @@ class ClassificationMixin:
     @property
     def classification_results(self) -> Optional[ClassificationResult]:
         """Returns the ClassificationResult from the *default* ('classification') key, or None."""
-        if not hasattr(self, 'analyses') or self.analyses is None:
+        if not hasattr(self, "analyses") or self.analyses is None:
             return None
         # Return the result object directly from the default key
-        return self.analyses.get('classification')
+        return self.analyses.get("classification")
 
     @property
     def category(self) -> Optional[str]:
         """Returns the top category label from the *default* ('classification') key, or None."""
-        result_obj = self.classification_results # Uses the property above
+        result_obj = self.classification_results  # Uses the property above
         # Access the property on the result object
         return result_obj.top_category if result_obj else None
 
     @property
     def category_confidence(self) -> Optional[float]:
         """Returns the top category confidence from the *default* ('classification') key, or None."""
-        result_obj = self.classification_results # Uses the property above
+        result_obj = self.classification_results  # Uses the property above
         # Access the property on the result object
         return result_obj.top_confidence if result_obj else None
 
     # Maybe add a helper to get results by specific key?
-    def get_classification_result(self, analysis_key: str = 'classification') -> Optional[ClassificationResult]:
-         """Gets a classification result object stored under a specific key."""
-         if not hasattr(self, 'analyses') or self.analyses is None:
-             return None
-         result = self.analyses.get(analysis_key)
-         if result is not None and not isinstance(result, ClassificationResult):
-              logger.warning(f"Item found under key '{analysis_key}' is not a ClassificationResult (type: {type(result)}). Returning None.")
-              return None
-         return result
+    def get_classification_result(
+        self, analysis_key: str = "classification"
+    ) -> Optional[ClassificationResult]:
+        """Gets a classification result object stored under a specific key."""
+        if not hasattr(self, "analyses") or self.analyses is None:
+            return None
+        result = self.analyses.get(analysis_key)
+        if result is not None and not isinstance(result, ClassificationResult):
+            logger.warning(
+                f"Item found under key '{analysis_key}' is not a ClassificationResult (type: {type(result)}). Returning None."
+            )
+            return None
+        return result

natural_pdf/classification/results.py

@@ -1,62 +1,80 @@
 # natural_pdf/classification/results.py
-from typing import List, Optional, Dict, Any
-from datetime import datetime
 import logging
+from dataclasses import dataclass
+from datetime import datetime
+from typing import Any, Dict, List, Optional
 
 logger = logging.getLogger(__name__)
 
+
+@dataclass
 class CategoryScore:
-    """Represents the score for a single category."""
-    label: str
-    confidence: float # Score between 0.0 and 1.0
-
-    def __init__(self, label: str, confidence: float):
-        # Basic validation
-        if not isinstance(label, str) or not label:
-             logger.warning(f"Initializing CategoryScore with invalid label: {label}")
-             # Fallback or raise? For now, allow but log.
-             # raise ValueError("Category label must be a non-empty string.")
-        if not isinstance(confidence, (float, int)) or not (0.0 <= confidence <= 1.0):
-             logger.warning(f"Initializing CategoryScore with invalid confidence: {confidence} for label '{label}'. Clamping to [0, 1].")
-             confidence = max(0.0, min(1.0, float(confidence)))
-             # raise ValueError("Category confidence must be a float between 0.0 and 1.0.")
-        
-        self.label = str(label)
-        self.confidence = float(confidence)
-
-    def __repr__(self):
-        return f"<CategoryScore label='{self.label}' confidence={self.confidence:.3f}>"
+    """Represents a category and its confidence score from classification."""
+
+    category: str
+    score: float
+
+    def to_dict(self) -> Dict[str, Any]:
+        """Convert to dictionary for serialization."""
+        return {"category": self.category, "score": self.score}
 
+
+@dataclass
 class ClassificationResult:
-    """Holds the structured results of a classification task."""
+    """Results from a classification operation."""
+
+    category: str
+    score: float
+    scores: List[CategoryScore]
     model_id: str
-    using: str # Renamed from engine_type ('text' or 'vision')
     timestamp: datetime
-    parameters: Dict[str, Any] # e.g., {'categories': [...], 'min_confidence': 0.1}
-    scores: List[CategoryScore] # List of scores above threshold, sorted by confidence
-
-    def __init__(self, model_id: str, using: str, timestamp: datetime, parameters: Dict[str, Any], scores: List[CategoryScore]):
-        if not isinstance(scores, list) or not all(isinstance(s, CategoryScore) for s in scores):
-             raise TypeError("Scores must be a list of CategoryScore objects.")
-             
-        self.model_id = str(model_id)
-        self.using = str(using) # Renamed from engine_type
-        self.timestamp = timestamp
-        self.parameters = parameters if parameters is not None else {}
-        # Ensure scores are sorted descending by confidence
-        self.scores = sorted(scores, key=lambda s: s.confidence, reverse=True)
+    using: str  # 'text' or 'vision'
+    parameters: Optional[Dict[str, Any]] = None
+
+    def __init__(
+        self,
+        category: str,
+        score: float,
+        scores: List[CategoryScore],
+        model_id: str,
+        using: str,
+        parameters: Optional[Dict[str, Any]] = None,
+        timestamp: Optional[datetime] = None,
+    ):
+        self.category = category
+        self.score = score
+        self.scores = scores
+        self.model_id = model_id
+        self.using = using
+        self.parameters = parameters or {}
+        self.timestamp = timestamp or datetime.now()
+
+    def to_dict(self) -> Dict[str, Any]:
+        """
+        Convert the classification result to a dictionary for serialization.
+
+        Returns:
+            Dictionary representation of the classification result
+        """
+        return {
+            "category": self.category,
+            "score": self.score,
+            "scores": [s.to_dict() for s in self.scores],
+            "model_id": self.model_id,
+            "using": self.using,
+            "parameters": self.parameters,
+            "timestamp": self.timestamp.isoformat(),
+        }
 
     @property
-    def top_category(self) -> Optional[str]:
-        """Returns the label of the category with the highest confidence."""
-        return self.scores[0].label if self.scores else None
+    def top_category(self) -> str:
+        """Returns the category with the highest score."""
+        return self.category
 
     @property
-    def top_confidence(self) -> Optional[float]:
+    def top_confidence(self) -> float:
         """Returns the confidence score of the top category."""
-        return self.scores[0].confidence if self.scores else None
+        return self.score
 
-    def __repr__(self):
-        top_cat = f" top='{self.top_category}' ({self.top_confidence:.2f})" if self.scores else ""
-        num_scores = len(self.scores)
-        return f"<ClassificationResult model='{self.model_id}' using='{self.using}' scores={num_scores}{top_cat}>" 
+    def __repr__(self) -> str:
+        return f"<ClassificationResult category='{self.category}' score={self.score:.3f} model='{self.model_id}'>"

natural_pdf/collections/mixins.py

@@ -1,10 +1,38 @@
 import logging
-from typing import Callable, Iterable, Any, TypeVar
+from typing import Any, Callable, Iterable, TypeVar
+
 from tqdm.auto import tqdm
 
 logger = logging.getLogger(__name__)
 
-T = TypeVar("T") # Generic type for items in the collection
+T = TypeVar("T")  # Generic type for items in the collection
+
+
+class DirectionalCollectionMixin:
+    """
+    Mixin providing directional methods for collections of elements/regions.
+    """
+
+    def below(self, **kwargs) -> "ElementCollection":
+        """Find regions below all elements in this collection."""
+        return self.apply(lambda element: element.below(**kwargs))
+
+    def above(self, **kwargs) -> "ElementCollection":
+        """Find regions above all elements in this collection."""
+        return self.apply(lambda element: element.above(**kwargs))
+
+    def left(self, **kwargs) -> "ElementCollection":
+        """Find regions to the left of all elements in this collection."""
+        return self.apply(lambda element: element.left(**kwargs))
+
+    def right(self, **kwargs) -> "ElementCollection":
+        """Find regions to the right of all elements in this collection."""
+        return self.apply(lambda element: element.right(**kwargs))
+
+    def expand(self, **kwargs) -> "ElementCollection":
+        """Expand all elements in this collection."""
+        return self.apply(lambda element: element.expand(**kwargs))
+
 
 class ApplyMixin:
     """
@@ -13,6 +41,7 @@ class ApplyMixin:
     Assumes the inheriting class implements `__iter__` and `__len__` appropriately
     for the items to be processed by `apply`.
     """
+
     def _get_items_for_apply(self) -> Iterable[Any]:
         """
         Returns the iterable of items to apply the function to.
@@ -22,7 +51,7 @@ class ApplyMixin:
         # Default to standard iteration over the collection itself
         return iter(self)
 
-    def apply(self: Any, func: Callable[[Any, ...], Any], *args, **kwargs) -> None:
+    def apply(self: Any, func: Callable[[Any, ...], Any], *args, **kwargs) -> Iterable[Any]:
         """
         Applies a function to each item in the collection.
 
@@ -34,7 +63,7 @@ class ApplyMixin:
                       A special keyword argument 'show_progress' (bool, default=False)
                       can be used to display a progress bar.
         """
-        show_progress = kwargs.pop('show_progress', False)
+        show_progress = kwargs.pop("show_progress", False)
         # Derive unit name from class name
         unit_name = self.__class__.__name__.lower()
         items_iterable = self._get_items_for_apply()
@@ -42,22 +71,41 @@ class ApplyMixin:
         # Need total count for tqdm, assumes __len__ is implemented by the inheriting class
         total_items = 0
         try:
-             total_items = len(self)
-        except TypeError: # Handle cases where __len__ might not be defined on self
-             logger.warning(f"Could not determine collection length for progress bar.")
+            total_items = len(self)
+        except TypeError:  # Handle cases where __len__ might not be defined on self
+            logger.warning(f"Could not determine collection length for progress bar.")
 
         if show_progress and total_items > 0:
-            items_iterable = tqdm(items_iterable, total=total_items, desc=f"Applying {func.__name__}", unit=unit_name)
+            items_iterable = tqdm(
+                items_iterable, total=total_items, desc=f"Applying {func.__name__}", unit=unit_name
+            )
         elif show_progress:
-             logger.info(f"Applying {func.__name__} (progress bar disabled for zero/unknown length).")
-
-        for item in items_iterable:
-            try:
-                # Apply the function with the item and any extra args/kwargs
-                func(item, *args, **kwargs)
-            except Exception as e:
-                # Log and continue for batch operations
-                logger.error(f"Error applying {func.__name__} to {item}: {e}", exc_info=True)
-                # Optionally add a mechanism to collect errors
-
-        # Returns None, primarily used for side effects.
+            logger.info(
+                f"Applying {func.__name__} (progress bar disabled for zero/unknown length)."
+            )
+
+        results = [func(item, *args, **kwargs) for item in items_iterable]
+
+        # If results is empty, return an empty list
+        if not results:
+            return []
+
+        # Import here to avoid circular imports
+        from natural_pdf import PDF, Page
+        from natural_pdf.collections.pdf_collection import PDFCollection
+        from natural_pdf.elements.base import Element
+        from natural_pdf.elements.collections import ElementCollection, PageCollection
+        from natural_pdf.elements.region import Region
+
+        first_non_none = next((r for r in results if r is not None), None)
+        first_type = type(first_non_none) if first_non_none is not None else None
+
+        # Return the appropriate collection based on result type (...generally)
+        if issubclass(first_type, Element) or issubclass(first_type, Region):
+            return ElementCollection(results)
+        elif first_type == PDF:
+            return PDFCollection(results)
+        elif first_type == Page:
+            return PageCollection(results)
+
+        return results