import gradio as gr
import cv2
import numpy as np
import mediapipe as mp

mp_fd = mp.solutions.face_detection.FaceDetection(model_selection=1, min_detection_confidence=0.6)

def read_rgb(image):
    # gradioはnumpy(H,W,3) RGBで来る想定
    if image is None:
        return None
    return image.copy()

def detect_face_bbox(rgb):
    h, w, _ = rgb.shape
    bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR)
    res = mp_fd.process(cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB))
    if not res.detections:
        return None
    # 最もスコアが高い顔を使う
    det = sorted(res.detections, key=lambda d: d.score[0], reverse=True)[0]
    bb = det.location_data.relative_bounding_box
    x1 = max(int(bb.xmin * w), 0)
    y1 = max(int(bb.ymin * h), 0)
    x2 = min(int((bb.xmin + bb.width) * w), w - 1)
    y2 = min(int((bb.ymin + bb.height) * h), h - 1)
    if x2 <= x1 or y2 <= y1:
        return None
    return (x1, y1, x2, y2)

def naive_paste(source_rgb, target_rgb):
    """
    まずは“動く”こと優先の簡易版：
    targetの顔bboxにsourceの顔bboxをリサイズして貼り付け、簡易フェザーでなじませる。
    （高品質にするには、この後にFaceMesh + 三角ワープ + seamlessCloneへ進む）
    """
    sb = detect_face_bbox(source_rgb)
    tb = detect_face_bbox(target_rgb)
    if sb is None:
        raise gr.Error("source画像から顔を検出できませんでした。正面顔・明るさを調整して再試行してください。")
    if tb is None:
        raise gr.Error("target画像から顔を検出できませんでした。正面顔・明るさを調整して再試行してください。")

    sx1, sy1, sx2, sy2 = sb
    tx1, ty1, tx2, ty2 = tb

    s_face = source_rgb[sy1:sy2, sx1:sx2]
    t_h, t_w = (ty2 - ty1), (tx2 - tx1)
    s_face_rs = cv2.resize(s_face, (t_w, t_h), interpolation=cv2.INTER_LINEAR)

    out = target_rgb.copy()

    # マスク（楕円 + ぼかし）で簡易ブレンド
    mask = np.zeros((t_h, t_w), dtype=np.uint8)
    cv2.ellipse(mask, (t_w // 2, t_h // 2), (int(t_w * 0.45), int(t_h * 0.50)), 0, 0, 360, 255, -1)
    mask = cv2.GaussianBlur(mask, (0, 0), sigmaX=max(t_w, t_h) * 0.02)

    roi = out[ty1:ty2, tx1:tx2]
    # alpha合成
    alpha = (mask.astype(np.float32) / 255.0)[..., None]
    blended = (alpha * s_face_rs + (1 - alpha) * roi).astype(np.uint8)
    out[ty1:ty2, tx1:tx2] = blended

    return out

def run(source_img, target_img):
    s = read_rgb(source_img)
    t = read_rgb(target_img)
    if s is None or t is None:
        raise gr.Error("source/target の両方の画像をアップロードしてください。")
    return naive_paste(s, t)

with gr.Blocks() as demo:
    gr.Markdown("# Face Swap (HF Spaces / Python)\nアップロードした2枚の画像から顔を検出し、target側にsource顔を当て込みます。")
    with gr.Row():
        src = gr.Image(label="Source（当てたい顔）", type="numpy")
        tgt = gr.Image(label="Target（当てられる側）", type="numpy")
    btn = gr.Button("Swap")
    out = gr.Image(label="Output", type="numpy")
    btn.click(run, inputs=[src, tgt], outputs=out)

demo.launch()