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
import zipfile
import shutil

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.7
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
    )

    # 创建一个临时文件夹，用来存放每个人物的去除背景的图片
    temp_dir = "outputs"
    # 删除临时文件夹，避免占用空间
    if os.path.isdir(temp_dir):
        shutil.rmtree(temp_dir)
    os.makedirs(temp_dir, exist_ok=True)

    images = []

    # instances.bboxes, instances.masks will be None, None if no obj is detected
    if instances.bboxes is None:
        return None

    img2 = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

    for ii, (xywh, mask) in enumerate(zip(instances.bboxes, instances.masks)):
        # 创建一个空白的白色图像，和原图大小一致
        white = np.full_like(img2, 255)

        # 把mask转换为bool类型，方便后续操作
        mask = mask.astype(np.bool_)

        # 对掩码进行高斯模糊，平滑边缘
        mask_smoothed = cv2.GaussianBlur(mask.astype(np.float32), (3, 3), 0)

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

        # 给每个人物编号，然后用cv2.imwrite函数来保存图片到文件夹中
        filename = f"person_{ii+1}.png"
        filepath = os.path.join(temp_dir, filename)
        images.append(white)
        cv2.imwrite(filepath, white[..., ::-1])

    # 创建一个压缩包，然后用zipfile.write函数来把文件夹中的所有图片添加到压缩包中
    zip_name = "persons.zip"
    zip_path = os.path.join(temp_dir, zip_name)
    with zipfile.ZipFile(zip_path, "w") as zf:
        for file in os.listdir(temp_dir):
            if file != zip_name:
                zf.write(os.path.join(temp_dir, file), file)

    # 返回一个图片文件和一个压缩包文件
    return images, zip_path


# 在gr.Interface中添加outputs参数，把fn函数的输出作为一个列表传入
iface = gr.Interface(
    # design titles and text descriptions
    title="Anime Subject Instance Segmentation",
    fn=fn,
    inputs=gr.Image(type="numpy"),
    outputs=[
        gr.Gallery(
            label="Anime Subject Instance Segmentation",
            show_label=False,
            elem_id="gallery",
            columns=[4],
            rows=[4],
            object_fit="contain",
            height="auto",
        ),
        gr.File(type="filepath", label="Download Zip"),
    ],  # 添加一个压缩包组件，给它一个名称
    examples=["1562990.jpg", "612989.jpg", "sample_3.jpg"],
)

iface.launch()