analyzer.py

"""
analyzer.py
------------
Top-level orchestration. Given raw bytes of an uploaded sheet (single
sprite image or full sprite sheet), runs the complete pipeline:

    load -> detect background -> segment -> extract each sprite ->
    quality analysis -> hashing -> duplicate check -> persist to
    staging + DB

Every sprite ends up as a `SpriteRecord` row with status='pending' in
staging/, ready for manual review (accept/reject/edit) in app.py.
"""

from __future__ import annotations

from dataclasses import dataclass

import database as db
from auto_tagger import suggest_tags
from background import BackgroundInfo, detect_background
from config import AUTO_TAG_ENABLED, LOW_CONFIDENCE_THRESHOLD, STAGING_DIR
from dataset_manager import staging_path_for
from duplicate_detector import compute_dhash, compute_sha256, find_duplicates
from image_utils import UnsupportedImageError, image_to_png_bytes, load_image_rgba, to_numpy, to_pil
from scoring import analyze_sprite
from segmentation import segment_sheet
from sprite_extractor import extract_all


@dataclass
class SheetProcessingResult:
    sheet_id: str
    filename: str
    width: int
    height: int
    sprite_count: int
    sprite_ids: list[str]
    background_color: tuple
    low_confidence_count: int
    error: str | None = None


def process_sheet(
    file_bytes: bytes,
    filename: str,
    padding: int = 0,
) -> SheetProcessingResult:
    """
    Process a single uploaded file (sprite sheet OR a lone sprite -- a lone
    sprite is just a sheet that segments into exactly one region) end to
    end and persist every detected sprite to staging + the DB.
    """
    sheet_id = db.new_id("sheet")

    try:
        img = load_image_rgba(file_bytes)
    except UnsupportedImageError as exc:
        return SheetProcessingResult(sheet_id, filename, 0, 0, 0, [], (0, 0, 0, 0), 0, error=str(exc))

    rgba = to_numpy(img)
    h, w = rgba.shape[:2]

    if h == 0 or w == 0:
        return SheetProcessingResult(sheet_id, filename, w, h, 0, [], (0, 0, 0, 0), 0, error="Image has zero dimensions")

    db.insert_sheet(sheet_id, filename, w, h)

    # Cache the original sheet bytes so later manual edits (move/resize/
    # merge/split bounding boxes) can re-crop EXACT pixels from the true
    # source instead of re-cropping from an already-cropped sprite (which
    # would either clip needed pixels or require resampling). Sheets are
    # typically small (a few hundred KB to a few MB), so this is cheap
    # relative to the sprite PNGs already being written to staging.
    (STAGING_DIR / f"{sheet_id}_source.png").write_bytes(image_to_png_bytes(img))

    bg_info: BackgroundInfo = detect_background(rgba)
    foreground = ~bg_info.mask

    if not foreground.any():
        # Entire sheet matched the background -- nothing to extract, but
        # this is not a crash, just zero sprites found.
        db.update_sheet_sprite_count(sheet_id, 0)
        return SheetProcessingResult(sheet_id, filename, w, h, 0, [], bg_info.color, 0)

    # The sheet's own pixels are passed alongside the mask so segmentation
    # can measure edge periodicity -- the only signal that recovers the
    # cell pitch of a tileset drawn edge to edge, where the mask alone is
    # one undifferentiated blob.
    detections, _labels, _groups = segment_sheet(foreground, rgba=rgba)
    extracted = extract_all(rgba, bg_info, detections, source_file=filename, padding=padding)

    existing_hashes = db.existing_hashes_for_dup_check()

    sprite_ids: list[str] = []
    low_confidence_count = 0

    for sprite in extracted:
        sprite_id = db.new_id("sprite")

        quality = analyze_sprite(sprite.rgba)
        sha256 = compute_sha256(image_to_png_bytes(to_pil(sprite.rgba)))
        phash = compute_dhash(sprite.rgba)

        dup_matches = find_duplicates(sha256, phash, existing_hashes)
        is_duplicate_of = dup_matches[0].sprite_id if dup_matches else None
        duplicate_similarity = dup_matches[0].similarity_percent if dup_matches else None

        # Persist both variants (background-kept, background-transparent)
        # to staging so the export choice can be made later without
        # recomputation.
        kept_path = staging_path_for(sprite_id, transparent=False)
        transparent_path = staging_path_for(sprite_id, transparent=True)
        to_pil(sprite.rgba).save(kept_path, format="PNG")
        to_pil(sprite.rgba_transparent_bg).save(transparent_path, format="PNG")

        if sprite.confidence < LOW_CONFIDENCE_THRESHOLD:
            low_confidence_count += 1

        # A machine suggestion so the sprite doesn't land in the dataset
        # as a bare "unknown" nobody will ever get around to labelling by
        # hand. It is always overridable, is marked as unreviewed, and
        # leaves the category blank rather than guessing when it isn't
        # reasonably sure -- see auto_tagger's module docstring.
        suggestion = (
            suggest_tags(
                sprite.rgba,
                segmentation_source=sprite.source,
                component_count=sprite.component_count,
                sheet_sprite_count=len(extracted),
            )
            if AUTO_TAG_ENABLED
            else None
        )

        rec = db.SpriteRecord(
            id=sprite_id,
            sheet_id=sheet_id,
            source_file=filename,
            sprite_index=sprite.sprite_index,
            x=sprite.bbox.x, y=sprite.bbox.y, width=sprite.bbox.width, height=sprite.bbox.height,
            png_path=str(kept_path),
            png_path_transparent=str(transparent_path),
            sha256=sha256,
            phash=phash,
            quality=quality.to_dict(),
            confidence=sprite.confidence,
            confidence_reasons=sprite.confidence_reasons,
            component_count=sprite.component_count,
            category=suggestion.category if suggestion else "unknown",
            style=suggestion.style if suggestion else "unknown",
            tags=suggestion.tags if suggestion else [],
            status="pending",
            is_duplicate_of=is_duplicate_of,
            duplicate_similarity=duplicate_similarity,
            segmentation_source=sprite.source,
            auto_tagged=bool(suggestion),
            auto_tag_confidence=suggestion.confidence if suggestion else None,
        )
        db.insert_sprite(rec)
        sprite_ids.append(sprite_id)

        # Newly staged sprite becomes comparable for subsequent sprites
        # within the SAME sheet too (e.g. two identical tree sprites on one
        # sheet should flag each other).
        existing_hashes.append({"id": sprite_id, "sha256": sha256, "phash": phash})

    db.update_sheet_sprite_count(sheet_id, len(sprite_ids))

    return SheetProcessingResult(
        sheet_id=sheet_id,
        filename=filename,
        width=w,
        height=h,
        sprite_count=len(sprite_ids),
        sprite_ids=sprite_ids,
        background_color=bg_info.color,
        low_confidence_count=low_confidence_count,
    )


def process_multiple(files: list[tuple[str, bytes]], padding: int = 0) -> list[SheetProcessingResult]:
    """Process several sheets (e.g. everything extracted from one ZIP)."""
    results = []
    for filename, data in files:
        results.append(process_sheet(data, filename, padding=padding))
    return results