Spaces:
Runtime error
Runtime error
Update app.py
Browse files
app.py
CHANGED
@@ -9,60 +9,101 @@ os.system("mim install mmdet")
|
|
9 |
import cv2
|
10 |
from PIL import Image
|
11 |
import numpy as np
|
|
|
|
|
12 |
|
13 |
from animeinsseg import AnimeInsSeg, AnimeInstances
|
14 |
from animeinsseg.anime_instances import get_color
|
15 |
|
16 |
|
17 |
-
|
18 |
if not os.path.exists("models"):
|
19 |
os.mkdir("models")
|
20 |
|
21 |
os.system("huggingface-cli lfs-enable-largefiles .")
|
22 |
-
os.system(
|
|
|
|
|
23 |
|
24 |
-
ckpt = r
|
25 |
|
26 |
mask_thres = 0.3
|
27 |
instance_thres = 0.3
|
28 |
-
refine_kwargs = {
|
|
|
|
|
29 |
# refine_kwargs = None
|
30 |
|
31 |
net = AnimeInsSeg(ckpt, mask_thr=mask_thres, refine_kwargs=refine_kwargs)
|
32 |
|
|
|
33 |
def fn(image):
|
34 |
img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
|
35 |
instances: AnimeInstances = net.infer(
|
36 |
-
img,
|
37 |
-
output_type='numpy',
|
38 |
-
pred_score_thr=instance_thres
|
39 |
)
|
40 |
|
41 |
-
#
|
42 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
43 |
|
44 |
# instances.bboxes, instances.masks will be None, None if no obj is detected
|
45 |
if instances.bboxes is None:
|
46 |
-
return
|
|
|
|
|
47 |
|
48 |
for ii, (xywh, mask) in enumerate(zip(instances.bboxes, instances.masks)):
|
|
|
|
|
|
|
49 |
# 把mask转换为bool类型,方便后续操作
|
50 |
mask = mask.astype(np.bool_)
|
51 |
|
52 |
# 用原图中对应的区域替换白色图像中的区域,实现去除背景的效果
|
53 |
-
white[mask] =
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
54 |
|
55 |
-
|
|
|
56 |
|
|
|
|
|
57 |
iface = gr.Interface(
|
58 |
# design titles and text descriptions
|
59 |
title="Anime Subject Instance Segmentation",
|
60 |
-
description="Segment image subjects with the proposed model in the paper [*Instance-guided Cartoon Editing with a Large-scale Dataset*](https://cartoonsegmentation.github.io/).",
|
61 |
fn=fn,
|
62 |
inputs=gr.Image(type="numpy"),
|
63 |
-
outputs=
|
64 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
65 |
)
|
66 |
|
67 |
-
iface.launch()
|
68 |
-
|
|
|
9 |
import cv2
|
10 |
from PIL import Image
|
11 |
import numpy as np
|
12 |
+
import zipfile
|
13 |
+
import shutil
|
14 |
|
15 |
from animeinsseg import AnimeInsSeg, AnimeInstances
|
16 |
from animeinsseg.anime_instances import get_color
|
17 |
|
18 |
|
|
|
19 |
if not os.path.exists("models"):
|
20 |
os.mkdir("models")
|
21 |
|
22 |
os.system("huggingface-cli lfs-enable-largefiles .")
|
23 |
+
os.system(
|
24 |
+
"git clone https://huggingface.co/dreMaz/AnimeInstanceSegmentation models/AnimeInstanceSegmentation"
|
25 |
+
)
|
26 |
|
27 |
+
ckpt = r"models/AnimeInstanceSegmentation/rtmdetl_e60.ckpt"
|
28 |
|
29 |
mask_thres = 0.3
|
30 |
instance_thres = 0.3
|
31 |
+
refine_kwargs = {
|
32 |
+
"refine_method": "refinenet_isnet"
|
33 |
+
} # set to None if not using refinenet
|
34 |
# refine_kwargs = None
|
35 |
|
36 |
net = AnimeInsSeg(ckpt, mask_thr=mask_thres, refine_kwargs=refine_kwargs)
|
37 |
|
38 |
+
|
39 |
def fn(image):
|
40 |
img = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
|
41 |
instances: AnimeInstances = net.infer(
|
42 |
+
img, output_type="numpy", pred_score_thr=instance_thres
|
|
|
|
|
43 |
)
|
44 |
|
45 |
+
# 创建一个临时文件夹,用来存放每个人物的去除背景的图片
|
46 |
+
temp_dir = "outputs"
|
47 |
+
# 删除临时文件夹,避免占用空间
|
48 |
+
if os.path.isdir(temp_dir):
|
49 |
+
shutil.rmtree(temp_dir)
|
50 |
+
os.makedirs(temp_dir, exist_ok=True)
|
51 |
+
|
52 |
+
images = []
|
53 |
|
54 |
# instances.bboxes, instances.masks will be None, None if no obj is detected
|
55 |
if instances.bboxes is None:
|
56 |
+
return None
|
57 |
+
|
58 |
+
img2 = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
59 |
|
60 |
for ii, (xywh, mask) in enumerate(zip(instances.bboxes, instances.masks)):
|
61 |
+
# 创建一个空白的白色图像,和原图大小一致
|
62 |
+
white = np.full_like(img2, 255)
|
63 |
+
|
64 |
# 把mask转换为bool类型,方便后续操作
|
65 |
mask = mask.astype(np.bool_)
|
66 |
|
67 |
# 用原图中对应的区域替换白色图像中的区域,实现去除背景的效果
|
68 |
+
white[mask] = img2[mask]
|
69 |
+
|
70 |
+
# 给每个人物编号,然后用cv2.imwrite函数来保存图片到文件夹中
|
71 |
+
filename = f"person_{ii+1}.png"
|
72 |
+
filepath = os.path.join(temp_dir, filename)
|
73 |
+
images.append(white)
|
74 |
+
cv2.imwrite(filepath, white[..., ::-1])
|
75 |
+
|
76 |
+
# 创建一个压缩包,然后用zipfile.write函数来把文件夹中的所有图片添加到压缩包中
|
77 |
+
zip_name = "persons.zip"
|
78 |
+
zip_path = os.path.join(temp_dir, zip_name)
|
79 |
+
with zipfile.ZipFile(zip_path, "w") as zf:
|
80 |
+
for file in os.listdir(temp_dir):
|
81 |
+
if file != zip_name:
|
82 |
+
zf.write(os.path.join(temp_dir, file), file)
|
83 |
|
84 |
+
# 返回一个图片文件和一个压缩包文件
|
85 |
+
return images, zip_path
|
86 |
|
87 |
+
|
88 |
+
# 在gr.Interface中添加outputs参数,把fn函数的输出作为一个列表传入
|
89 |
iface = gr.Interface(
|
90 |
# design titles and text descriptions
|
91 |
title="Anime Subject Instance Segmentation",
|
|
|
92 |
fn=fn,
|
93 |
inputs=gr.Image(type="numpy"),
|
94 |
+
outputs=[
|
95 |
+
gr.Gallery(
|
96 |
+
label="Anime Subject Instance Segmentation",
|
97 |
+
show_label=False,
|
98 |
+
elem_id="gallery",
|
99 |
+
columns=[4],
|
100 |
+
rows=[4],
|
101 |
+
object_fit="contain",
|
102 |
+
height="auto",
|
103 |
+
),
|
104 |
+
gr.File(type="filepath", label="Download Zip"),
|
105 |
+
], # 添加一个压缩包组件,给它一个名称
|
106 |
+
examples=["1562990.jpg", "612989.jpg", "sample_3.jpg"],
|
107 |
)
|
108 |
|
109 |
+
iface.launch()
|
|