import spaces
import gradio as gr

import torch
import numpy as np
import cv2
import mediapipe as mp

from PIL import Image
from torchvision import transforms
from transformers import (
    AutoModelForImageSegmentation,
    VitMatteForImageMatting,
    AutoImageProcessor,
)

# ============================================================
# CONFIG
# ============================================================

torch.set_float32_matmul_precision("high")

DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

# ------------------------------------------------------------
# PRIMARY MODELS
# ------------------------------------------------------------

GLOBAL_MODEL_ID = "ZhengPeng7/BiRefNet_HR"
MATTING_MODEL_ID = "ZhengPeng7/BiRefNet_HR-matting"

# ------------------------------------------------------------
# FALLBACK MODELS
# ------------------------------------------------------------

GLOBAL_FALLBACK_ID = "ZhengPeng7/BiRefNet"
MATTING_FALLBACK_ID = "ZhengPeng7/BiRefNet-matting"

# ------------------------------------------------------------
# INFERENCE
# ------------------------------------------------------------

HR_MAX_SIZE = 2048
FALLBACK_MAX_SIZE = 1536

# ------------------------------------------------------------
# HEAD / BEARD ROI
# ------------------------------------------------------------

HEAD_PAD_X = 0.85
HEAD_PAD_TOP = 0.90
HEAD_PAD_BOTTOM = 1.90

HEAD_UPSCALE = 1.50
HEAD_MAX_SIZE = 2048

# ------------------------------------------------------------
# NORMALIZATION
# ------------------------------------------------------------

MEAN = [0.485, 0.456, 0.406]
STD = [0.229, 0.224, 0.225]

# ============================================================
# MODEL LOADING
# ============================================================


def load_model(primary_id, fallback_id, name):

    print()
    print("=" * 70)
    print(f"LOADING {name}")
    print(f"Primary:  {primary_id}")
    print("=" * 70)

    try:

        model = AutoModelForImageSegmentation.from_pretrained(
            primary_id,
            trust_remote_code=True
        )

        model.to(DEVICE)
        model.eval()

        print(f"{name}: PRIMARY MODEL READY")

        return model, True

    except Exception as primary_error:

        print()
        print(f"WARNING: {name} primary model failed:")
        print(primary_error)
        print()
        print(f"Loading fallback: {fallback_id}")

        model = AutoModelForImageSegmentation.from_pretrained(
            fallback_id,
            trust_remote_code=True
        )

        model.to(DEVICE)
        model.eval()

        print(f"{name}: FALLBACK MODEL READY")

        return model, False


global_model, GLOBAL_HR = load_model(
    GLOBAL_MODEL_ID,
    GLOBAL_FALLBACK_ID,
    "GLOBAL SEGMENTATION"
)

matting_model, MATTING_HR = load_model(
    MATTING_MODEL_ID,
    MATTING_FALLBACK_ID,
    "BEARD / HAIR MATTING"
)

# ------------------------------------------------------------
# VITMATTE — trimap-guided fine matting.
#
# This is the missing piece. BiRefNet (global + matting) is a
# trimap-FREE model: it has to guess, for every pixel, whether
# it's fg/bg/edge, with no help. That's fine for the interior
# of the beard but it's exactly why thin flyaway hairs get
# killed (model guesses "background" and commits to 0) and why
# edges look blocky (no explicit unknown-region reasoning).
#
# ViTMatte instead takes an explicit trimap:
#   0   = definitely background
#   255 = definitely foreground
#   128 = "figure this out carefully"
# and only has to solve alpha inside the gray band. That's a
# fundamentally easier, more accurate problem than solving alpha
# for the whole image blind, and it's why it recovers thin hair
# structure that BiRefNet alone misses.
# ------------------------------------------------------------

VITMATTE_ID = "hustvl/vitmatte-base-composition-1k"

print()
print("=" * 70)
print("LOADING VITMATTE (trimap-guided hair/beard refinement)")
print("=" * 70)

try:
    vitmatte_processor = AutoImageProcessor.from_pretrained(VITMATTE_ID)
    vitmatte_model = VitMatteForImageMatting.from_pretrained(VITMATTE_ID)
    vitmatte_model.to(DEVICE)
    vitmatte_model.eval()
    VITMATTE_READY = True
    print("VITMATTE: READY")
except Exception as e:
    print("WARNING: ViTMatte failed to load, falling back to BiRefNet-only path.")
    print(e)
    vitmatte_processor = None
    vitmatte_model = None
    VITMATTE_READY = False

print()
print("=" * 70)
print("MODEL STATUS")
print("Global HR :", GLOBAL_HR)
print("Matte HR  :", MATTING_HR)
print("ViTMatte  :", VITMATTE_READY)
print("=" * 70)
print()

# ============================================================
# MEDIAPIPE
# ============================================================

mp_face_detection = mp.solutions.face_detection

face_detector = mp_face_detection.FaceDetection(
    model_selection=1,
    min_detection_confidence=0.45
)

# ============================================================
# TRANSFORM
# ============================================================


def make_transform(size):

    return transforms.Compose([
        transforms.Resize(
            size,
            interpolation=transforms.InterpolationMode.BILINEAR
        ),

        transforms.ToTensor(),

        transforms.Normalize(
            MEAN,
            STD
        )
    ])


# ============================================================
# MODEL SIZE
# ============================================================


def inference_size(h, w, max_size):

    largest = max(h, w)

    scale = min(
        max_size / float(largest),
        1.0
    )

    nh = int(
        round(
            (h * scale) / 32
        ) * 32
    )

    nw = int(
        round(
            (w * scale) / 32
        ) * 32
    )

    nh = max(32, nh)
    nw = max(32, nw)

    return nh, nw


# ============================================================
# MODEL RUNNER
# ============================================================


def run_model(
    model,
    image,
    max_size
):

    h = image.height
    w = image.width

    target_h, target_w = inference_size(
        h,
        w,
        max_size
    )

    transform = make_transform(
        (target_h, target_w)
    )

    tensor = transform(
        image
    ).unsqueeze(0).to(DEVICE)

    # --------------------------------------------------------
    # Do not force FP16 on CPU.
    # --------------------------------------------------------

    if DEVICE == "cuda":
        tensor = tensor.half()

    with torch.inference_mode():

        prediction = model(
            tensor
        )

        # Most BiRefNet variants return a list/tuple
        # of progressively refined predictions.
        if isinstance(prediction, (list, tuple)):

            prediction = prediction[-1]

        elif isinstance(prediction, dict):

            values = list(prediction.values())
            prediction = values[-1]

        prediction = prediction.sigmoid()

    alpha = (
        prediction[0]
        .squeeze()
        .float()
        .cpu()
        .numpy()
    )

    # Free GPU memory immediately — with 3 models running back to
    # back on ZeroGPU's shared allocation, not doing this can push
    # peak memory over budget and cause a silent worker kill with
    # no Python traceback.
    del tensor, prediction
    if DEVICE == "cuda":
        torch.cuda.empty_cache()

    # --------------------------------------------------------
    # Back to EXACT original dimensions.
    # --------------------------------------------------------

    alpha = cv2.resize(
        alpha,
        (w, h),
        interpolation=cv2.INTER_CUBIC
    )

    return np.clip(
        alpha,
        0.0,
        1.0
    ).astype(np.float32)


# ============================================================
# FACE / BEARD REGION
# ============================================================


def detect_head_region(rgb):

    h, w = rgb.shape[:2]

    try:

        result = face_detector.process(rgb)

        if not result.detections:
            return None

        detection = max(
            result.detections,
            key=lambda d: (
                max(
                    0.0,
                    d.location_data
                    .relative_bounding_box
                    .width
                )
                *
                max(
                    0.0,
                    d.location_data
                    .relative_bounding_box
                    .height
                )
            )
        )

        box = (
            detection
            .location_data
            .relative_bounding_box
        )

        x = int(box.xmin * w)
        y = int(box.ymin * h)

        bw = int(box.width * w)
        bh = int(box.height * h)

        # ----------------------------------------------------
        # BIG REGION.
        #
        # Especially important:
        # beard extends far below normal face box.
        # ----------------------------------------------------

        padx = int(bw * HEAD_PAD_X)
        padtop = int(bh * HEAD_PAD_TOP)
        padbottom = int(bh * HEAD_PAD_BOTTOM)

        x0 = max(
            0,
            x - padx
        )

        y0 = max(
            0,
            y - padtop
        )

        x1 = min(
            w,
            x + bw + padx
        )

        y1 = min(
            h,
            y + bh + padbottom
        )

        if x1 <= x0 or y1 <= y0:
            return None

        return (
            x0,
            y0,
            x1,
            y1
        )

    except Exception as e:

        print(
            "Face detection error:",
            e
        )

        return None


# ============================================================
# ALPHA HELPERS
# ============================================================


def alpha_uncertainty(alpha):

    return np.clip(
        1.0 - np.abs(alpha - 0.5) * 2.0,
        0.0,
        1.0
    ).astype(np.float32)


def alpha_boundary(alpha):

    fg = (
        alpha > 0.015
    ).astype(np.uint8)

    kernel = cv2.getStructuringElement(
        cv2.MORPH_ELLIPSE,
        (9, 9)
    )

    dilated = cv2.dilate(
        fg,
        kernel
    )

    eroded = cv2.erode(
        fg,
        kernel
    )

    band = (
        dilated.astype(np.float32)
        -
        eroded.astype(np.float32)
    )

    uncertain = alpha_uncertainty(
        alpha
    )

    band = np.maximum(
        band,
        uncertain
    )

    return np.clip(
        band,
        0.0,
        1.0
    )


# ============================================================
# VERY LIGHT GUIDED EDGE ALIGNMENT
# ============================================================


def guided_alpha(
    alpha,
    rgb,
    radius=2,
    eps=0.00035
):

    if not hasattr(cv2, "ximgproc"):
        return alpha

    if not hasattr(
        cv2.ximgproc,
        "guidedFilter"
    ):
        return alpha

    guide = (
        rgb.astype(np.float32)
        / 255.0
    )

    try:

        filtered = cv2.ximgproc.guidedFilter(
            guide=guide,
            src=alpha.astype(np.float32),
            radius=radius,
            eps=eps
        )

        # IMPORTANT:
        # Only 20% guided correction.
        #
        # Strong guided filtering can destroy
        # individual beard hairs.

        result = (
            alpha * 0.80
            +
            filtered * 0.20
        )

        return np.clip(
            result,
            0.0,
            1.0
        ).astype(np.float32)

    except Exception:

        return alpha


# ============================================================
# TRIMAP GENERATION
# ============================================================


def build_trimap(
    coarse,
    matte,
    hair_reach=25,
):
    """
    Build a trimap from the two alpha estimates we already have.

    FIX for "missing hairs": we deliberately mark a WIDE unknown
    band that extends further outward than the previous boundary
    band did (hair_reach controls this). Real beard flyaways can
    stick out well past where BiRefNet thinks the silhouette
    ends. If we only trust BiRefNet's own boundary estimate, any
    hair beyond it is unrecoverable no matter what runs next.
    Widening the "unknown" region costs nothing — ViTMatte will
    correctly resolve those pixels back to 0 if there's truly no
    hair there.
    """

    # Definite foreground: both models agree strongly.
    definite_fg = (
        (coarse > 0.92) & (matte > 0.85)
    )

    # Definite background: both models agree strongly it's empty.
    definite_bg = (
        (coarse < 0.02) & (matte < 0.05)
    )

    unknown = ~(definite_fg | definite_bg)

    # Widen the unknown band outward so distant flyaway hairs
    # that either model marks as pure background still get a
    # chance to be re-examined by ViTMatte.
    unknown_u8 = unknown.astype(np.uint8)
    kernel = cv2.getStructuringElement(
        cv2.MORPH_ELLIPSE,
        (hair_reach, hair_reach),
    )
    unknown_u8 = cv2.dilate(unknown_u8, kernel)

    trimap = np.full(coarse.shape, 128, dtype=np.uint8)
    trimap[definite_fg & (unknown_u8 == 0)] = 255
    trimap[definite_bg & (unknown_u8 == 0)] = 0
    trimap[unknown_u8 == 1] = 128

    return trimap


def run_vitmatte_refine(
    crop_rgb,
    trimap,
):
    """
    Run ViTMatte on the crop using the trimap. Returns a full
    alpha for the crop; only the trimap's 128 region actually
    matters (0/255 pass straight through), but ViTMatte needs
    the full image + trimap as input.
    """

    if not VITMATTE_READY:
        return None

    image_pil = Image.fromarray(crop_rgb)
    trimap_pil = Image.fromarray(trimap, mode="L")

    inputs = vitmatte_processor(
        images=image_pil,
        trimaps=trimap_pil,
        return_tensors="pt",
    ).to(DEVICE)

    with torch.inference_mode():
        out = vitmatte_model(**inputs)

    alpha = out.alphas[0, 0].float().cpu().numpy()

    del inputs, out
    if DEVICE == "cuda":
        torch.cuda.empty_cache()

    # ViTMatte processor pads to a stride-32 canvas internally;
    # resize back to the exact crop size.
    alpha = cv2.resize(
        alpha,
        (crop_rgb.shape[1], crop_rgb.shape[0]),
        interpolation=cv2.INTER_LINEAR,
    )

    return np.clip(alpha, 0.0, 1.0).astype(np.float32)


# ============================================================
# HEAD MATTE
# ============================================================


def run_head_matte(
    rgb,
    global_alpha
):

    box = detect_head_region(
        rgb
    )

    if box is None:

        print(
            "No face detected -> no local beard crop."
        )

        return global_alpha

    x0, y0, x1, y1 = box

    crop = rgb[
        y0:y1,
        x0:x1
    ]

    crop_h, crop_w = crop.shape[:2]

    if crop.size == 0:
        return global_alpha

    print(
        f"HEAD/BEARD ROI: {crop_w}x{crop_h}"
    )

    # --------------------------------------------------------
    # High-resolution local image.
    #
    # This is ONLY used to estimate alpha.
    # It is never used as final RGB.
    # --------------------------------------------------------

    scale = HEAD_UPSCALE

    high_w = min(
        HEAD_MAX_SIZE,
        max(
            crop_w,
            int(crop_w * scale)
        )
    )

    high_h = min(
        HEAD_MAX_SIZE,
        max(
            crop_h,
            int(crop_h * scale)
        )
    )

    highres = cv2.resize(
        crop,
        (high_w, high_h),
        interpolation=cv2.INTER_CUBIC
    )

    matte_input = Image.fromarray(
        highres
    )

    print(
        f"BEARD MATTE INPUT: {high_w}x{high_h}"
    )

    matte_high = run_model(
        matting_model,
        matte_input,
        (
            HR_MAX_SIZE
            if MATTING_HR
            else FALLBACK_MAX_SIZE
        )
    )

    # --------------------------------------------------------
    # Return matte to original crop dimensions.
    # --------------------------------------------------------

    matte = cv2.resize(
        matte_high,
        (crop_w, crop_h),
        interpolation=cv2.INTER_LANCZOS4
    )

    coarse = global_alpha[
        y0:y1,
        x0:x1
    ]

    # ========================================================
    # CRITICAL PART
    # ========================================================
    #
    # - global controls confident interior/background
    # - local matte controls boundary + fine hairs
    # ========================================================

    uncertain = alpha_uncertainty(
        coarse
    )

    boundary = alpha_boundary(
        coarse
    )

    local_weight = (
        uncertain * 0.70
        +
        boundary * 0.55
    )

    local_weight = np.clip(
        local_weight,
        0.0,
        1.0
    )

    # --------------------------------------------------------
    # If global is strongly confident foreground,
    # preserve global interior.
    # --------------------------------------------------------

    local_weight[
        coarse > 0.96
    ] = np.minimum(
        local_weight[
            coarse > 0.96
        ],
        0.12
    )

    # --------------------------------------------------------
    # If global is strongly background, still allow the
    # matting model to recover tiny beard/hair strands.
    # --------------------------------------------------------

    tiny_hair_zone = (
        (coarse > 0.003)
        &
        (coarse < 0.35)
    )

    local_weight[
        tiny_hair_zone
    ] = np.maximum(
        local_weight[
            tiny_hair_zone
        ],
        0.72
    )

    # --------------------------------------------------------
    # FIX: this blur was previously part of a 3-pass blur
    # stack (here + preserve_fine_alpha + micro_sharpen_alpha)
    # that softened 1-3px hair strands. Reduced further since
    # the other two passes downstream are now toned down too.
    # --------------------------------------------------------

    local_weight = cv2.GaussianBlur(
        local_weight,
        (0, 0),
        0.20
    )

    # --------------------------------------------------------
    # Matte is authoritative around beard/hair.
    # --------------------------------------------------------

    refined = (
        matte * local_weight
        +
        coarse * (1.0 - local_weight)
    )

    # --------------------------------------------------------
    # Restore very confident global interior.
    # --------------------------------------------------------

    refined[
        coarse > 0.985
    ] = coarse[
        coarse > 0.985
    ]

    # --------------------------------------------------------
    # Edge alignment is deliberately tiny.
    # --------------------------------------------------------

    refined = guided_alpha(
        refined,
        crop,
        radius=2,
        eps=0.0003
    )

    # ========================================================
    # VITMATTE REFINEMENT — this is the new, real fix.
    #
    # BiRefNet (global + local matte) has now given us its best
    # guess ("refined"). We build a trimap from that guess, then
    # let ViTMatte solve the actual unknown/hair band explicitly
    # instead of us continuing to hand-tune blend weights.
    #
    # This directly targets all three problems reported:
    #   - missing hairs: wide unknown band + trimap-guided model
    #     built specifically to recover fine structure in gray
    #     regions, rather than committing to 0 early.
    #   - blocky edges: ViTMatte outputs continuous soft alpha
    #     for the whole unknown band, not a boolean/dilated mask
    #     blended in with hard weights.
    #   - color tint: a cleaner alpha here means less "wrong"
    #     alpha for decontamination to try to compensate for
    #     downstream — most of the residual tint was fallout
    #     from imperfect alpha, not just the color-correction
    #     math.
    # ========================================================

    if VITMATTE_READY:

        trimap = build_trimap(
            coarse,
            matte,
            hair_reach=25,
        )

        vit_alpha = run_vitmatte_refine(
            crop,
            trimap,
        )

        if vit_alpha is not None:

            unknown_band = (trimap == 128).astype(np.float32)

            # Soften the seam between "definite" regions and the
            # ViTMatte-solved region so there's no visible ring.
            unknown_band = cv2.GaussianBlur(
                unknown_band,
                (0, 0),
                1.5,
            )

            refined = (
                vit_alpha * unknown_band
                +
                refined * (1.0 - unknown_band)
            )

    output = global_alpha.copy()

    output[
        y0:y1,
        x0:x1
    ] = refined

    return np.clip(
        output,
        0.0,
        1.0
    ).astype(np.float32)


# ============================================================
# FINE DETAIL RECOVERY
# ============================================================


def preserve_fine_alpha(alpha):

    # --------------------------------------------------------
    # FIX: sigma reduced from 0.55 -> 0.35 and strength halved.
    # This was the second of three stacked blurs that were
    # collectively softening thin hair strands. It still
    # recovers real high-frequency detail, just less
    # aggressively, so it doesn't fight the matte's own edges.
    # --------------------------------------------------------

    smooth = cv2.GaussianBlur(
        alpha,
        (0, 0),
        0.35
    )

    detail = (
        alpha
        -
        smooth
    )

    uncertainty = alpha_uncertainty(
        alpha
    )

    strength = (
        0.05
        +
        0.05 * uncertainty
    )

    result = (
        alpha
        +
        detail * strength
    )

    return np.clip(
        result,
        0.0,
        1.0
    ).astype(np.float32)


# ============================================================
# MICRO EDGE SHARPEN
# ============================================================


def micro_sharpen_alpha(alpha):

    # --------------------------------------------------------
    # FIX: this was the third blur in the stack. Sigma reduced
    # and strength reduced — this stage was contributing very
    # little sharpening but a full extra blur pass of softening
    # to every hair edge.
    # --------------------------------------------------------

    blurred = cv2.GaussianBlur(
        alpha,
        (0, 0),
        0.25
    )

    detail = (
        alpha
        -
        blurred
    )

    uncertainty = alpha_uncertainty(
        alpha
    )

    result = (
        alpha
        +
        detail
        *
        (
            0.04
            *
            uncertainty
        )
    )

    return np.clip(
        result,
        0.0,
        1.0
    ).astype(np.float32)


# ============================================================
# ALPHA CLEANUP
# ============================================================


def cleanup_alpha(alpha):

    result = np.clip(
        alpha,
        0.0,
        1.0
    ).astype(np.float32)

    # Only remove almost-zero numerical noise.
    result[
        result < 0.002
    ] = 0.0

    result[
        result > 0.999
    ] = 1.0

    return result


# ============================================================
# BACKGROUND ESTIMATION
# ============================================================


def estimate_background_color(
    rgb,
    alpha
):

    h, w = alpha.shape

    border = max(
        4,
        int(min(h, w) * 0.025)
    )

    border_mask = np.zeros(
        (h, w),
        dtype=bool
    )

    border_mask[
        :border,
        :
    ] = True

    border_mask[
        -border:,
        :
    ] = True

    border_mask[
        :,
        :border
    ] = True

    border_mask[
        :,
        -border:
    ] = True

    mask = (
        border_mask
        &
        (alpha < 0.02)
    )

    pixels = rgb[
        mask
    ]

    if len(pixels) < 100:

        pixels = rgb[
            alpha < 0.02
        ]

    if len(pixels) < 100:

        return np.array(
            [255.0, 255.0, 255.0],
            dtype=np.float32
        )

    return np.median(
        pixels.astype(np.float32),
        axis=0
    )


# ============================================================
# LOCAL BACKGROUND COLOR
# ============================================================


def local_background_color(
    rgb,
    alpha
):
    """
    Estimate background around each pixel.

    This is better than using one global color because
    real photos may have gradients/shadows.
    """

    background_mask = (
        alpha < 0.025
    ).astype(np.float32)

    weight = cv2.GaussianBlur(
        background_mask,
        (0, 0),
        7.0
    )

    weight = np.maximum(
        weight,
        1e-5
    )

    rgb_float = rgb.astype(
        np.float32
    )

    channels = []

    for c in range(3):

        channel = (
            rgb_float[..., c]
            *
            background_mask
        )

        smooth_channel = cv2.GaussianBlur(
            channel,
            (0, 0),
            7.0
        )

        local = (
            smooth_channel
            /
            weight
        )

        channels.append(
            local
        )

    bg = np.stack(
        channels,
        axis=2
    )

    global_bg = estimate_background_color(
        rgb,
        alpha
    )

    # If local estimate has no real background nearby,
    # fall back toward global estimate.
    bg = np.where(
        np.isfinite(bg),
        bg,
        global_bg.reshape(1, 1, 3)
    )

    bg = np.clip(
        bg,
        0,
        255
    )

    return bg.astype(np.float32)


# ============================================================
# BACKGROUND COLOR DECONTAMINATION  (FIXED)
# ============================================================


def decontaminate_edges(
    rgb,
    alpha
):
    """
    Remove background color contamination from semi-transparent
    edge pixels.

    Formula:
        C = alpha * F + (1-alpha) * B
    therefore:
        F = (C - (1-alpha)B) / alpha

    ------------------------------------------------------------
    THE BUG THAT WAS HERE:

    The original code applied strong correction to mid-alpha
    pixels but then explicitly SUPPRESSED correction on the
    lowest-alpha pixels:

        strength[a < 0.10] *= 0.35

    and used a high safe_alpha floor (0.16), which under-divides
    the recovered color for anything below that.

    But thin beard/hair flyaways sit exactly in that a < 0.10
    range. That's precisely where the blue/cyan background tint
    was showing up — the math correctly detected the
    contamination, then was told to only partially fix it.

    THE FIX:
    - Lower the safe_alpha floor so correction isn't artificially
      weakened for real hair pixels.
    - Do NOT crush "strength" for low alpha. Instead correct at
      (near) full strength, then denoise the recovered color in
      the low-alpha band with a small median filter. This keeps
      the instability under control without leaving background
      color baked into the hair.
    ------------------------------------------------------------
    """

    image = rgb.astype(
        np.float32
    )

    a = np.clip(
        alpha.astype(np.float32),
        0.0,
        1.0
    )

    background = local_background_color(
        rgb,
        alpha
    )

    a3 = a[..., None]

    # --------------------------------------------------------
    # FIX for blockiness: replaced the boolean edge mask with a
    # continuous weight. A hard (a > 0.03) & (a < 0.94) mask means
    # a pixel at a=0.029 gets ZERO correction while its neighbor
    # at a=0.031 gets full correction — that step is visible as
    # blocky/banded edges once composited on a new background.
    # A smooth ramp removes the seam entirely.
    # --------------------------------------------------------

    edge_weight = np.clip(
        np.minimum(
            (a - 0.02) / 0.06,      # ramp up from a=0.02 to 0.08
            (0.96 - a) / 0.10,      # ramp down from a=0.86 to 0.96
        ),
        0.0,
        1.0,
    )

    edge = edge_weight > 0.0

    if not np.any(edge):

        return rgb

    # --------------------------------------------------------
    # FIX: lower floor (was 0.16). Let the correction actually
    # happen for thin hair instead of pretending alpha is higher
    # than it is.
    # --------------------------------------------------------

    safe_alpha = np.maximum(
        a3,
        0.06
    )

    recovered = (
        image
        -
        (1.0 - a3)
        *
        background
    ) / safe_alpha

    recovered = np.clip(
        recovered,
        0.0,
        255.0
    )

    # --------------------------------------------------------
    # FIX: denoise the recovered color ONLY in the low-alpha
    # band. Division by a small number amplifies per-pixel
    # sensor noise; a 3x3 median kills that noise without
    # re-introducing background color (median doesn't blend
    # in neighboring background pixels' color the way a
    # Gaussian blur would).
    # --------------------------------------------------------

    low_alpha_mask = (a < 0.20) & edge

    if np.any(low_alpha_mask):

        recovered_u8 = np.clip(
            recovered, 0, 255
        ).astype(np.uint8)

        recovered_denoised = cv2.medianBlur(
            recovered_u8,
            3
        ).astype(np.float32)

        recovered = np.where(
            low_alpha_mask[..., None],
            recovered_denoised,
            recovered
        )

    # --------------------------------------------------------
    # Determine how strongly to correct.
    #
    # FIX: strength now RISES toward low alpha instead of being
    # cut down. Low-alpha pixels are the ones most likely to be
    # majority background, so they need the most correction, not
    # the least.
    # --------------------------------------------------------

    low_alpha_strength = np.clip(
        (0.94 - a) / 0.70,
        0.0,
        1.0
    )

    distance = np.linalg.norm(
        image - background,
        axis=2
    )

    similarity = np.exp(
        -(
            distance / 85.0
        ) ** 2
    )

    strength = (
        low_alpha_strength
        *
        (
            0.35
            +
            0.55 * similarity
        )
    )

    # Allow near-full correction now (was capped at 0.58).
    strength = np.clip(
        strength,
        0.0,
        0.95
    )

    # Strong foreground should barely change.
    strength[
        a > 0.82
    ] *= 0.20

    # --------------------------------------------------------
    # FIX: removed the old line that crushed strength for
    # a < 0.10 (`strength[a < 0.10] *= 0.35`). That line was
    # the main cause of the residual blue fringe. The median
    # denoise above handles the instability instead.
    # --------------------------------------------------------

    # FIX: multiply by the continuous edge_weight instead of a
    # boolean np.where — this removes the hard seam at the band
    # boundary entirely, since strength now fades to exactly 0
    # smoothly rather than snapping off.
    strength = strength * edge_weight

    strength3 = strength[..., None]

    output = (
        image * (1.0 - strength3)
        +
        recovered * strength3
    )

    return np.clip(
        output,
        0,
        255
    ).astype(np.uint8)


# ============================================================
# ADDITIONAL BLUE-SPILL SAFETY  (FIXED)
# ============================================================


def remove_background_tint(
    rgb,
    alpha
):
    """
    Conservative second pass for strongly colored backgrounds
    (blue/green screen etc).

    FIX: previously this ALSO crushed correction for the
    lowest-alpha pixels (`strength[a < 0.08] *= 0.25`), stacking
    on top of the same bug in decontaminate_edges. Removed that
    line so this pass actually helps rather than mostly no-oping
    on exactly the pixels that show the tint.
    """

    image = rgb.astype(
        np.float32
    )

    a = alpha.astype(
        np.float32
    )

    bg = estimate_background_color(
        rgb,
        alpha
    )

    bg_mean = float(
        np.mean(bg)
    )

    chroma = float(
        np.max(bg) - np.min(bg)
    )

    # If background isn't strongly colored, use only tiny
    # correction. If it IS strongly colored (like your blue
    # studio background), allow more correction than before.
    if chroma < 18.0:
        max_strength = 0.10
    else:
        max_strength = 0.32

    bg3 = bg.reshape(
        1,
        1,
        3
    )

    distance = np.linalg.norm(
        image - bg3,
        axis=2
    )

    similarity = np.exp(
        -(
            distance / 95.0
        ) ** 2
    )

    edge_weight = np.clip(
        (0.85 - a) / 0.70,
        0.0,
        1.0
    )

    strength = (
        similarity
        *
        edge_weight
        *
        max_strength
    )

    # Do not touch solid beard/body.
    strength[
        a > 0.80
    ] = 0.0

    # --------------------------------------------------------
    # FIX: removed `strength[a < 0.08] *= 0.25`. That was
    # suppressing correction on exactly the thinnest hairs,
    # doubling up with the same bug in decontaminate_edges.
    # --------------------------------------------------------

    strength3 = strength[..., None]

    corrected = (
        image
        +
        (
            image - bg3
        )
        *
        strength3
    )

    corrected = np.clip(
        corrected,
        0,
        255
    )

    return corrected.astype(
        np.uint8
    )


# ============================================================
# FINAL EDGE PROTECTION
# ============================================================


def protect_fine_hair(
    original_rgb,
    cleaned_rgb,
    alpha
):
    """
    Prevent the color cleanup stage from changing solid
    foreground pixels. Hair/beard edge pixels remain corrected,
    while solid subject pixels stay essentially original.
    """

    a = alpha.astype(
        np.float32
    )

    original = original_rgb.astype(
        np.float32
    )

    cleaned = cleaned_rgb.astype(
        np.float32
    )

    strength = np.clip(
        (0.92 - a) / 0.75,
        0.0,
        1.0
    )

    strength[
        a > 0.86
    ] = 0.0

    strength *= 0.85

    strength3 = strength[..., None]

    result = (
        original * (1.0 - strength3)
        +
        cleaned * strength3
    )

    return np.clip(
        result,
        0,
        255
    ).astype(np.uint8)


# ============================================================
# FINAL ALPHA
# ============================================================


def finalize_alpha(alpha):

    alpha = np.clip(
        alpha,
        0.0,
        1.0
    ).astype(np.float32)

    alpha[
        alpha < 0.0015
    ] = 0.0

    alpha[
        alpha > 0.9995
    ] = 1.0

    return alpha


# ============================================================
# MAIN PIPELINE
# ============================================================


@spaces.GPU(duration=280)
def remove_background(
    image
):

    if image is None:

        raise gr.Error(
            "Please upload a photo first."
        )

    try:

        # ====================================================
        # ORIGINAL
        # ====================================================

        original = image.convert(
            "RGB"
        )

        rgb = np.array(
            original,
            dtype=np.uint8
        )

        h, w = rgb.shape[:2]

        print()
        print("=" * 70)
        print(f"INPUT: {w}x{h}")
        print("=" * 70)

        # ====================================================
        # STAGE 1 — GLOBAL SEGMENTATION
        # ====================================================

        print("STAGE 1: Global subject segmentation")

        global_alpha = run_model(
            global_model,
            original,
            (
                HR_MAX_SIZE
                if GLOBAL_HR
                else FALLBACK_MAX_SIZE
            )
        )

        global_alpha = guided_alpha(
            global_alpha,
            rgb,
            radius=2,
            eps=0.0005
        )

        # ====================================================
        # STAGE 2 — BEARD / HAIR MATTE
        # ====================================================

        print("STAGE 2: High-resolution beard/hair matting")

        alpha = run_head_matte(
            rgb,
            global_alpha
        )

        # ====================================================
        # STAGE 3 — FINE DETAIL
        # ====================================================

        print("STAGE 3: Fine hair/beard detail preservation")

        alpha = preserve_fine_alpha(
            alpha
        )

        # ====================================================
        # STAGE 4 — MICRO SHARPEN
        # ====================================================

        print("STAGE 4: Micro alpha sharpening")

        alpha = micro_sharpen_alpha(
            alpha
        )

        # ====================================================
        # STAGE 5 — FINAL ALPHA
        # ====================================================

        alpha = cleanup_alpha(
            alpha
        )

        alpha = finalize_alpha(
            alpha
        )

        # ====================================================
        # STAGE 6 — BACKGROUND COLOR REMOVAL (fixed)
        # ====================================================

        print("STAGE 5: Background color decontamination")

        clean_rgb = decontaminate_edges(
            rgb,
            alpha
        )

        # ====================================================
        # STAGE 7 — COLOR-SPILL SAFETY (fixed)
        # ====================================================

        print("STAGE 6: Fine background-tint cleanup")

        clean_rgb = remove_background_tint(
            clean_rgb,
            alpha
        )

        # ====================================================
        # STAGE 8 — PROTECT SUBJECT
        # ====================================================

        clean_rgb = protect_fine_hair(
            rgb,
            clean_rgb,
            alpha
        )

        # ====================================================
        # FINAL PNG
        # ====================================================

        print("Creating final transparent PNG...")

        result = Image.fromarray(
            clean_rgb,
            mode="RGB"
        )

        alpha_image = Image.fromarray(
            np.clip(
                alpha * 255.0,
                0,
                255
            ).astype(np.uint8),
            mode="L"
        )

        result.putalpha(
            alpha_image
        )

        print("DONE.")
        print("=" * 70)

        return result

    except Exception as e:

        import traceback

        traceback.print_exc()

        raise gr.Error(
            f"Background removal failed: {e}"
        )


# ============================================================
# UI
# ============================================================

css = """
#header {
    text-align: center;
    padding: 26px 0 10px;
}

#header h1 {
    font-size: 32px;
    font-weight: 700;

    background: linear-gradient(
        135deg,
        #6366f1,
        #06b6d4
    );

    -webkit-background-clip: text;
    -webkit-text-fill-color: transparent;

    margin-bottom: 5px;
}

#header p {
    color: #888;
    font-size: 14px;
}

#run-btn {
    background: linear-gradient(
        135deg,
        #6366f1,
        #06b6d4
    ) !important;

    color: white !important;
    font-weight: 600 !important;
    border: none !important;
}
"""

with gr.Blocks(
    title="Peace Network BG Remover"
) as demo:

    gr.HTML(
        """
        <div id="header">
            <h1>
                Peace Network BG Remover
            </h1>

            <p>
                High-resolution background removal
                with fine hair & beard preservation.
            </p>
        </div>
        """
    )

    with gr.Row():

        inp = gr.Image(
            type="pil",
            label="Upload photo"
        )

        out = gr.Image(
            type="pil",
            label="Transparent PNG",
            format="png"
        )

    btn = gr.Button(
        "Remove Background",
        variant="primary",
        elem_id="run-btn"
    )

    btn.click(
        remove_background,
        inputs=inp,
        outputs=out
    )

# ============================================================
# LAUNCH
# ============================================================

demo.queue().launch(
    css=css
)