engine/mapper.py

import os
from PIL import Image, ImageStat

Image.MAX_IMAGE_PIXELS = 200_000_000

from engine.hoodie_rules import AOP_HOODIE_MAPPING_RULES
from engine.hf_processors import run_optional_ai_processors

def norm(name):
    return name.lower().replace(" ", "_")

def pct_box(w, h, x1, y1, x2, y2):
    return (int(w * x1), int(h * y1), int(w * x2), int(h * y2))

def region_score(img):
    stat = ImageStat.Stat(img.convert("RGB"))
    contrast = sum(stat.stddev) / 3
    saturation_hint = max(stat.mean) - min(stat.mean)
    brightness = sum(stat.mean) / 3
    return contrast + saturation_hint + (255 - abs(140 - brightness)) * 0.12

def analyze_regions(image, prompt, ai_results=None):
    ai_results = ai_results or {}
    w, h = image.size
    p = (prompt or "").lower()

    base = [
        ("PRIMARY_FOCAL_REGION", pct_box(w, h, 0.18, 0.10, 0.82, 0.82)),
        ("SECONDARY_FOCAL_REGION", pct_box(w, h, 0.04, 0.00, 0.75, 0.50)),
        ("ENVIRONMENT_REGION", pct_box(w, h, 0.00, 0.00, 1.00, 1.00)),
        ("LOWER_ANCHOR_REGION", pct_box(w, h, 0.05, 0.58, 0.95, 1.00)),
        ("LEFT_EDGE_FLOW_REGION", pct_box(w, h, 0.00, 0.04, 0.38, 0.96)),
        ("RIGHT_EDGE_FLOW_REGION", pct_box(w, h, 0.62, 0.04, 1.00, 0.96)),
    ]

    regions = []

    for region_type, box in base:
        crop = image.crop(box)
        score = region_score(crop)

        if region_type == "PRIMARY_FOCAL_REGION":
            if any(k in p for k in ["face", "head", "person", "character", "dj", "main", "center", "figure", "subject"]):
                score += 45

        if region_type == "SECONDARY_FOCAL_REGION":
            if any(k in p for k in ["eye", "monster", "upper", "teeth", "dramatic", "background"]):
                score += 35

        if region_type == "LOWER_ANCHOR_REGION":
            if any(k in p for k in ["pocket", "deck", "table", "lower", "base", "ground", "smoke"]):
                score += 30

        if region_type == "ENVIRONMENT_REGION":
            if any(k in p for k in ["pattern", "repeat", "tile", "smoke", "flame", "abstract", "texture"]):
                score += 30

        if region_type == "LEFT_EDGE_FLOW_REGION":
            if any(k in p for k in ["left sleeve", "left edge", "flame", "smoke"]):
                score += 20

        if region_type == "RIGHT_EDGE_FLOW_REGION":
            if any(k in p for k in ["right sleeve", "right edge", "purple", "smoke"]):
                score += 20

        regions.append({
            "type": region_type,
            "box": box,
            "score": round(score, 2),
            "source": "local_region_analysis",
        })

    for category, result in ai_results.items():
        if not result.get("ok"):
            continue

        for det in result.get("detections", []):
            det_type = str(det.get("type", category)).lower()
            box = det.get("box")

            if not box or len(box) != 4:
                continue

            region_type = None

            if det_type in ["face", "person", "object", "subject"]:
                region_type = "PRIMARY_FOCAL_REGION"
            elif det_type in ["text", "ocr"]:
                region_type = "TEXT_REGION"
            elif det_type == "logo":
                region_type = "LOGO_REGION"

            if region_type:
                regions.append({
                    "type": region_type,
                    "box": tuple(map(int, box)),
                    "score": float(det.get("score", 85)),
                    "source": "hf_optional_ai",
                })

    return sorted(regions, key=lambda r: r["score"], reverse=True)

def cover_resize_no_padding(img, target_w, target_h):
    src_w, src_h = img.size
    scale = max(target_w / src_w, target_h / src_h)
    new_w, new_h = int(src_w * scale), int(src_h * scale)

    resized = img.resize((new_w, new_h), Image.LANCZOS)

    left = max(0, (new_w - target_w) // 2)
    top = max(0, (new_h - target_h) // 2)

    return resized.crop((left, top, left + target_w, top + target_h))

def contain_resize_with_transparency(img, target_w, target_h):
    src_w, src_h = img.size
    scale = min(target_w / src_w, target_h / src_h)
    new_w, new_h = int(src_w * scale), int(src_h * scale)

    resized = img.resize((new_w, new_h), Image.LANCZOS)

    canvas = Image.new("RGBA", (target_w, target_h), (0, 0, 0, 0))
    canvas.paste(resized, ((target_w - new_w) // 2, (target_h - new_h) // 2))

    return canvas

def tile_texture(img, target_w, target_h):
    tile_size = max(300, min(900, target_w, target_h))
    tile = cover_resize_no_padding(img, tile_size, tile_size)

    out = Image.new("RGBA", (target_w, target_h))

    for x in range(0, target_w, tile.width):
        for y in range(0, target_h, tile.height):
            out.paste(tile, (x, y))

    return out.crop((0, 0, target_w, target_h))

def get_mask_path(template_id, piece_name):
    candidates = [
        f"template_assets/{template_id}/masks/{piece_name}.png",
        f"template_assets/{template_id}/masks/{norm(piece_name)}.png",
        f"template_assets/{template_id}/{piece_name}.png",
        f"template_assets/{template_id}/{norm(piece_name)}.png",
    ]

    for path in candidates:
        if os.path.exists(path):
            return path

    return None

def apply_piece_mask(mapped_img, template_id, piece_name):
    mask_path = get_mask_path(template_id, piece_name)

    if not mask_path:
        return mapped_img, None

    mask_img = Image.open(mask_path).convert("RGBA").resize(mapped_img.size, Image.LANCZOS)
    alpha = mask_img.getchannel("A")

    out = mapped_img.convert("RGBA")
    out.putalpha(alpha)

    return out, mask_path

def choose_region(piece_name, regions):
    rule = AOP_HOODIE_MAPPING_RULES.get(norm(piece_name))

    if not rule:
        return regions[0], "safe_cover"

    for preferred in rule["preferred"]:
        for region in regions:
            if region["type"] == preferred and region["type"] not in rule["forbidden"]:
                return region, rule["strategy"]

    for region in regions:
        if region["type"] not in rule["forbidden"]:
            return region, rule["strategy"]

    env = next((r for r in regions if r["type"] == "ENVIRONMENT_REGION"), regions[0])
    return env, "fallback_texture"

def apply_strategy(image, region, strategy, target_w, target_h):
    crop = image.crop(region["box"]).convert("RGBA")

    if strategy == "texture_band":
        return tile_texture(crop, target_w, target_h)

    if strategy == "vertical_texture":
        return cover_resize_no_padding(crop, target_w, target_h)

    if strategy == "dark_texture":
        return cover_resize_no_padding(crop, target_w, target_h)

    if strategy == "fallback_texture":
        return cover_resize_no_padding(crop, target_w, target_h)

    if strategy == "centered_focal":
        return cover_resize_no_padding(crop, target_w, target_h)

    if strategy == "wide_dramatic":
        return cover_resize_no_padding(crop, target_w, target_h)

    if strategy == "lower_anchor":
        return cover_resize_no_padding(crop, target_w, target_h)

    return cover_resize_no_padding(crop, target_w, target_h)

def map_all(image_path, template_id, pieces, prompt, use_optional_ai=False):
    image = Image.open(image_path).convert("RGBA")

    ai_results = run_optional_ai_processors(image_path, prompt) if use_optional_ai else {}
    regions = analyze_regions(image, prompt, ai_results)

    outputs = {}
    decisions = []
    warnings = []
    used_boxes = {}
    mask_usage = {}

    for piece_name, size in pieces.items():
        target_w, target_h = size

        region, strategy = choose_region(piece_name, regions)
        mapped = apply_strategy(image, region, strategy, target_w, target_h)
        mapped, mask_path = apply_piece_mask(mapped, template_id, piece_name)

        outputs[piece_name] = mapped

        decision = {
            "piece": piece_name,
            "region": region["type"],
            "strategy": strategy,
            "box": region["box"],
            "score": region["score"],
            "output_size": [target_w, target_h],
            "mask_applied": bool(mask_path),
            "mask_path": mask_path,
        }

        decisions.append(decision)

        key = str(region["box"])
        used_boxes[key] = used_boxes.get(key, 0) + 1
        mask_usage[piece_name] = bool(mask_path)

        risky_small = norm(piece_name) in [
            "left_sleeve",
            "right_sleeve",
            "collar",
            "left_cuff",
            "right_cuff",
            "waistband",
            "kangaroo_pocket",
        ]

        if risky_small and region["type"] in [
            "TEXT_REGION",
            "LOGO_REGION",
            "PRIMARY_FOCAL_REGION",
        ]:
            warnings.append(
                f"{piece_name}: readable text, logo, or primary focal content may be on a risky small panel."
            )

        if not mask_path:
            warnings.append(
                f"{piece_name}: no piece mask found, output is rectangular. Add template_assets/{template_id}/masks/{piece_name}.png for true shaped AOP output."
            )

    front_decisions = [
        d for d in decisions if norm(d["piece"]) in ["front", "front_body"]
    ]

    if front_decisions and front_decisions[0]["region"] != "PRIMARY_FOCAL_REGION":
        warnings.append("Main focal point was not placed on the front panel.")

    if used_boxes and max(used_boxes.values()) > 3:
        warnings.append("Too many panels are using the same crop. Design may look repetitive.")

    if not any(mask_usage.values()):
        warnings.append(
            "No template masks were applied. Current outputs are rectangular crops, not true garment-shaped production pieces."
        )

    return outputs, {
        "prompt": prompt,
        "template_id": template_id,
        "plain_engine_logic": [
            "Big panels get meaning.",
            "Small panels get texture.",
            "Risky seams get atmosphere.",
            "Pockets get lower anchors.",
            "Cuffs get color bands.",
            "Hoods get continuation, not the whole story.",
            "True AOP piece shape requires alpha masks from template PNGs.",
        ],
        "regions": regions,
        "decisions": decisions,
        "warnings": warnings,
        "optional_ai_processors": ai_results,
    }