import os
import gradio as gr
import cv2
import numpy as np
from PIL import Image

os.makedirs("./SAM2-Video-Predictor/checkpoints/", exist_ok=True)
os.makedirs("./model/", exist_ok=True)

from huggingface_hub import snapshot_download

def download_sam2():
    snapshot_download(repo_id="facebook/sam2-hiera-large", local_dir="./SAM2-Video-Predictor/checkpoints/")
    print("Download sam2 completed")

def download_remover():
    snapshot_download(repo_id="zibojia/minimax-remover", local_dir="./model/")
    print("Download minimax remover completed")

download_sam2()
download_remover()

import torch
import argparse
import random

import torch.nn.functional as F
import time
import random
from omegaconf import OmegaConf
from einops import rearrange
from diffusers.models import AutoencoderKLWan
import scipy
from transformer_minimax_remover import Transformer3DModel
from einops import rearrange
from diffusers.schedulers import UniPCMultistepScheduler
from pipeline_minimax_remover import Minimax_Remover_Pipeline

from diffusers.utils import export_to_video
from decord import VideoReader, cpu
from moviepy.editor import ImageSequenceClip

from sam2 import load_model

from sam2.build_sam import build_sam2, build_sam2_video_predictor
from sam2.sam2_image_predictor import SAM2ImagePredictor

import spaces

COLOR_PALETTE = [
    (255, 0, 0), (0, 255, 0), (0, 0, 255), (255, 255, 0), (255, 0, 255),
    (0, 255, 255), (255, 128, 0), (128, 0, 255), (0, 128, 255), (128, 255, 0)
]

random_seed = 42
video_length = 9000
W = 1024
H = W
device = "cpu"

def get_pipe_image_and_video_predictor():
    vae = AutoencoderKLWan.from_pretrained("./model/vae", torch_dtype=torch.float16)
    transformer = Transformer3DModel.from_pretrained("./model/transformer", torch_dtype=torch.float16)
    scheduler = UniPCMultistepScheduler.from_pretrained("./model/scheduler")

    pipe = Minimax_Remover_Pipeline(transformer=transformer, vae=vae, scheduler=scheduler)

    sam2_checkpoint = "./SAM2-Video-Predictor/checkpoints/sam2_hiera_large.pt"
    config = "sam2_hiera_l.yaml"

    video_predictor = build_sam2_video_predictor(config, sam2_checkpoint, device=device)
    model = build_sam2(config, sam2_checkpoint, device=device)
    model.image_size = 1024
    image_predictor = SAM2ImagePredictor(sam_model=model)

    return pipe, image_predictor, video_predictor

def get_video_info(video_path, video_state):
    video_state["input_points"] = []
    video_state["scaled_points"] = []
    video_state["input_labels"] = []
    video_state["frame_idx"] = 0
    vr = VideoReader(video_path, ctx=cpu(0))
    first_frame = vr[0].asnumpy()
    del vr

    if first_frame.shape[0] > first_frame.shape[1]:
        W_ = W
        H_ = int(W_ * first_frame.shape[0] / first_frame.shape[1])
    else:
        H_ = H
        W_ = int(H_ * first_frame.shape[1] / first_frame.shape[0])

    first_frame = cv2.resize(first_frame, (W_, H_))
    video_state["origin_images"] = np.expand_dims(first_frame, axis=0)
    video_state["inference_state"] = None
    video_state["video_path"] = video_path
    video_state["masks"] = None
    video_state["painted_images"] = None
    image = Image.fromarray(first_frame)
    return image

def segment_frame(evt: gr.SelectData, label, video_state):
    if video_state["origin_images"] is None:
        return None
    x, y = evt.index
    new_point = [x, y]
    label_value = 1 if label == "Positive" else 0

    video_state["input_points"].append(new_point)
    video_state["input_labels"].append(label_value)
    height, width = video_state["origin_images"][0].shape[0:2]
    scaled_points = []
    for pt in video_state["input_points"]:
        sx = pt[0] / width
        sy = pt[1] / height
        scaled_points.append([sx, sy])

    video_state["scaled_points"] = scaled_points

    image_predictor.set_image(video_state["origin_images"][0])
    mask, _, _ = image_predictor.predict(
        point_coords=video_state["scaled_points"],
        point_labels=video_state["input_labels"],
        multimask_output=False,
        normalize_coords=False,
    )

    mask = np.squeeze(mask)
    mask = cv2.resize(mask, (width, height))
    mask = mask[:,:,None]

    color = np.array(COLOR_PALETTE[int(time.time()) % len(COLOR_PALETTE)], dtype=np.float32) / 255.0
    color = color[None, None, :]
    org_image = video_state["origin_images"][0].astype(np.float32) / 255.0
    painted_image = (1 - mask * 0.5) * org_image + mask * 0.5 * color
    painted_image = np.uint8(np.clip(painted_image * 255, 0, 255))
    video_state["painted_images"] = np.expand_dims(painted_image, axis=0)
    video_state["masks"] = np.expand_dims(mask[:,:,0], axis=0)

    for i in range(len(video_state["input_points"])):
        point = video_state["input_points"][i]
        if video_state["input_labels"][i] == 0:
            cv2.circle(painted_image, point, radius=3, color=(0, 0, 255), thickness=-1)  # 红色点，半径为3
        else:
            cv2.circle(painted_image, point, radius=3, color=(255, 0, 0), thickness=-1)

    return Image.fromarray(painted_image)

def clear_clicks(video_state):
    video_state["input_points"] = []
    video_state["input_labels"] = []
    video_state["scaled_points"] = []
    video_state["inference_state"] = None
    video_state["masks"] = None
    video_state["painted_images"] = None
    return Image.fromarray(video_state["origin_images"][0]) if video_state["origin_images"] is not None else None

    
def load_uploaded_mask_video(video_path, mask_video_path, n_frames, video_state):
    """Load a user-supplied mask video (white = object) and use it directly,
    skipping click segmentation and SAM2 tracking."""
    src = video_state.get("video_path") or video_path
    if src is None or mask_video_path is None:
        raise gr.Error("Upload both the source video and the mask video first.")
    video_state["video_path"] = src   # <-- persist for the Remove step

    # --- load source frames (same logic as track_video) ---
    vr = VideoReader(src, ctx=cpu(0))
    images = [vr[i].asnumpy() for i in range(min(len(vr), int(n_frames)))]
    del vr

    if images[0].shape[0] > images[0].shape[1]:
        W_ = W
        H_ = int(W_ * images[0].shape[0] / images[0].shape[1])
    else:
        H_ = H
        W_ = int(H_ * images[0].shape[1] / images[0].shape[0])
    images = [cv2.resize(img, (W_, H_)) for img in images]
    video_state["origin_images"] = images

    # --- load mask frames ---
    mvr = VideoReader(mask_video_path, ctx=cpu(0))
    mask_frames = []
    for i in range(min(len(mvr), len(images))):
        m = mvr[i].asnumpy()
        if m.ndim == 3:
            m = cv2.cvtColor(m, cv2.COLOR_RGB2GRAY)
        m = cv2.resize(m, (W_, H_), interpolation=cv2.INTER_NEAREST)
        m = (m > 127).astype(np.float32)
        # match the (H_, W_, 3) shape that track_video produces
        mask_frames.append(np.repeat(m[..., None], 3, axis=2))
    del mvr

    while len(mask_frames) < len(images):      # pad if mask video is shorter
        mask_frames.append(mask_frames[-1].copy())
    video_state["masks"] = mask_frames

    # preview: overlay first mask on first frame
    color = np.array(COLOR_PALETTE[0], dtype=np.float32) / 255.0
    frame = images[0].astype(np.float32) / 255.0
    painted = (1 - mask_frames[0] * 0.5) * frame + mask_frames[0] * 0.5 * color
    preview = np.uint8(np.clip(painted * 255, 0, 255))
    return video_state, Image.fromarray(preview)

def get_source_info(video_state):
    path = video_state.get("video_path")
    if path is None:
        raise gr.Error(
            "No source video found. Upload a video, then click "
            "'Extract First Frame' or 'Use Uploaded Mask' before Remove."
        )
    vr = VideoReader(path, ctx=cpu(0))
    fps = float(vr.get_avg_fps())
    h, w = vr[0].shape[:2]
    del vr
    return (fps if fps and fps > 0 else 15.0), (w, h)

def on_video_change(video_path):
    if not video_path:
        return gr.update()
    vr = VideoReader(video_path, ctx=cpu(0))
    n = len(vr)
    del vr
    return gr.update(maximum=max(9000, n), value=n, label=f"Tracking Frames N (video has {n})")

def preprocess_for_removal(images, masks):
    out_images = []
    out_masks = []
    for img, msk in zip(images, masks):
        if img.shape[0] > img.shape[1]:
            img_resized = cv2.resize(img, (480, 832), interpolation=cv2.INTER_LINEAR)
        else:
            img_resized = cv2.resize(img, (832, 480), interpolation=cv2.INTER_LINEAR)
        img_resized = img_resized.astype(np.float32) / 127.5 - 1.0  # [-1, 1]
        out_images.append(img_resized)
        if msk.shape[0] > msk.shape[1]:
            msk_resized = cv2.resize(msk, (480, 832), interpolation=cv2.INTER_NEAREST)
        else:
            msk_resized = cv2.resize(msk, (832, 480), interpolation=cv2.INTER_NEAREST)
        msk_resized = msk_resized.astype(np.float32)
        msk_resized = (msk_resized > 0.5).astype(np.float32)
        out_masks.append(msk_resized)
    arr_images = np.stack(out_images)
    arr_masks = np.stack(out_masks)
    return torch.from_numpy(arr_images).half(), torch.from_numpy(arr_masks).half()

MAX_PROC_SIDE = None  # None = native resolution; set e.g. 1280 if you hit OOM

@spaces.GPU(duration=1200)  # 1080p is heavy — give it time
def inference_and_return_video(dilation_iterations, num_inference_steps, video_path, video_state):
    if video_state.get("video_path") is None:
        video_state["video_path"] = video_path
    if video_state["origin_images"] is None or video_state["masks"] is None:
        return None

    masks = video_state["masks"]
    fps, (orig_w, orig_h) = get_source_info(video_state)  # helper from before

    # --- load full-res frames straight from the source video ---
    vr = VideoReader(video_state["video_path"], ctx=cpu(0))
    images = [vr[i].asnumpy() for i in range(min(len(masks), len(vr)))]
    del vr
    n_frames = len(images)

    # --- processing size: native (optionally capped), rounded UP to a multiple of 16 ---
    proc_w, proc_h = orig_w, orig_h
    if MAX_PROC_SIDE and max(proc_w, proc_h) > MAX_PROC_SIDE:
        s = MAX_PROC_SIDE / max(proc_w, proc_h)
        proc_w, proc_h = round(proc_w * s), round(proc_h * s)
    W16 = (proc_w + 15) // 16 * 16   # 1920 -> 1920
    H16 = (proc_h + 15) // 16 * 16   # 1080 -> 1088
    pad_w, pad_h = W16 - proc_w, H16 - proc_h

    proc_images, proc_masks = [], []
    for img, msk in zip(images, masks):
        img = cv2.resize(img, (proc_w, proc_h), interpolation=cv2.INTER_LINEAR)
        msk = cv2.resize(msk, (proc_w, proc_h), interpolation=cv2.INTER_NEAREST)
        img = img.astype(np.float32) / 127.5 - 1.0          # [-1, 1]
        msk = (msk > 0.5).astype(np.float32)
        if msk.ndim == 2:
            msk = msk[..., None]
        if pad_h or pad_w:  # reflect-pad image, zero-pad mask (background)
            img = cv2.copyMakeBorder(img, 0, pad_h, 0, pad_w, cv2.BORDER_REFLECT)
            msk = cv2.copyMakeBorder(msk, 0, pad_h, 0, pad_w, cv2.BORDER_CONSTANT, value=0)
        proc_images.append(img)
        proc_masks.append(msk)

    img_tensor = torch.from_numpy(np.stack(proc_images)).half().to("cuda")
    mask_tensor = torch.from_numpy(np.stack(proc_masks)).half()[..., :1].to("cuda")

    pipe.to("cuda")
    with torch.no_grad():
        out = pipe(
            images=img_tensor,
            masks=mask_tensor,
            num_frames=n_frames,
            height=H16,
            width=W16,
            num_inference_steps=int(num_inference_steps),
            generator=torch.Generator(device=device).manual_seed(random_seed),
            iterations=int(dilation_iterations),
        ).frames[0]

    out = np.uint8(out * 255)
    out = out[:, :proc_h, :proc_w]                       # strip the padding back off
    if (proc_w, proc_h) != (orig_w, orig_h):             # only if MAX_PROC_SIDE kicked in
        out = [cv2.resize(f, (orig_w, orig_h), interpolation=cv2.INTER_CUBIC) for f in out]

    video_file = f"/tmp/{time.time()}-{random.random()}-removed_output.mp4"
    clip = ImageSequenceClip(list(out), fps=fps)
    clip.write_videofile(
        video_file,
        codec="libx264",
        audio=False,
        fps=fps,
        preset="medium",
        ffmpeg_params=["-crf", "17", "-pix_fmt", "yuv420p"],
        verbose=False, logger=None
    )
    return video_file

@spaces.GPU(duration=40)
def track_video(n_frames,video_state):
    input_points = video_state["input_points"]
    input_labels = video_state["input_labels"]
    frame_idx = video_state["frame_idx"]
    obj_id = video_state["obj_id"]
    scaled_points = video_state["scaled_points"]

    vr = VideoReader(video_state["video_path"], ctx=cpu(0))
    height, width = vr[0].shape[0:2]
    images = [vr[i].asnumpy() for i in range(min(len(vr), n_frames))]
    del vr

    if images[0].shape[0] > images[0].shape[1]:
        W_ = W
        H_ = int(W_ * images[0].shape[0] / images[0].shape[1])
    else:
        H_ = H
        W_ = int(H_ * images[0].shape[1] / images[0].shape[0])

    images = [cv2.resize(img, (W_, H_)) for img in images]
    video_state["origin_images"] = images
    images = np.array(images)

    sam2_checkpoint = "./SAM2-Video-Predictor/checkpoints/sam2_hiera_large.pt"
    config = "sam2_hiera_l.yaml"
    video_predictor_local = build_sam2_video_predictor(config, sam2_checkpoint, device="cuda")
    
    inference_state = video_predictor_local.init_state(images=images/255, device="cuda")
    #video_state["inference_state"] = inference_state #cause bug

    if len(torch.from_numpy(video_state["masks"][0]).shape) == 3:
        mask = torch.from_numpy(video_state["masks"][0])[:,:,0]
    else:
        mask = torch.from_numpy(video_state["masks"][0])

    video_predictor_local.add_new_mask(
        inference_state=inference_state,
        frame_idx=0,
        obj_id=obj_id,
        mask=mask
    )

    output_frames = []
    mask_frames = []
    color = np.array(COLOR_PALETTE[int(time.time()) % len(COLOR_PALETTE)], dtype=np.float32) / 255.0
    color = color[None, None, :]
    for out_frame_idx, out_obj_ids, out_mask_logits in video_predictor_local.propagate_in_video(inference_state):
        frame = images[out_frame_idx].astype(np.float32) / 255.0
        mask = np.zeros((H, W, 3), dtype=np.float32)
        for i, logit in enumerate(out_mask_logits):
            out_mask = logit.cpu().squeeze().detach().numpy()
            out_mask = (out_mask[:,:,None] > 0).astype(np.float32)
            mask += out_mask
        mask = np.clip(mask, 0, 1)
        mask = cv2.resize(mask, (W_, H_))
        mask_frames.append(mask)
        painted = (1 - mask * 0.5) * frame + mask * 0.5 * color
        painted = np.uint8(np.clip(painted * 255, 0, 255))
        output_frames.append(painted)
    video_state["masks"] =mask_frames
    video_file = f"/tmp/{time.time()}-{random.random()}-tracked_output.mp4"
    clip = ImageSequenceClip(output_frames, fps=15)
    clip.write_videofile(video_file, codec='libx264', audio=False, verbose=False, logger=None)
    print("line 286 done")
    return video_file,video_state

text = """
<div style='text-align:center; font-size:32px; font-family: Arial, Helvetica, sans-serif;'>
  Minimax-Remover: Taming Bad Noise Helps Video Object Removal
</div>
<div style="display: flex; justify-content: center; align-items: center; gap: 10px; flex-wrap: nowrap;">
  <a href="https://huggingface.co/zibojia/minimax-remover"><img alt="Huggingface Model" src="https://img.shields.io/badge/%F0%9F%A4%97%20Huggingface-Model-brightgreen"></a>
  <a href="https://github.com/zibojia/MiniMax-Remover"><img alt="Github" src="https://img.shields.io/badge/MiniMaxRemover-github-black"></a>
  <a href="https://huggingface.co/spaces/PengWeixuanSZU/MiniMax-Remover"><img alt="Huggingface Space" src="https://img.shields.io/badge/%F0%9F%A4%97%20Huggingface-Space-1e90ff"></a>
  <a href="https://arxiv.org/abs/2505.24873"><img alt="arXiv" src="https://img.shields.io/badge/MiniMaxRemover-arXiv-b31b1b"></a>
  <a href="https://www.youtube.com/watch?v=KaU5yNl6CTc"><img alt="YouTube" src="https://img.shields.io/badge/Youtube-video-ff0000"></a>
  <a href="https://minimax-remover.github.io"><img alt="Demo Page" src="https://img.shields.io/badge/Website-Demo%20Page-yellow"></a>
</div>
<div style='text-align:center; font-size:20px; margin-top: 10px; font-family: Arial, Helvetica, sans-serif;'>
  Bojia Zi<sup>*</sup>, Weixuan Peng<sup>*</sup>, Xianbiao Qi<sup>†</sup>, Jianan Wang, Shihao Zhao, Rong Xiao, Kam-Fai Wong
</div>
<div style='text-align:center; font-size:14px; color: #888; margin-top: 5px; font-family: Arial, Helvetica, sans-serif;'>
  <sup>*</sup> Equal contribution &nbsp; &nbsp; <sup>†</sup> Corresponding author
</div>
"""

pipe, image_predictor, video_predictor = get_pipe_image_and_video_predictor()

with gr.Blocks() as demo:
    video_state = gr.State({
        "origin_images": None,
        "inference_state": None,
        "masks": None,  # Store user-generated masks
        "painted_images": None,
        "video_path": None,
        "input_points": [],
        "scaled_points": [],
        "input_labels": [],
        "frame_idx": 0,
        "obj_id": 1
    })
    gr.Markdown(f"<div style='text-align:center;'>{text}</div>")

    with gr.Column():
        video_input = gr.Video(label="Upload Video", elem_id="my-video1")
        get_info_btn = gr.Button("Extract First Frame", elem_id="my-btn")

        gr.Examples(
            examples=[
                ["./cartoon/0.mp4"],
                ["./cartoon/1.mp4"],
                ["./cartoon/2.mp4"],
                ["./cartoon/3.mp4"],
                ["./cartoon/4.mp4"],
                ["./normal_videos/0.mp4"],
                ["./normal_videos/1.mp4"],
                ["./normal_videos/3.mp4"],
                ["./normal_videos/4.mp4"],
                ["./normal_videos/5.mp4"],
            ],
            inputs=[video_input],
            label="Choose a video to remove.",
            elem_id="my-btn2"
        )

        image_output = gr.Image(label="First Frame Segmentation", interactive=True, elem_id="my-video")#, height="35%", width="60%")
        demo.css = """
        #my-btn {
           width: 60% !important;
           margin: 0 auto;
        }

        #my-video1 {
           width: 60% !important;
           height: 35% !important;
           margin: 0 auto;
        }
        #my-video {
           width: 60% !important;
           height: 35% !important;
           margin: 0 auto;
        }
        #my-md {
           margin: 0 auto;
        }
        #my-btn2 {
            width: 60% !important;
            margin: 0 auto;
        }
        #my-btn2 button {
            width: 120px !important;
            max-width: 120px !important;
            min-width: 120px !important;
            height: 70px !important;
            max-height: 70px !important;
            min-height: 70px !important;
            margin: 8px !important;
            border-radius: 8px !important;
            overflow: hidden !important;
            white-space: normal !important;
        }
        """
        with gr.Row(elem_id="my-btn"):
            point_prompt = gr.Radio(["Positive", "Negative"], label="Click Type", value="Positive")
            clear_btn = gr.Button("Clear All Clicks")

        with gr.Row(elem_id="my-btn"):
            n_frames_slider = gr.Slider(minimum=1, maximum=9000, value=81, step=1, label="Tracking Frames N")
            track_btn = gr.Button("Tracking")
        video_output = gr.Video(label="Tracking Result", elem_id="my-video")
        video_input.change(on_video_change, inputs=video_input, outputs=n_frames_slider)

        mask_video_input = gr.Video(label="(Optional) Upload Mask Video — white = object to remove")
        use_mask_btn = gr.Button("Use Uploaded Mask (skips clicking & tracking)", elem_id="my-btn")
        
        use_mask_btn.click(
            load_uploaded_mask_video,
            inputs=[video_input, mask_video_input, n_frames_slider, video_state],
            outputs=[video_state, image_output],
        )

        with gr.Column(elem_id="my-btn"):
            dilation_slider = gr.Slider(minimum=1, maximum=20, value=6, step=1, label="Mask Dilation")
            inference_steps_slider = gr.Slider(minimum=1, maximum=100, value=6, step=1, label="Num Inference Steps")
          
        remove_btn = gr.Button("Remove", elem_id="my-btn")
        remove_video = gr.Video(label="Remove Results", elem_id="my-video")
        remove_btn.click(
            inference_and_return_video,
            inputs=[dilation_slider, inference_steps_slider, video_input, video_state],
            outputs=remove_video
        )
        
        get_info_btn.click(get_video_info, inputs=[video_input, video_state], \
                       outputs=image_output)
        image_output.select(fn=segment_frame, inputs=[point_prompt, video_state], outputs=image_output)
        clear_btn.click(clear_clicks, inputs=video_state, outputs=image_output)
        track_btn.click(track_video, inputs=[n_frames_slider,video_state], outputs=[video_output,video_state])

demo.launch()