opsd_utils/privileged/image_utils.py

"""Image loading and crop utilities for privileged teacher dual-image forward."""
from __future__ import annotations

from typing import Any, Optional

from PIL import Image

from data_utils.paths import resolve_image_path
from data_utils.privileged_schema import resolve_crop_bbox
from opsd_utils import debug_log as opsd_debug


def load_rgb(image: Any) -> Optional[Image.Image]:
    """Load sample image as RGB PIL from path or in-memory object."""
    if image is None:
        return None
    if isinstance(image, Image.Image):
        return image.convert("RGB") if image.mode != "RGB" else image
    if isinstance(image, str):
        path = resolve_image_path(image)
        try:
            img = Image.open(path)
            return img.convert("RGB")
        except (FileNotFoundError, OSError):
            opsd_debug.log("privileged_image", "load_rgb failed", path=path)
            return None
    return None


def center_crop(img: Image.Image, margin_ratio: float = 0.25) -> Image.Image:
    w, h = img.size
    margin_w = int(w * margin_ratio)
    margin_h = int(h * margin_ratio)
    return img.crop((margin_w, margin_h, w - margin_w, h - margin_h))


def crop_image(
    img: Image.Image,
    bbox_norm: Optional[list[float]] = None,
    strategy: str = "center",
    margin_ratio: float = 0.25,
    fallback_reason: Optional[str] = None,
) -> tuple[Image.Image, str]:
    """
    Crop image using C2 normalized bbox or center fallback.
    Returns (cropped_image, crop_strategy_used).
    """
    if bbox_norm is not None and strategy in ("bbox", "heuristic", "bbox_then_center"):
        w, h = img.size
        x0 = int(bbox_norm[0] * w)
        y0 = int(bbox_norm[1] * h)
        x1 = int(bbox_norm[2] * w)
        y1 = int(bbox_norm[3] * h)
        x0, x1 = max(0, min(x0, w - 1)), max(1, min(x1, w))
        y0, y1 = max(0, min(y0, h - 1)), max(1, min(y1, h))
        if x1 > x0 and y1 > y0:
            used = strategy if strategy != "bbox_then_center" else "bbox"
            opsd_debug.log(
                "privileged_image",
                "crop_image bbox",
                strategy=used,
                bbox_norm=bbox_norm,
                crop_px=(x0, y0, x1, y1),
                fallback_reason=fallback_reason,
            )
            return img.crop((x0, y0, x1, y1)), used

    crop = center_crop(img, margin_ratio=margin_ratio)
    used = "center_fallback" if fallback_reason else "center"
    opsd_debug.log(
        "privileged_image",
        "crop_image center",
        strategy=used,
        bbox_norm=bbox_norm,
        margin_ratio=margin_ratio,
        fallback_reason=fallback_reason,
    )
    return crop, used


def heuristic_crop_from_visual_fact(
    img: Image.Image,
    sample: dict[str, Any],
    crop_cfg: Optional[dict[str, Any]] = None,
) -> tuple[Image.Image, str, Optional[list[float]]]:
    """D2 with D1 fallback: heuristic bbox from visual_fact, else center crop."""
    crop_cfg = crop_cfg or {}
    margin_ratio = float(crop_cfg.get("margin_ratio", 0.25))
    bbox_norm, strategy = resolve_crop_bbox(sample, crop_cfg)
    fallback_reason = None
    if strategy == "center" and sample.get("visual_fact"):
        fallback_reason = "heuristic_failed"
    crop, used = crop_image(
        img,
        bbox_norm=bbox_norm,
        strategy=strategy if bbox_norm else "center",
        margin_ratio=margin_ratio,
        fallback_reason=fallback_reason,
    )
    return crop, used, bbox_norm


def resolve_teacher_images(
    sample: dict[str, Any],
    profile: str,
    crop_cfg: Optional[dict[str, Any]] = None,
) -> tuple[list[Image.Image], dict[str, Any]]:
    """
    Build teacher image list for privileged forward.
    text -> [full]; visual/hybrid + mode=dual -> [full, crop]; otherwise [full].
    Returns (images, debug_meta).
    """
    crop_cfg = crop_cfg or {}
    image = sample.get("image")
    if image is None:
        return [], {"crop_strategy": "none", "num_teacher_images": 0, "has_bbox": False}

    full = load_rgb(image)
    if full is None:
        return [], {"crop_strategy": "load_failed", "num_teacher_images": 0, "has_bbox": False}

    image_mode = str(crop_cfg.get("mode", "single")).strip().lower()
    if profile == "text" or image_mode in ("single", "full", "off", "none"):
        meta = {
            "crop_strategy": "single_full",
            "num_teacher_images": 1,
            "has_bbox": False,
            "bbox_norm": None,
            "image_mode": image_mode,
        }
        return [full], meta

    crop, crop_strategy, bbox_norm = heuristic_crop_from_visual_fact(full, sample, crop_cfg)
    meta = {
        "crop_strategy": crop_strategy,
        "num_teacher_images": 2,
        "has_bbox": bbox_norm is not None,
        "bbox_norm": bbox_norm,
        "full_size": full.size,
        "crop_size": crop.size,
        "image_mode": image_mode,
    }
    return [full, crop], meta