slahiri-ComfyUI-AI-Photography-Toolkit/core/compose/classifier/base.py

"""Base data structures for token classification.

This module defines the classification result structure and
the classified image container.
"""

from dataclasses import dataclass, field
from typing import Dict, List, Optional, Any
from datetime import datetime

from .categories import (
    CanonicalCategory,
    SubjectDetailType,
    CATEGORY_ORDER,
    SUBJECT_DETAIL_ORDER,
)
from ..tokenizer.base import ImageToken


@dataclass
class TokenClassification:
    """Classification result for a single token.

    A token may belong to multiple categories (e.g., "looking at viewer"
    is both Action/Pose and Subject Details/eyes). The primary category
    is used for assembly ordering.

    Attributes:
        token: The original ImageToken being classified
        primary_category: Main category assignment
        secondary_categories: Additional category assignments
        subcategory: For SUBJECT_DETAILS, the specific subcategory
        confidence: Classification confidence (0.0-1.0)
        classifier_layer: Which layer made the classification (1-5)
        metadata: Additional classification context
    """
    token: ImageToken
    primary_category: CanonicalCategory
    secondary_categories: List[CanonicalCategory] = field(default_factory=list)
    subcategory: Optional[SubjectDetailType] = None
    confidence: float = 1.0
    classifier_layer: int = 0  # 1=deterministic, 2=source, 3=dictionary, 4=embedding, 5=uncategorized
    metadata: Dict[str, Any] = field(default_factory=dict)

    @property
    def all_categories(self) -> List[CanonicalCategory]:
        """Get all categories (primary + secondary)."""
        return [self.primary_category] + self.secondary_categories

    @property
    def is_subject_detail(self) -> bool:
        """Check if this is a subject detail token."""
        return self.primary_category == CanonicalCategory.SUBJECT_DETAILS

    @property
    def is_filtered(self) -> bool:
        """Check if this token was filtered (META pattern)."""
        return self.primary_category == CanonicalCategory.FILTERED_META

    @property
    def is_uncategorized(self) -> bool:
        """Check if this token couldn't be categorized."""
        return self.primary_category == CanonicalCategory.UNCATEGORIZED

    def __repr__(self):
        cat = self.primary_category.value
        if self.subcategory:
            cat = f"{cat}/{self.subcategory.value}"
        return f"TokenClassification({self.token.text!r} → {cat})"


@dataclass
class ClassifiedImage:
    """Collection of classified tokens for a single image.

    Organizes tokens by category for efficient prompt assembly.
    """
    # Tokens organized by category
    categories: Dict[CanonicalCategory, List[TokenClassification]] = field(default_factory=dict)

    # Subject details organized by subcategory
    subject_details: Dict[SubjectDetailType, List[TokenClassification]] = field(default_factory=dict)

    # All classifications (for iteration)
    all_classifications: List[TokenClassification] = field(default_factory=list)

    # Filtered tokens (META patterns)
    filtered: List[TokenClassification] = field(default_factory=list)

    # Metadata
    source_file: Optional[str] = None
    extraction_timestamp: Optional[datetime] = None
    total_tokens_input: int = 0
    image_info: Dict[str, Any] = field(default_factory=dict)

    def __post_init__(self):
        """Initialize category dictionaries."""
        # Ensure all categories have empty lists
        for cat in CanonicalCategory:
            if cat not in self.categories:
                self.categories[cat] = []

        # Ensure all subcategories have empty lists
        for subcat in SubjectDetailType:
            if subcat not in self.subject_details:
                self.subject_details[subcat] = []

    def add(self, classification: TokenClassification) -> None:
        """Add a classification to the appropriate category."""
        self.all_classifications.append(classification)

        # Handle filtered tokens
        if classification.is_filtered:
            self.filtered.append(classification)
            return

        # Add to main category
        cat = classification.primary_category
        if cat not in self.categories:
            self.categories[cat] = []
        self.categories[cat].append(classification)

        # Add to subcategory if subject detail
        if classification.is_subject_detail and classification.subcategory:
            subcat = classification.subcategory
            if subcat not in self.subject_details:
                self.subject_details[subcat] = []
            self.subject_details[subcat].append(classification)

    def get_category(self, category: CanonicalCategory) -> List[TokenClassification]:
        """Get all classifications for a category."""
        return self.categories.get(category, [])

    def get_subject_detail(self, subcategory: SubjectDetailType) -> List[TokenClassification]:
        """Get all classifications for a subject detail subcategory."""
        return self.subject_details.get(subcategory, [])

    def get_tokens_by_category(self, category: CanonicalCategory) -> List[ImageToken]:
        """Get just the tokens (not full classifications) for a category."""
        return [c.token for c in self.get_category(category)]

    def get_texts_by_category(self, category: CanonicalCategory) -> List[str]:
        """Get just the text values for a category."""
        return [c.token.text for c in self.get_category(category)]

    def iter_by_order(self):
        """Iterate categories in canonical output order."""
        for cat in CATEGORY_ORDER:
            classifications = self.get_category(cat)
            if classifications:
                yield cat, classifications

    def iter_subject_details_by_order(self):
        """Iterate subject details in canonical order."""
        for subcat in SUBJECT_DETAIL_ORDER:
            classifications = self.get_subject_detail(subcat)
            if classifications:
                yield subcat, classifications

    @property
    def has_subject(self) -> bool:
        """Check if any subject tokens exist."""
        return bool(self.get_category(CanonicalCategory.SUBJECT))

    @property
    def subject_count(self) -> int:
        """Get estimated subject count from subject tokens."""
        subjects = self.get_texts_by_category(CanonicalCategory.SUBJECT)
        # Look for count indicators
        for text in subjects:
            if "2girl" in text or "2boy" in text or "couple" in text:
                return 2
            if "3girl" in text or "3boy" in text or "group" in text:
                return 3
            if "multiple" in text or "crowd" in text:
                return 4
        return 1  # Default to solo

    @property
    def subject_gender(self) -> Optional[str]:
        """Infer subject gender from subject tokens."""
        subjects = self.get_texts_by_category(CanonicalCategory.SUBJECT)
        for text in subjects:
            if "girl" in text or "woman" in text or "female" in text or "lady" in text:
                return "female"
            if "boy" in text or "man" in text or "male" in text or "gentleman" in text:
                return "male"
        return None

    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary format for JSON serialization."""
        result = {
            "total_tokens": len(self.all_classifications),
            "filtered_count": len(self.filtered),
            "categories": {},
        }

        for cat in CATEGORY_ORDER:
            classifications = self.get_category(cat)
            if classifications:
                result["categories"][cat.value] = [
                    {
                        "text": c.token.text,
                        "confidence": c.token.confidence,
                        "source": c.token.source.value,
                        "subcategory": c.subcategory.value if c.subcategory else None,
                    }
                    for c in classifications
                ]

        if self.filtered:
            result["filtered"] = [c.token.text for c in self.filtered]

        return result

    def __len__(self):
        return len(self.all_classifications)

    def __repr__(self):
        counts = {cat.value: len(cls) for cat, cls in self.categories.items() if cls}
        return f"ClassifiedImage({counts})"