import os
os.system("wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth -O sam_vit_b.pth")
import gradio as gr
import numpy as np
import torch
import cv2
from PIL import Image
from segment_anything import sam_model_registry, SamPredictor

# SAM モデルの準備（適宜モデルパスを指定）
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
sam = sam_model_registry["vit_b"](checkpoint="sam_vit_b.pth").to(DEVICE)
predictor = SamPredictor(sam)

def process_image_with_mask(editor_value):
    # 'composite' から合成画像を取得
    composite = editor_value.get("composite", None)
    if composite is None:
        return None  # 合成画像がない場合は処理を終了

    # 'layers' からマスクを取得（最初のレイヤーを使用）
    layers = editor_value.get("layers", [])
    if not layers:
        return composite  # レイヤーがない場合は合成画像をそのまま返す

    mask = layers[0]

    # 画像とマスクを numpy 配列に変換
    img_np = np.array(composite.convert("RGB"))
    mask_np = np.array(mask.convert("L"))

    # SAM による予測対象のセット
    predictor.set_image(img_np)

    # マスクされたピクセル座標を取得
    yx = np.argwhere(mask_np > 127)
    if len(yx) == 0:
        return composite  # マスクがない場合は合成画像をそのまま返す

    input_points = np.flip(yx, axis=1)  # (y, x) → (x, y)
    input_labels = np.ones(len(input_points))

    masks, _, _ = predictor.predict(
        point_coords=input_points,
        point_labels=input_labels,
        multimask_output=False
    )

    mask_out = masks[0]

    # 背景を透明に
    img_rgba = composite.convert("RGBA")
    img_array = np.array(img_rgba)
    img_array[mask_out] = [0, 0, 0, 0]  # 該当領域を透明に

    return Image.fromarray(img_array)

# UI 側
with gr.Blocks() as demo:
    gr.Markdown("## 手描きマスクから物体を削除・透明化")

    editor = gr.ImageEditor(label="画像にマスクを描いてください", type="pil")

    with gr.Row():
        process_btn = gr.Button("物体を削除（透明化）")
        output = gr.Image(label="出力画像", type="pil")

    process_btn.click(
        fn=process_image_with_mask,
        inputs=[editor],
        outputs=output
    )


demo.launch()