Update app.py

app.py

@@ -9,60 +9,101 @@ 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")
+os.system(
+    "git clone https://huggingface.co/dreMaz/AnimeInstanceSegmentation models/AnimeInstanceSegmentation"
+)
 
-ckpt = r'models/AnimeInstanceSegmentation/rtmdetl_e60.ckpt'
+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 = {
+    "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
+        img, output_type="numpy", pred_score_thr=instance_thres
     )
 
-    # 创建一个空白的白色图像，和原图大小一致
-    white = np.full_like(img, 255)
+    # 创建一个临时文件夹，用来存放每个人物的去除背景的图片
+    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 Image.fromarray(white)
+        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_)
 
         # 用原图中对应的区域替换白色图像中的区域，实现去除背景的效果
-        white[mask] = img[mask]
+        white[mask] = img2[mask]
+
+        # 给每个人物编号，然后用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 Image.fromarray(white[..., ::-1])
+    # 返回一个图片文件和一个压缩包文件
+    return images, zip_path
 
+
+# 在gr.Interface中添加outputs参数，把fn函数的输出作为一个列表传入
 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"]
+    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()
-
+iface.launch()