src/pipeline/input_loader.py

"""Input loader stage for the character attribute extraction pipeline."""

import os
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple, Union, Iterator
from PIL import Image
import logging
import torch
from torch.utils.data import Dataset, DataLoader
try:
    from datasets import Dataset as HFDataset, load_dataset
    HF_AVAILABLE = True
except ImportError:
    HF_AVAILABLE = False
    HFDataset = None

from .base import PipelineStage

logger = logging.getLogger(__name__)

class DatasetItem:
    """Represents a single item from the dataset."""
    
    def __init__(self, image_path: str, text_path: Optional[str] = None, tags: Optional[str] = None):
        self.image_path = image_path
        self.text_path = text_path
        self.tags = tags
        self.item_id = Path(image_path).stem
    
    def load_image(self) -> Image.Image:
        """Load and return the PIL Image."""
        try:
            return Image.open(self.image_path).convert('RGB')
        except Exception as e:
            logger.error(f"Failed to load image {self.image_path}: {e}")
            raise
    
    def load_tags(self) -> str:
        """Load tags from text file or return provided tags."""
        if self.tags:
            return self.tags
        
        if self.text_path and os.path.exists(self.text_path):
            try:
                with open(self.text_path, 'r', encoding='utf-8') as f:
                    return f.read().strip()
            except Exception as e:
                logger.error(f"Failed to load tags from {self.text_path}: {e}")
                return ""
        
        return ""

class CharacterDataset(Dataset):
    """PyTorch Dataset for character images and attributes."""
    
    def __init__(self, items: List[DatasetItem], transform=None):
        self.items = items
        self.transform = transform
    
    def __len__(self):
        return len(self.items)
    
    def __getitem__(self, idx):
        item = self.items[idx]
        image = item.load_image()
        tags = item.load_tags()
        
        if self.transform:
            image = self.transform(image)
        
        return {
            'image': image,
            'tags': tags,
            'item_id': item.item_id,
            'image_path': item.image_path
        }

class InputLoader(PipelineStage):
    """Loads images and associated text data from the dataset."""
    
    def __init__(self, config: Optional[Dict[str, Any]] = None):
        super().__init__("InputLoader", config)
        if config:
            self.dataset_path = config.get('dataset_path', './continued/sensitive')
            self.batch_size = config.get('batch_size', 32)
            self.num_workers = config.get('num_workers', 4)
        else:
            self.dataset_path = './continued/sensitive'
            self.batch_size = 32
            self.num_workers = 4
        self.supported_image_formats = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff'}
    
    def discover_dataset_items(self) -> List[DatasetItem]:
        """Discover all image-text pairs in the dataset directory."""
        dataset_path = Path(self.dataset_path)
        items = []
        
        if not dataset_path.exists():
            self.logger.error(f"Dataset path does not exist: {dataset_path}")
            return items
        
        # Find all image files
        image_files = []
        for ext in self.supported_image_formats:
            image_files.extend(dataset_path.glob(f"*{ext}"))
        
        self.logger.info(f"Found {len(image_files)} image files")
        
        for image_path in image_files:
            # Look for corresponding text file
            text_path = image_path.with_suffix('.txt')
            
            item = DatasetItem(
                image_path=str(image_path),
                text_path=str(text_path) if text_path.exists() else None
            )
            items.append(item)
        
        self.logger.info(f"Created {len(items)} dataset items")
        return items
    
    def load_single_item(self, item_path: Union[str, Path]) -> DatasetItem:
        """Load a single item by path."""
        item_path = Path(item_path)
        
        if not item_path.exists():
            raise FileNotFoundError(f"Item not found: {item_path}")
        
        # If it's an image, look for corresponding text
        if item_path.suffix.lower() in self.supported_image_formats:
            text_path = item_path.with_suffix('.txt')
            return DatasetItem(
                image_path=str(item_path),
                text_path=str(text_path) if text_path.exists() else None
            )
        
        # If it's a text file, look for corresponding image
        elif item_path.suffix.lower() == '.txt':
            for ext in self.supported_image_formats:
                image_path = item_path.with_suffix(ext)
                if image_path.exists():
                    return DatasetItem(
                        image_path=str(image_path),
                        text_path=str(item_path)
                    )
            
            # No corresponding image found, create text-only item
            with open(item_path, 'r', encoding='utf-8') as f:
                tags = f.read().strip()
            
            return DatasetItem(
                image_path=None,
                text_path=str(item_path),
                tags=tags
            )
        
        else:
            raise ValueError(f"Unsupported file format: {item_path.suffix}")
    
    def process(self, input_data: Any) -> Dict[str, Any]:
        """Process input data and return loaded content."""
        if isinstance(input_data, (str, Path)):
            # Single item path provided
            item = self.load_single_item(input_data)
        elif isinstance(input_data, DatasetItem):
            # DatasetItem provided directly
            item = input_data
        else:
            raise ValueError(f"Unsupported input type: {type(input_data)}")
        
        result = {
            'item_id': item.item_id,
            'image': None,
            'tags': '',
            'image_path': item.image_path,
            'text_path': item.text_path
        }
        
        # Load image if available
        if item.image_path and os.path.exists(item.image_path):
            try:
                result['image'] = item.load_image()
                self.logger.debug(f"Loaded image: {item.image_path}")
            except Exception as e:
                self.logger.error(f"Failed to load image {item.image_path}: {e}")
                result['image'] = None
        
        # Load tags if available
        try:
            result['tags'] = item.load_tags()
            self.logger.debug(f"Loaded tags: {len(result['tags'])} characters")
        except Exception as e:
            self.logger.error(f"Failed to load tags: {e}")
            result['tags'] = ''
        
        return result
    
    def validate_input(self, input_data: Any) -> bool:
        """Validate that input can be processed."""
        if isinstance(input_data, (str, Path)):
            return Path(input_data).exists()
        elif isinstance(input_data, DatasetItem):
            return True
        return False
    
    def get_sample_items(self, n: int = 10) -> List[DatasetItem]:
        """Get a sample of dataset items for testing."""
        all_items = self.discover_dataset_items()
        return all_items[:n] if len(all_items) >= n else all_items
    
    def create_pytorch_dataset(self, items: Optional[List[DatasetItem]] = None, transform=None) -> CharacterDataset:
        """Create a PyTorch Dataset from dataset items."""
        if items is None:
            items = self.discover_dataset_items()
        return CharacterDataset(items, transform=transform)
    
    def create_dataloader(self, items: Optional[List[DatasetItem]] = None, transform=None, 
                         batch_size: Optional[int] = None, shuffle: bool = True) -> DataLoader:
        """Create a PyTorch DataLoader for batch processing."""
        dataset = self.create_pytorch_dataset(items, transform)
        batch_size = batch_size or self.batch_size
        
        return DataLoader(
            dataset,
            batch_size=batch_size,
            shuffle=shuffle,
            num_workers=self.num_workers,
            collate_fn=self._collate_fn
        )
    
    def _collate_fn(self, batch):
        """Custom collate function for DataLoader."""
        images = [item['image'] for item in batch]
        tags = [item['tags'] for item in batch]
        item_ids = [item['item_id'] for item in batch]
        image_paths = [item['image_path'] for item in batch]
        
        return {
            'images': images,
            'tags': tags,
            'item_ids': item_ids,
            'image_paths': image_paths
        }
    
    def create_huggingface_dataset(self, items: Optional[List[DatasetItem]] = None) -> Optional[HFDataset]:
        """Create a HuggingFace Dataset for efficient batch processing."""
        if not HF_AVAILABLE:
            logger.warning("HuggingFace datasets not available. Install with: pip install datasets")
            return None
        
        if items is None:
            items = self.discover_dataset_items()
        
        # Prepare data for HuggingFace Dataset
        data = {
            'image_path': [item.image_path for item in items],
            'tags': [item.load_tags() for item in items],
            'item_id': [item.item_id for item in items]
        }
        
        return HFDataset.from_dict(data)
    
    def process_with_hf_map(self, processing_fn, items: Optional[List[DatasetItem]] = None, 
                           batch_size: Optional[int] = None, num_proc: int = 4) -> Optional[HFDataset]:
        """Process dataset using HuggingFace datasets.map() for efficient batch inference."""
        if not HF_AVAILABLE:
            logger.warning("HuggingFace datasets not available")
            return None
        
        dataset = self.create_huggingface_dataset(items)
        if dataset is None:
            return None
        
        batch_size = batch_size or self.batch_size
        
        # Apply processing function using datasets.map()
        processed_dataset = dataset.map(
            processing_fn,
            batched=True,
            batch_size=batch_size,
            num_proc=num_proc,
            remove_columns=['image_path'],  # Remove to save memory
            desc="Processing character attributes"
        )
        
        return processed_dataset
    
    def batch_process_iterator(self, items: Optional[List[DatasetItem]] = None, 
                              batch_size: Optional[int] = None) -> Iterator[List[DatasetItem]]:
        """Create an iterator for batch processing without loading all data into memory."""
        if items is None:
            items = self.discover_dataset_items()
        
        batch_size = batch_size or self.batch_size
        
        for i in range(0, len(items), batch_size):
            yield items[i:i + batch_size]