Spaces:
Running
Running
Upload seg_script.py
Browse files- seg_script.py +104 -0
seg_script.py
ADDED
@@ -0,0 +1,104 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
'''
|
2 |
+
python seg_script.py Genshin_Impact_Images Genshin_Impact_Images_Seg
|
3 |
+
'''
|
4 |
+
|
5 |
+
import os
|
6 |
+
import cv2
|
7 |
+
import argparse
|
8 |
+
from PIL import Image
|
9 |
+
import numpy as np
|
10 |
+
from tqdm import tqdm
|
11 |
+
from pathlib import Path
|
12 |
+
|
13 |
+
from animeinsseg import AnimeInsSeg, AnimeInstances
|
14 |
+
from animeinsseg.anime_instances import get_color
|
15 |
+
|
16 |
+
# 设置模型路径
|
17 |
+
ckpt = r'models/AnimeInstanceSegmentation/rtmdetl_e60.ckpt'
|
18 |
+
mask_thres = 0.3
|
19 |
+
instance_thres = 0.3
|
20 |
+
refine_kwargs = {'refine_method': 'refinenet_isnet'} # 如果不使用 refinenet,设置为 None
|
21 |
+
# refine_kwargs = None
|
22 |
+
|
23 |
+
# 初始化模型
|
24 |
+
net = AnimeInsSeg(ckpt, mask_thr=mask_thres, refine_kwargs=refine_kwargs)
|
25 |
+
|
26 |
+
def process_image(image_path, output_dir):
|
27 |
+
# 读取图像
|
28 |
+
img = cv2.imread(image_path)
|
29 |
+
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
30 |
+
|
31 |
+
# 推理
|
32 |
+
instances: AnimeInstances = net.infer(
|
33 |
+
img,
|
34 |
+
output_type='numpy',
|
35 |
+
pred_score_thr=instance_thres
|
36 |
+
)
|
37 |
+
|
38 |
+
# 初始化输出图像
|
39 |
+
drawed = img.copy()
|
40 |
+
im_h, im_w = img.shape[:2]
|
41 |
+
|
42 |
+
# 如果没有检测到对象,直接返回原图
|
43 |
+
if instances.bboxes is None:
|
44 |
+
return
|
45 |
+
|
46 |
+
# 保存绘制后的图像(只保存一次)
|
47 |
+
base_name = Path(image_path).stem
|
48 |
+
output_path = Path(output_dir)
|
49 |
+
output_path.mkdir(parents=True, exist_ok=True)
|
50 |
+
Image.fromarray(drawed).save(output_path / f"{base_name}_drawed.png")
|
51 |
+
|
52 |
+
# 处理每个实例
|
53 |
+
for ii, (xywh, mask) in enumerate(zip(instances.bboxes, instances.masks)):
|
54 |
+
color = get_color(ii)
|
55 |
+
|
56 |
+
mask_alpha = 0.5
|
57 |
+
linewidth = max(round(sum(img.shape) / 2 * 0.003), 2)
|
58 |
+
|
59 |
+
# 绘制边界框
|
60 |
+
p1, p2 = (int(xywh[0]), int(xywh[1])), (int(xywh[2] + xywh[0]), int(xywh[3] + xywh[1]))
|
61 |
+
cv2.rectangle(drawed, p1, p2, color, thickness=linewidth, lineType=cv2.LINE_AA)
|
62 |
+
|
63 |
+
# 绘制掩码
|
64 |
+
p = mask.astype(np.float32)
|
65 |
+
blend_mask = np.full((im_h, im_w, 3), color, dtype=np.float32)
|
66 |
+
alpha_msk = (mask_alpha * p)[..., None]
|
67 |
+
alpha_ori = 1 - alpha_msk
|
68 |
+
drawed = drawed * alpha_ori + alpha_msk * blend_mask
|
69 |
+
|
70 |
+
drawed = drawed.astype(np.uint8)
|
71 |
+
|
72 |
+
# 裁剪图像
|
73 |
+
x1, y1, x2, y2 = int(xywh[0]), int(xywh[1]), int(xywh[0] + xywh[2]), int(xywh[1] + xywh[3])
|
74 |
+
cropped_img = img[y1:y2, x1:x2]
|
75 |
+
cropped_mask = mask[y1:y2, x1:x2]
|
76 |
+
|
77 |
+
# 创建透明通道的边缘图
|
78 |
+
alpha_channel = (cropped_mask * 255).astype(np.uint8)
|
79 |
+
rgba_image = np.dstack((cropped_img, alpha_channel))
|
80 |
+
|
81 |
+
# 保存裁剪后的图像和分割后的图像(文件名包含实例下标)
|
82 |
+
Image.fromarray(cropped_img).save(output_path / f"{base_name}_cropped_{ii}.png")
|
83 |
+
Image.fromarray(rgba_image, 'RGBA').save(output_path / f"{base_name}_segmented_{ii}.png")
|
84 |
+
|
85 |
+
def main():
|
86 |
+
parser = argparse.ArgumentParser(description="Anime Instance Segmentation")
|
87 |
+
parser.add_argument("input_path", type=str, help="Path to the input image or folder")
|
88 |
+
parser.add_argument("output_dir", type=str, help="Path to the output directory")
|
89 |
+
args = parser.parse_args()
|
90 |
+
|
91 |
+
input_path = Path(args.input_path)
|
92 |
+
output_dir = Path(args.output_dir)
|
93 |
+
|
94 |
+
if input_path.is_file():
|
95 |
+
process_image(input_path, output_dir)
|
96 |
+
elif input_path.is_dir():
|
97 |
+
image_paths = list(input_path.rglob("*.png")) + list(input_path.rglob("*.jpg"))
|
98 |
+
for image_path in tqdm(image_paths, desc="Processing images"):
|
99 |
+
process_image(image_path, output_dir)
|
100 |
+
else:
|
101 |
+
print("Invalid input path. Please provide a valid image or folder path.")
|
102 |
+
|
103 |
+
if __name__ == "__main__":
|
104 |
+
main()
|