vision_extractor.py

"""
Visual ROI (Region of Interest) Extractor for TextStruct.

This module extracts figures, tables, and diagrams from page images based on
Phase 4.x metadata (layout detection, table detection, figure detection).

Extracted ROIs are saved as separate image files for vision embedding.
"""

from pathlib import Path
from typing import List, Dict, Tuple, Optional
import numpy as np
from PIL import Image

from logger import get_logger

logger = get_logger()


class VisualROIExtractor:
    """
    Extract and crop visual content regions from page images.

    Features:
    - Extracts figures, tables, diagrams based on metadata
    - Crops with configurable padding
    - Saves as separate image files
    - Links to parent chunks for hybrid retrieval
    """

    def __init__(
        self,
        padding_px: int = 10,
        min_width: int = 50,
        min_height: int = 50,
        max_width: int = 4000,
        max_height: int = 4000,
    ):
        """
        Initialize ROI extractor.

        Args:
            padding_px: Padding around ROI in pixels (default: 10)
            min_width: Minimum ROI width in pixels (default: 50)
            min_height: Minimum ROI height in pixels (default: 50)
            max_width: Maximum ROI width in pixels (default: 4000)
            max_height: Maximum ROI height in pixels (default: 4000)
        """
        self.padding_px = padding_px
        self.min_width = min_width
        self.min_height = min_height
        self.max_width = max_width
        self.max_height = max_height

    def extract_rois(
        self,
        page_image_path: Path,
        blocks: List[Dict],
        output_dir: Path,
        page_idx: int = 0,
    ) -> List[Dict]:
        """
        Extract ROIs from a page image based on block metadata.

        Args:
            page_image_path: Path to page image (PNG/JPEG)
            blocks: List of OCR blocks with metadata
            output_dir: Directory to save extracted ROI images
            page_idx: Page index for naming

        Returns:
            List of ROI metadata dictionaries:
                {
                    "roi_id": str,
                    "image_path": Path,
                    "bbox": (x0, y0, x1, y1),  # pixel coordinates
                    "visual_type": str,  # "figure", "table", "diagram"
                    "page": int,
                    "parent_block_id": str,
                    "confidence": float,
                }
        """
        page_image_path = Path(page_image_path)
        output_dir = Path(output_dir)
        output_dir.mkdir(parents=True, exist_ok=True)

        if not page_image_path.exists():
            logger.warning(f"Page image not found: {page_image_path}")
            return []

        # Load page image
        try:
            img = Image.open(page_image_path)
            img_width, img_height = img.size
        except Exception as e:
            logger.error(f"Failed to load page image {page_image_path}: {e}")
            return []

        rois = []
        roi_count = 0

        for block_idx, block in enumerate(blocks):
            metadata = block.get("metadata", {})

            # Check if block has visual content
            visual_type = self._determine_visual_type(metadata)

            if visual_type is None:
                continue

            # Get bounding box
            bbox = block.get("bbox", None)

            if bbox is None:
                logger.debug(f"Block {block_idx} on page {page_idx} missing bbox")
                continue

            # Denormalize bbox if needed
            if self._is_normalized_bbox(bbox):
                bbox_px = self._denormalize_bbox(bbox, img_width, img_height)
            else:
                bbox_px = bbox

            # Apply padding
            bbox_padded = self._apply_padding(bbox_px, img_width, img_height)

            # Validate bbox dimensions
            if not self._is_valid_bbox(bbox_padded):
                logger.debug(f"Invalid bbox dimensions for block {block_idx}")
                continue

            # Crop ROI
            try:
                roi_img = img.crop(bbox_padded)
            except Exception as e:
                logger.warning(f"Failed to crop ROI from page {page_idx}, block {block_idx}: {e}")
                continue

            # Generate ROI ID and save
            roi_id = f"roi_page{page_idx:04d}_block{block_idx:04d}_{visual_type}"
            roi_filename = f"{roi_id}.png"
            roi_path = output_dir / roi_filename

            try:
                roi_img.save(roi_path, "PNG")
            except Exception as e:
                logger.warning(f"Failed to save ROI {roi_id}: {e}")
                continue

            # Create ROI metadata
            roi_meta = {
                "roi_id": roi_id,
                "image_path": str(roi_path),
                "bbox": bbox_padded,  # Pixel coordinates
                "visual_type": visual_type,
                "page": page_idx,
                "parent_block_id": block.get("id", f"block_{block_idx}"),
                "confidence": metadata.get("figure_confidence", metadata.get("table_confidence", 0.0)),
            }

            # Add type-specific metadata
            if visual_type == "table":
                roi_meta["table_metadata"] = metadata.get("table_metadata", {})
            elif visual_type == "figure":
                roi_meta["figure_metadata"] = metadata.get("figure_metadata", {})

            rois.append(roi_meta)
            roi_count += 1

        if roi_count > 0:
            logger.info(f"Extracted {roi_count} ROIs from page {page_idx}")

        return rois

    def _determine_visual_type(self, metadata: Dict) -> Optional[str]:
        """
        Determine if block contains visual content and classify type.

        Args:
            metadata: Block metadata dictionary

        Returns:
            "figure", "table", "diagram", or None
        """
        # Check for table
        if metadata.get("table_metadata") is not None:
            return "table"

        # Check for figure
        if metadata.get("figure_metadata") is not None:
            fig_meta = metadata["figure_metadata"]
            fig_type = fig_meta.get("type", "unknown")

            if fig_type in ["diagram", "chart", "graph", "plot"]:
                return "diagram"
            else:
                return "figure"

        # Check for has_visual_content flag
        if metadata.get("has_visual_content", False):
            return "figure"

        return None

    def _is_normalized_bbox(self, bbox: Tuple) -> bool:
        """
        Check if bbox coordinates are normalized (0-1) or pixel coordinates.

        Args:
            bbox: (x0, y0, x1, y1) tuple

        Returns:
            True if normalized, False if pixel coordinates
        """
        x0, y0, x1, y1 = bbox

        # If all values are between 0 and 1, assume normalized
        return all(0 <= val <= 1 for val in [x0, y0, x1, y1])

    def _denormalize_bbox(
        self,
        bbox: Tuple[float, float, float, float],
        img_width: int,
        img_height: int
    ) -> Tuple[int, int, int, int]:
        """
        Convert normalized bbox (0-1) to pixel coordinates.

        Args:
            bbox: Normalized (x0, y0, x1, y1) tuple
            img_width: Image width in pixels
            img_height: Image height in pixels

        Returns:
            Pixel coordinates (x0, y0, x1, y1)
        """
        x0, y0, x1, y1 = bbox

        x0_px = int(x0 * img_width)
        y0_px = int(y0 * img_height)
        x1_px = int(x1 * img_width)
        y1_px = int(y1 * img_height)

        return (x0_px, y0_px, x1_px, y1_px)

    def _apply_padding(
        self,
        bbox: Tuple[int, int, int, int],
        img_width: int,
        img_height: int
    ) -> Tuple[int, int, int, int]:
        """
        Apply padding to bbox while staying within image bounds.

        Args:
            bbox: Pixel coordinates (x0, y0, x1, y1)
            img_width: Image width
            img_height: Image height

        Returns:
            Padded bbox coordinates
        """
        x0, y0, x1, y1 = bbox

        # Apply padding
        x0_padded = max(0, x0 - self.padding_px)
        y0_padded = max(0, y0 - self.padding_px)
        x1_padded = min(img_width, x1 + self.padding_px)
        y1_padded = min(img_height, y1 + self.padding_px)

        return (x0_padded, y0_padded, x1_padded, y1_padded)

    def _is_valid_bbox(self, bbox: Tuple[int, int, int, int]) -> bool:
        """
        Validate bbox dimensions.

        Args:
            bbox: Pixel coordinates (x0, y0, x1, y1)

        Returns:
            True if valid, False otherwise
        """
        x0, y0, x1, y1 = bbox

        width = x1 - x0
        height = y1 - y0

        # Check minimum dimensions
        if width < self.min_width or height < self.min_height:
            return False

        # Check maximum dimensions
        if width > self.max_width or height > self.max_height:
            return False

        # Check that x1 > x0 and y1 > y0
        if x1 <= x0 or y1 <= y0:
            return False

        return True

    def batch_extract_rois(
        self,
        page_images: List[Path],
        pages_blocks: List[List[Dict]],
        output_dir: Path,
    ) -> List[Dict]:
        """
        Extract ROIs from multiple pages.

        Args:
            page_images: List of paths to page images
            pages_blocks: List of block lists (one per page)
            output_dir: Directory to save extracted ROIs

        Returns:
            List of all ROI metadata dictionaries
        """
        all_rois = []

        if len(page_images) != len(pages_blocks):
            logger.warning(
                f"Mismatch: {len(page_images)} images but {len(pages_blocks)} block lists"
            )
            # Use minimum length to avoid index errors
            num_pages = min(len(page_images), len(pages_blocks))
        else:
            num_pages = len(page_images)

        logger.info(f"Extracting ROIs from {num_pages} pages...")

        for page_idx in range(num_pages):
            page_image = page_images[page_idx]
            blocks = pages_blocks[page_idx]

            rois = self.extract_rois(
                page_image_path=page_image,
                blocks=blocks,
                output_dir=output_dir,
                page_idx=page_idx,
            )

            all_rois.extend(rois)

        logger.info(f"Extracted {len(all_rois)} total ROIs from {num_pages} pages")

        return all_rois


def extract_rois_from_pipeline(
    image_paths: List[Path],
    pages_blocks: List[List[Dict]],
    output_dir: Path,
    padding_px: int = 10,
) -> List[Dict]:
    """
    Convenience function for pipeline integration.

    Args:
        image_paths: List of page image paths
        pages_blocks: List of block lists from OCR
        output_dir: Output directory for ROIs
        padding_px: Padding around ROIs (default: 10)

    Returns:
        List of ROI metadata dictionaries
    """
    extractor = VisualROIExtractor(padding_px=padding_px)

    rois = extractor.batch_extract_rois(
        page_images=image_paths,
        pages_blocks=pages_blocks,
        output_dir=output_dir,
    )

    return rois