svjack commited on
Commit
e81f668
·
verified ·
1 Parent(s): 596acf1

Upload seg_script.py

Browse files
Files changed (1) hide show
  1. 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()