|
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" |
|
} |
|
|
|
|
|
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 = [] |
|
|
|
|
|
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 = 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] |
|
|
|
|
|
filename = f"person_{ii+1}.png" |
|
filepath = os.path.join(temp_dir, filename) |
|
images.append(white) |
|
cv2.imwrite(filepath, white[..., ::-1]) |
|
|
|
|
|
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 |
|
|
|
|
|
|
|
iface = gr.Interface( |
|
|
|
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() |