import gradio as gr

import os

os.system("mim install mmengine")
os.system('mim install "mmcv>=2.0.0"')
os.system("mim install mmdet")

import cv2
from PIL import Image
import numpy as np

from animeinsseg import AnimeInsSeg, AnimeInstances
from animeinsseg.anime_instances import get_color



if not os.path.exists("models"):
    os.mkdir("models")

os.system("huggingface-cli lfs-enable-largefiles .")
os.system("git clone https://huggingface.co/dreMaz/AnimeInstanceSegmentation models/AnimeInstanceSegmentation")

ckpt = r'models/AnimeInstanceSegmentation/rtmdetl_e60.ckpt'

mask_thres = 0.3
instance_thres = 0.3
refine_kwargs = {'refine_method': 'refinenet_isnet'} # set to None if not using refinenet
# refine_kwargs = None

net = AnimeInsSeg(ckpt, mask_thr=mask_thres, refine_kwargs=refine_kwargs)

def fn(image):
    img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    instances: AnimeInstances = net.infer(
        img,
        output_type='numpy',
        pred_score_thr=instance_thres
    )

    # 创建一个空白的白色图像，和原图大小一致
    white = np.full_like(img, 255)

    # instances.bboxes, instances.masks will be None, None if no obj is detected
    if instances.bboxes is None:
        return Image.fromarray(white)

    for ii, (xywh, mask) in enumerate(zip(instances.bboxes, instances.masks)):
        # 把mask转换为bool类型，方便后续操作
        mask = mask.astype(np.bool_)

        # 用原图中对应的区域替换白色图像中的区域，实现去除背景的效果
        white[mask] = img[mask]

    return Image.fromarray(white[..., ::-1])

iface = gr.Interface(
    # design titles and text descriptions
    title="Anime Subject Instance Segmentation",
    description="Segment image subjects with the proposed model in the paper [*Instance-guided Cartoon Editing with a Large-scale Dataset*](https://cartoonsegmentation.github.io/).",
    fn=fn,
    inputs=gr.Image(type="numpy"),
    outputs=gr.Image(type="pil"),
    examples=["1562990.jpg", "612989.jpg", "sample_3.jpg"]
)

iface.launch()