caption-flow 0.1.0__py3-none-any.whl → 0.2.1__py3-none-any.whl

caption_flow/utils/image_processor.py

@@ -1,14 +1,20 @@
 """Image preprocessing utilities."""
 
 import asyncio
+import logging
 from concurrent.futures import ProcessPoolExecutor
+from io import BytesIO
 from pathlib import Path
-from typing import List, Any
+from typing import List, Any, Optional, Tuple, Union
 
 import numpy as np
+import requests
 from PIL import Image
 
 
+logger = logging.getLogger(__name__)
+
+
 class ImageProcessor:
     """Handles image loading and preprocessing."""
 
@@ -46,6 +52,120 @@ class ImageProcessor:
 
         return arr
 
+    @staticmethod
+    def process_image_data(img_data: Union[str, bytes, Image.Image]) -> Optional[bytes]:
+        """
+        Process various types of image data into bytes.
+
+        Args:
+            img_data: Can be a URL string, bytes, or PIL Image
+
+        Returns:
+            Image data as bytes, or None if processing failed
+        """
+        try:
+            if isinstance(img_data, str):
+                # It's a URL - download the image
+                try:
+                    # Download with timeout
+                    response = requests.get(
+                        img_data,
+                        timeout=30,
+                        headers={"User-Agent": "Mozilla/5.0 (captionflow-dataset-loader)"},
+                    )
+                    response.raise_for_status()
+                    image_data = response.content
+
+                    # Verify it's an image by trying to open it
+                    img = Image.open(BytesIO(image_data))
+                    img.verify()  # Verify it's a valid image
+
+                    return image_data
+
+                except Exception as e:
+                    logger.error(f"Failed to download image from {img_data}: {e}")
+                    return None
+
+            elif hasattr(img_data, "__class__") and "Image" in str(img_data.__class__):
+                # It's a PIL Image object
+                import io
+
+                # Save as PNG bytes
+                img_bytes = io.BytesIO()
+                # Convert to RGB
+                img_data = img_data.convert("RGB")
+                img_data.save(img_bytes, format="PNG")
+                return img_bytes.getvalue()
+
+            elif isinstance(img_data, bytes):
+                # Already bytes - validate it's an image
+                try:
+                    img = Image.open(BytesIO(img_data))
+                    img.verify()
+                    return img_data
+                except Exception as e:
+                    logger.error(f"Invalid image data: {e}")
+                    return None
+
+            else:
+                logger.warning(f"Unknown image data type: {type(img_data)}")
+                return None
+
+        except Exception as e:
+            logger.error(f"Error processing image data: {e}")
+            import traceback
+
+            logger.error(traceback.format_exc())
+            return None
+
+    @staticmethod
+    def prepare_for_inference(image: Image.Image) -> Image.Image:
+        """
+        Prepare image for inference, handling transparency and mostly black/white images.
+
+        Args:
+            image: PIL Image to prepare
+
+        Returns:
+            Prepared PIL Image
+        """
+        # Convert to RGBA to handle transparency
+        img_rgba = image.convert("RGBA")
+        rgb_img = img_rgba.convert("RGB")
+        np_img = np.array(rgb_img)
+
+        # Calculate percentage of pixels that are (0,0,0) or (255,255,255)
+        total_pixels = np_img.shape[0] * np_img.shape[1]
+        black_pixels = np.all(np_img == [0, 0, 0], axis=-1).sum()
+        white_pixels = np.all(np_img == [255, 255, 255], axis=-1).sum()
+        black_pct = black_pixels / total_pixels
+        white_pct = white_pixels / total_pixels
+
+        threshold = 0.90  # 90% threshold
+
+        is_mostly_black = black_pct >= threshold
+        is_mostly_white = white_pct >= threshold
+
+        if is_mostly_black or is_mostly_white:
+            # Replace background with opposite color for better contrast
+            bg_color = (255, 255, 255) if is_mostly_black else (0, 0, 0)
+            background = Image.new("RGB", img_rgba.size, bg_color)
+            # Use alpha channel as mask if present
+            if img_rgba.mode == "RGBA":
+                background.paste(img_rgba.convert("RGB"), mask=img_rgba.split()[3])
+            else:
+                background.paste(img_rgba.convert("RGB"))
+
+            color_type = "black" if is_mostly_black else "white"
+            pct = black_pct if is_mostly_black else white_pct
+            logger.debug(
+                f"Image is {pct*100:.1f}% {color_type}; background replaced with {bg_color}"
+            )
+
+            return background
+        else:
+            return rgb_img
+
     def shutdown(self):
         """Shutdown the executor."""
         self.executor.shutdown(wait=True)

caption_flow/utils/prompt_template.py

@@ -0,0 +1,137 @@
+"""Prompt template system for dynamic column substitution."""
+
+import re
+import logging
+from typing import Dict, Any, List, Optional
+
+logger = logging.getLogger(__name__)
+
+
+class PromptTemplate:
+    """Handles prompt templates with column substitution."""
+
+    # Pattern to match {column:column_name} or {col:column_name}
+    COLUMN_PATTERN = re.compile(r"\{(?:column|col):([\w-]+)\}")
+
+    def __init__(self, template: str):
+        """
+        Initialize with a prompt template.
+
+        Args:
+            template: Prompt template string, e.g.
+                     "describe this image. tags: {column:user_tags}"
+        """
+        self.template = template
+        self.required_columns = self._extract_columns()
+
+    def _extract_columns(self) -> List[str]:
+        """Extract required column names from template."""
+        matches = self.COLUMN_PATTERN.findall(self.template)
+        return list(set(matches))  # Remove duplicates
+
+    def format(self, item_data: Dict[str, Any]) -> str:
+        """
+        Format the template with actual column values.
+
+        Args:
+            item_data: Dictionary containing column values from dataset
+
+        Returns:
+            Formatted prompt string
+        """
+        prompt = self.template
+
+        # Replace all column references
+        for match in self.COLUMN_PATTERN.finditer(self.template):
+            full_match = match.group(0)  # e.g., {column:user_tags}
+            column_name = match.group(1)  # e.g., user_tags
+
+            # Get column value with fallback
+            value = item_data.get(column_name, "")
+
+            # Handle different value types
+            if value is None:
+                value = ""
+            elif isinstance(value, list):
+                # Join list items with commas
+                value = ", ".join(str(v) for v in value if v)
+            elif not isinstance(value, str):
+                value = str(value)
+
+            # Replace in prompt
+            prompt = prompt.replace(full_match, value)
+
+        return prompt.strip()
+
+    def validate_columns(self, available_columns: List[str]) -> List[str]:
+        """
+        Validate that required columns are available.
+
+        Returns:
+            List of missing column names
+        """
+        missing = []
+        for col in self.required_columns:
+            if col not in available_columns:
+                missing.append(col)
+        return missing
+
+
+class PromptTemplateManager:
+    """Manages multiple prompt templates."""
+
+    def __init__(self, prompts: List[str]):
+        """
+        Initialize with list of prompt strings (which may contain templates).
+
+        Args:
+            prompts: List of prompt strings
+        """
+        self.templates = [PromptTemplate(p) for p in prompts]
+        self._all_required_columns = None
+
+    @property
+    def required_columns(self) -> List[str]:
+        """Get all required columns across all templates."""
+        if self._all_required_columns is None:
+            cols = set()
+            for template in self.templates:
+                cols.update(template.required_columns)
+            self._all_required_columns = list(cols)
+        return self._all_required_columns
+
+    def format_all(self, item_data: Dict[str, Any]) -> List[str]:
+        """
+        Format all templates with item data.
+
+        Args:
+            item_data: Dictionary containing column values
+
+        Returns:
+            List of formatted prompts
+        """
+        formatted = []
+        for template in self.templates:
+            try:
+                prompt = template.format(item_data)
+                formatted.append(prompt)
+            except Exception as e:
+                logger.error(f"Error formatting prompt template '{template.template}': {e}")
+                # Fall back to raw template
+                formatted.append(template.template)
+
+        return formatted
+
+    def validate_all(self, available_columns: List[str]) -> Dict[str, List[str]]:
+        """
+        Validate all templates against available columns.
+
+        Returns:
+            Dict mapping template string to list of missing columns
+        """
+        issues = {}
+        for template in self.templates:
+            missing = template.validate_columns(available_columns)
+            if missing:
+                issues[template.template] = missing
+        return issues

caption_flow/utils/shard_processor.py

@@ -0,0 +1,315 @@
+"""Shard processing abstraction for different dataset types."""
+
+import io
+import logging
+import time
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Generator, Tuple, Optional, Dict, Any
+from dataclasses import dataclass
+from threading import Event
+import shlex
+
+import webdataset as wds
+from PIL import Image
+
+from .dataset_loader import DatasetLoader
+
+logger = logging.getLogger(__name__)
+
+
+class ShardProcessor(ABC):
+    """Abstract base for processing dataset shards."""
+
+    @abstractmethod
+    def iterate_chunk(
+        self,
+        chunk,
+        dataset_loader: Optional[DatasetLoader],
+        should_stop: Event,
+        connected: Event,
+    ) -> Generator[Tuple[str, str, bytes], None, None]:
+        """
+        Iterate through items in a chunk.
+
+        Yields:
+            Tuple of (key, url, image_data)
+        """
+        pass
+
+
+class HFDatasetShardProcessor(ShardProcessor):
+    """Processor for HuggingFace virtual dataset shards."""
+
+    def iterate_chunk(
+        self,
+        chunk,
+        dataset_loader: Optional[DatasetLoader],
+        should_stop: Event,
+        connected: Event,
+    ) -> Generator[Tuple[str, str, bytes], None, None]:
+        """Process HuggingFace virtual shard chunk."""
+        if not dataset_loader:
+            logger.error("No dataset loader configured for HuggingFace dataset shard")
+            return
+
+        # Get unprocessed ranges
+        unprocessed_ranges = getattr(chunk, "unprocessed_ranges", [(0, chunk.chunk_size - 1)])
+
+        logger.info(
+            f"Processing HF dataset chunk {chunk.chunk_id} with ranges: {unprocessed_ranges}"
+        )
+
+        items_processed = 0
+        current_idx = 0
+
+        # Construct proper virtual shard URL
+        parts = chunk.shard_url.split("_chunk_")
+        if len(parts) == 2:
+            base_path = parts[0]
+            virtual_shard_url = f"{base_path}:chunk:{chunk.start_index}"
+        else:
+            virtual_shard_url = chunk.shard_url
+
+        logger.debug(f"Using virtual shard URL: {virtual_shard_url}")
+
+        # Iterate through the virtual shard
+        for key, url, image_data in dataset_loader.iterate_shard(virtual_shard_url):
+            # Check if we should stop
+            if should_stop.is_set() or not connected.is_set():
+                logger.info(f"Stopping chunk processing early due to disconnect")
+                break
+
+            # Check if current index is in any unprocessed range
+            in_range = any(start <= current_idx <= end for start, end in unprocessed_ranges)
+
+            if not in_range:
+                current_idx += 1
+                continue  # Skip already processed items
+
+            # Check if we've processed enough for this chunk
+            if current_idx >= chunk.chunk_size:
+                break
+
+            items_processed += 1
+            current_idx += 1
+            yield key, url, image_data
+
+        logger.info(
+            f"HF dataset chunk {chunk.chunk_id}: yielded {items_processed} items "
+            f"from ranges {unprocessed_ranges}"
+        )
+
+    def iterate_chunk_with_metadata(
+        self,
+        chunk,
+        dataset_loader: Optional[DatasetLoader],
+        should_stop: Event,
+        connected: Event,
+    ) -> Generator[Tuple[str, str, bytes, Dict[str, Any]], None, None]:
+        """
+        Process HuggingFace virtual shard chunk with metadata.
+
+        Yields:
+            Tuple of (key, url, image_data, metadata)
+        """
+        if not dataset_loader:
+            logger.error("No dataset loader configured for HuggingFace dataset shard")
+            return
+
+        # Get unprocessed ranges
+        unprocessed_ranges = getattr(chunk, "unprocessed_ranges", [(0, chunk.chunk_size - 1)])
+
+        logger.info(
+            f"Processing HF dataset chunk {chunk.chunk_id} with ranges: {unprocessed_ranges}"
+        )
+
+        items_processed = 0
+        current_idx = 0
+
+        # Construct proper virtual shard URL
+        parts = chunk.shard_url.split("_chunk_")
+        if len(parts) == 2:
+            base_path = parts[0]
+            virtual_shard_url = f"{base_path}:chunk:{chunk.start_index}"
+        else:
+            virtual_shard_url = chunk.shard_url
+
+        logger.debug(f"Using virtual shard URL: {virtual_shard_url}")
+
+        # Use the new iterate method that includes metadata
+        for key, url, image_data, metadata in dataset_loader.iterate_shard_with_metadata(
+            virtual_shard_url
+        ):
+            # Check if we should stop
+            if should_stop.is_set() or not connected.is_set():
+                logger.info(f"Stopping chunk processing early due to disconnect")
+                break
+
+            # Check if current index is in any unprocessed range
+            in_range = any(start <= current_idx <= end for start, end in unprocessed_ranges)
+
+            if not in_range:
+                current_idx += 1
+                continue  # Skip already processed items
+
+            # Check if we've processed enough for this chunk
+            if current_idx >= chunk.chunk_size:
+                break
+
+            items_processed += 1
+            current_idx += 1
+            yield key, url, image_data, metadata
+
+        logger.info(
+            f"HF dataset chunk {chunk.chunk_id}: yielded {items_processed} items "
+            f"from ranges {unprocessed_ranges}"
+        )
+
+
+class WebDatasetShardProcessor(ShardProcessor):
+    """Processor for WebDataset tar shards with range support."""
+
+    def __init__(self, hf_token: Optional[str] = None, dataset_type: str = "local"):
+        self.hf_token = hf_token
+        self.dataset_type = dataset_type
+
+    def iterate_chunk(
+        self,
+        chunk,
+        dataset_loader: Optional[DatasetLoader],
+        should_stop: Event,
+        connected: Event,
+    ) -> Generator[Tuple[str, str, bytes], None, None]:
+        """Process WebDataset shard chunk with unprocessed ranges."""
+        # Get unprocessed ranges
+        unprocessed_ranges = getattr(chunk, "unprocessed_ranges", [(0, chunk.chunk_size - 1)])
+
+        logger.info(
+            f"Processing WebDataset chunk {chunk.chunk_id} with ranges: {unprocessed_ranges}"
+        )
+
+        # Create WebDataset pipeline
+        if self.dataset_type == "huggingface" and not chunk.shard_url.startswith("hf_dataset:"):
+            # Use curl with auth for HuggingFace WebDataset
+            url_cmd = f"pipe:curl -s -L -H 'Authorization:Bearer {shlex.quote(self.hf_token)}' {shlex.quote(chunk.shard_url)} || true"
+            ds = wds.DataPipeline(
+                wds.SimpleShardList(url_cmd),
+                wds.tarfile_to_samples(),
+                wds.to_tuple("__key__", "jpg;png;jpeg;webp;jxl"),
+            )
+        else:
+            # Local file
+            ds = wds.DataPipeline(
+                wds.SimpleShardList(chunk.shard_url),
+                wds.tarfile_to_samples(),
+                wds.to_tuple("__key__", "jpg;png;jpeg;webp;jxl"),
+            )
+
+        # Process items
+        current_idx = 0
+        items_yielded = 0
+
+        for key, image_data in ds:
+            # Check if we should stop
+            if should_stop.is_set() or not connected.is_set():
+                logger.info(f"Stopping WebDataset chunk processing early due to disconnect")
+                break
+
+            # Calculate relative index within chunk
+            relative_idx = current_idx - chunk.start_index
+
+            # Skip items before chunk start
+            if current_idx < chunk.start_index:
+                current_idx += 1
+                continue
+
+            # Stop if beyond chunk
+            if relative_idx >= chunk.chunk_size:
+                break
+
+            # Check if current index is in any unprocessed range
+            in_range = any(start <= relative_idx <= end for start, end in unprocessed_ranges)
+
+            if in_range:
+                items_yielded += 1
+                yield key, chunk.shard_url, image_data
+
+            current_idx += 1
+
+        logger.info(
+            f"WebDataset chunk {chunk.chunk_id}: yielded {items_yielded} items "
+            f"from ranges {unprocessed_ranges}"
+        )
+
+    def iterate_chunk_with_metadata(
+        self,
+        chunk,
+        dataset_loader: Optional[DatasetLoader],
+        should_stop: Event,
+        connected: Event,
+    ) -> Generator[Tuple[str, str, bytes, Dict[str, Any]], None, None]:
+        """Process WebDataset shard chunk with metadata and range support."""
+        # Get unprocessed ranges
+        unprocessed_ranges = getattr(chunk, "unprocessed_ranges", [(0, chunk.chunk_size - 1)])
+
+        logger.info(
+            f"Processing WebDataset chunk {chunk.chunk_id} with ranges: {unprocessed_ranges}"
+        )
+
+        # Create WebDataset pipeline
+        if self.dataset_type == "huggingface" and not chunk.shard_url.startswith("hf_dataset:"):
+            # Use curl with auth for HuggingFace WebDataset
+            url_cmd = f"pipe:curl -s -L -H 'Authorization:Bearer {shlex.quote(self.hf_token)}' {shlex.quote(chunk.shard_url)} || true"
+            ds = wds.DataPipeline(
+                wds.SimpleShardList(url_cmd),
+                wds.tarfile_to_samples(),
+                wds.to_tuple("__key__", "jpg;png;jpeg;webp;jxl"),
+            )
+        else:
+            # Local file
+            ds = wds.DataPipeline(
+                wds.SimpleShardList(chunk.shard_url),
+                wds.tarfile_to_samples(),
+                wds.to_tuple("__key__", "jpg;png;jpeg;webp;jxl"),
+            )
+
+        # Process items
+        absolute_idx = 0  # Absolute index in the shard
+        items_yielded = 0
+
+        for key, image_data in ds:
+            # Check if we should stop
+            if should_stop.is_set() or not connected.is_set():
+                logger.info(f"Stopping WebDataset chunk processing early due to disconnect")
+                break
+
+            # Skip items before chunk start
+            if absolute_idx < chunk.start_index:
+                absolute_idx += 1
+                continue
+
+            # Calculate relative index within chunk
+            relative_idx = absolute_idx - chunk.start_index
+
+            # Stop if beyond chunk
+            if relative_idx >= chunk.chunk_size:
+                break
+
+            # Check if current index is in any unprocessed range
+            in_range = any(start <= relative_idx <= end for start, end in unprocessed_ranges)
+
+            if in_range:
+                # Create metadata with the relative index
+                metadata = {
+                    "_chunk_relative_index": relative_idx,
+                }
+                items_yielded += 1
+                yield key, chunk.shard_url, image_data, metadata
+
+            absolute_idx += 1
+
+        logger.info(
+            f"WebDataset chunk {chunk.chunk_id}: yielded {items_yielded} items "
+            f"from ranges {unprocessed_ranges}"
+        )

caption_flow/utils/shard_tracker.py

@@ -0,0 +1,87 @@
+"""Shard tracking using CheckpointTracker base class."""
+
+from pathlib import Path
+from typing import Dict, Any, List, Set
+
+from .checkpoint_tracker import CheckpointTracker
+
+
+class ShardTracker(CheckpointTracker):
+    """Tracks shard processing progress."""
+
+    def __init__(self, checkpoint_path: Path):
+        """Initialize shard tracker with checkpoint file."""
+        self.completed_shards: Set[str] = set()
+        self.partial_shards: Dict[str, Dict[str, Any]] = {}
+        super().__init__(checkpoint_path)
+
+    def _get_default_state(self) -> Dict[str, Any]:
+        """Return default state structure for new checkpoints."""
+        return {"completed_shards": [], "partial_shards": {}}
+
+    def _deserialize_state(self, data: Dict[str, Any]) -> None:
+        """Deserialize loaded data into instance state."""
+        self.completed_shards = set(data.get("completed_shards", []))
+        self.partial_shards = data.get("partial_shards", {})
+
+    def _serialize_state(self) -> Dict[str, Any]:
+        """Serialize instance state for saving."""
+        return {
+            "completed_shards": list(self.completed_shards),
+            "partial_shards": self.partial_shards,
+        }
+
+    def mark_complete(self, shard_name: str) -> None:
+        """Mark a shard as complete."""
+        self.completed_shards.add(shard_name)
+        if shard_name in self.partial_shards:
+            del self.partial_shards[shard_name]
+        self.save()
+
+    def update_partial(self, shard_name: str, processed_keys: List[str]) -> None:
+        """Update partial progress for a shard."""
+        self.partial_shards[shard_name] = {"keys": processed_keys, "count": len(processed_keys)}
+        self.save()
+
+    def get_processed_keys(self, shard_name: str) -> Set[str]:
+        """Get set of processed keys for a shard."""
+        if shard_name in self.completed_shards:
+            return set()  # All done
+
+        if shard_name in self.partial_shards:
+            return set(self.partial_shards[shard_name].get("keys", []))
+
+        return set()
+
+    def is_complete(self, shard_name: str) -> bool:
+        """Check if a shard is complete."""
+        return shard_name in self.completed_shards
+
+    def get_remaining_shards(self, all_shards: List[str]) -> List[str]:
+        """Get list of shards that still need processing."""
+        remaining = []
+        for s in all_shards:
+            # Extract shard name properly for both regular and virtual shards
+            if s.startswith("hf_dataset:"):
+                shard_name = s  # Use full virtual shard ID
+            else:
+                shard_name = Path(s).stem
+
+            if shard_name not in self.completed_shards:
+                remaining.append(s)
+
+        return remaining
+
+    def get_stats(self) -> Dict[str, Any]:
+        """Get shard tracking statistics."""
+        base_stats = super().get_stats()
+        base_stats.update(
+            {
+                "completed_shards": len(self.completed_shards),
+                "partial_shards": len(self.partial_shards),
+                "total_partial_keys": sum(
+                    len(data.get("keys", [])) for data in self.partial_shards.values()
+                ),
+            }
+        )
+        return base_stats