plutosss commited on
Commit
99fac33
·
verified ·
1 Parent(s): 5cbca7c

Delete process_line修改版(一键生成黑底白线稿).py

Browse files
process_line修改版(一键生成黑底白线稿).py DELETED
@@ -1,367 +0,0 @@
1
- # 部署 teed、depth-anything
2
- # 腐蚀算法
3
- # 读取图片
4
- # 输出图片
5
- # 使用 depth-anything + teed 生成外轮廓
6
- # 使用 teed + 腐蚀算法 生成内边缘
7
- from PIL import Image
8
-
9
- import cv2
10
- import cv2_ext
11
- import numpy as np
12
- import os
13
- import torch
14
- import torch.nn.functional as F
15
- from torchvision.transforms import Compose
16
- from tqdm import tqdm
17
- import TEED.main as teed
18
- from TEED.main import parse_args
19
-
20
- from depthAnything.depth_anything.dpt import DepthAnything
21
- from depthAnything.depth_anything.util.transform import Resize, NormalizeImage, PrepareForNet
22
- import shutil
23
-
24
-
25
- def multiply_blend(image1, image2):
26
- # 将图片转换为浮点数,方便计算
27
- # Ensure image2 has the same shape as image1
28
- image2 = np.stack((image2,) * 3, axis=-1)
29
- # Perform the blending
30
- multiplied = np.multiply(image1 / 255.0, image2 / 255.0) * 255.0
31
- return multiplied.astype(np.uint8)
32
-
33
- # Example usage
34
-
35
-
36
- image1 = np.random.randint(0, 256, (717, 790, 3), dtype=np.uint8)
37
- image2 = np.random.randint(0, 256, (717, 790), dtype=np.uint8)
38
-
39
- result = multiply_blend(image1, image2)
40
- print(result.shape) # Should be (717, 790, 3)
41
-
42
- def screen_blend(image1, image2):
43
- # 将图片转换为浮点数,方便计算
44
- image1 = image1.astype(float)
45
- image2 = image2.astype(float)
46
-
47
- # 执行滤色操作
48
- screened = 1 - (1 - image1 / 255) * (1 - image2 / 255) * 255
49
-
50
- # 将结果转换回uint8
51
- result = np.clip(screened, 0, 255).astype('uint8')
52
- return result
53
-
54
- def erosion(img, kernel_size = 3, iterations = 1, dilate = False):
55
-
56
- # 灰度化
57
- if len(img.shape) == 3:
58
- img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
59
-
60
- # # 二值化
61
- # _, img = cv2.threshold(img, 200, 255, cv2.THRESH_BINARY)
62
-
63
- # 腐蚀
64
- kernel = np.ones((kernel_size, kernel_size), np.uint8)
65
- if dilate:
66
- img = cv2.dilate(img, kernel, iterations=iterations)
67
- else:
68
- img = cv2.erode(img, kernel, iterations=iterations)
69
-
70
- return img
71
-
72
- def erosion_img_from_path(img_path, output_dir = './output/erosion_img', kernel_size = 3, iterations = 1, dilate = False):
73
- # 读取图片
74
- if os.path.isfile(img_path):
75
- name, extension = os.path.splitext(img_path)
76
- if extension:
77
- if extension.lower() == 'txt':
78
- with open(img_path, 'r',encoding= 'utf-8') as f:
79
- filenames = f.read().splitlines()
80
- elif extension.lower() in ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp','tif']:
81
- filenames = [img_path]
82
- else:
83
- filenames = os.listdir(img_path)
84
- filenames = [os.path.join(img_path, filename) for filename in filenames if not filename.startswith('.') and filename.lower().endswith(('.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp','tif'))]
85
- filenames.sort()
86
-
87
- os.makedirs(output_dir, exist_ok=True)
88
-
89
- for filename in tqdm(filenames):
90
- img = cv2.imread(filename)
91
- img = erosion(img, kernel_size, iterations, dilate)
92
- cv2.imwrite(os.path.join(output_dir, os.path.basename(filename)), img)
93
-
94
-
95
- def copy_file(src, dest):
96
- # 移动文件
97
- source = src
98
- destination = dest
99
- try:
100
- shutil.copy(source, destination)
101
- except IOError as e:
102
- print("Unable to copy file. %s" % e)
103
-
104
-
105
- def guassian_blur_path(img_path, output_dir = './output/guassian_blur', kernel_size = 3, sigmaX = 0):
106
- # 读取图片
107
- if os.path.isfile(img_path):
108
- name, extension = os.path.splitext(img_path)
109
- if extension:
110
- if extension.lower() == 'txt':
111
- with open(img_path, 'r',encoding= 'utf-8') as f:
112
- filenames = f.read().splitlines()
113
- elif extension.lower() in ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp','tif']:
114
- filenames = [img_path]
115
- else:
116
- filenames = os.listdir(img_path)
117
- filenames = [os.path.join(img_path, filename) for filename in filenames if not filename.startswith('.') and filename.lower().endswith(('.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp','tif'))]
118
- filenames.sort()
119
-
120
- os.makedirs(output_dir, exist_ok=True)
121
-
122
- for filename in tqdm(filenames):
123
- img = cv2.imread(filename)
124
- img = cv2.GaussianBlur(img, (kernel_size,kernel_size), sigmaX)
125
- cv2.imwrite(os.path.join(output_dir, os.path.basename(filename)), img)
126
-
127
- def depth_anything(img_path = './input', outdir = './output/depth_anything', encoder = 'vitl', pred_only = True, grayscale = True):
128
- # parser = argparse.ArgumentParser()
129
- # parser.add_argument('--img-path', type=str)
130
- # parser.add_argument('--outdir', type=str, default='./vis_depth')
131
- # parser.add_argument('--encoder', type=str, default='vitl', choices=['vits', 'vitb', 'vitl'])
132
-
133
- # parser.add_argument('--pred-only', dest='pred_only', action='store_true', help='only display the prediction')
134
- # parser.add_argument('--grayscale', dest='grayscale', action='store_true', help='do not apply colorful palette')
135
-
136
- # args = parser.parse_args()
137
-
138
- margin_width = 50
139
- caption_height = 60
140
-
141
- font = cv2.FONT_HERSHEY_SIMPLEX
142
- font_scale = 1
143
- font_thickness = 2
144
-
145
- DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
146
-
147
- model_configs = {
148
- 'vitl': {'encoder': 'vitl', 'features': 256, 'out_channels': [256, 512, 1024, 1024]},
149
- 'vitb': {'encoder': 'vitb', 'features': 128, 'out_channels': [96, 192, 384, 768]},
150
- 'vits': {'encoder': 'vits', 'features': 64, 'out_channels': [48, 96, 192, 384]}
151
- }
152
-
153
- depth_anything = DepthAnything(model_configs[encoder])
154
- depth_anything.load_state_dict(torch.load('./checkpoints/depth_anything_{}14.pth'.format(encoder)))
155
- depth_anything = depth_anything.to(DEVICE).eval()
156
-
157
- total_params = sum(param.numel() for param in depth_anything.parameters())
158
- print('Total parameters: {:.2f}M'.format(total_params / 1e6))
159
-
160
- transform = Compose([
161
- Resize(
162
- width=518,
163
- height=518,
164
- resize_target=False,
165
- keep_aspect_ratio=True,
166
- ensure_multiple_of=14,
167
- resize_method='lower_bound',
168
- image_interpolation_method=cv2.INTER_CUBIC,
169
- ),
170
- NormalizeImage(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
171
- PrepareForNet(),
172
- ])
173
-
174
- if os.path.isfile(img_path):
175
- name, extension = os.path.splitext(img_path)
176
- if extension:
177
- if extension.lower() == 'txt':
178
- with open(img_path, 'r',encoding= 'utf-8') as f:
179
- filenames = f.read().splitlines()
180
- elif extension.lower() in ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp','tif']:
181
- filenames = [img_path]
182
- else:
183
- filenames = os.listdir(img_path)
184
- filenames = [os.path.join(img_path, filename) for filename in filenames if not filename.startswith('.') and filename.lower().endswith(('.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp','tif'))]
185
- filenames.sort()
186
-
187
- os.makedirs(outdir, exist_ok=True)
188
-
189
- for filename in tqdm(filenames):
190
- raw_image = cv2.imread(filename)
191
- image = cv2.cvtColor(raw_image, cv2.COLOR_BGR2RGB) / 255.0
192
-
193
- h, w = image.shape[:2]
194
-
195
- image = transform({'image': image})['image']
196
- image = torch.from_numpy(image).unsqueeze(0).to(DEVICE)
197
-
198
- with torch.no_grad():
199
- depth = depth_anything(image)
200
-
201
- depth = F.interpolate(depth[None], (h, w), mode='bilinear', align_corners=False)[0, 0]
202
- depth = (depth - depth.min()) / (depth.max() - depth.min()) * 255.0
203
-
204
- depth = depth.cpu().numpy().astype(np.uint8)
205
-
206
- if grayscale:
207
- depth = np.repeat(depth[..., np.newaxis], 3, axis=-1)
208
- else:
209
- depth = cv2.applyColorMap(depth, cv2.COLORMAP_INFERNO)
210
-
211
- filename = os.path.basename(filename)
212
-
213
- if pred_only:
214
- cv2.imwrite(os.path.join(outdir, filename[:filename.rfind('.')] + '_depth.png'), depth)
215
- else:
216
- split_region = np.ones((raw_image.shape[0], margin_width, 3), dtype=np.uint8) * 255
217
- combined_results = cv2.hconcat([raw_image, split_region, depth])
218
-
219
- caption_space = np.ones((caption_height, combined_results.shape[1], 3), dtype=np.uint8) * 255
220
- captions = ['Raw image', 'Depth Anything']
221
- segment_width = w + margin_width
222
-
223
- for i, caption in enumerate(captions):
224
- # Calculate text size
225
- text_size = cv2.getTextSize(caption, font, font_scale, font_thickness)[0]
226
-
227
- # Calculate x-coordinate to center the text
228
- text_x = int((segment_width * i) + (w - text_size[0]) / 2)
229
-
230
- # Add text caption
231
- cv2.putText(caption_space, caption, (text_x, 40), font, font_scale, (0, 0, 0), font_thickness)
232
-
233
- final_result = cv2.vconcat([caption_space, combined_results])
234
-
235
- cv2.imwrite(os.path.join(outdir, filename[:filename.rfind('.')] + '_img_depth.png'), final_result)
236
-
237
- def teed_imgs(img_path = './input', outdir = './output/teed_imgs',gaussianBlur = [0,3,0]):
238
- args, train_info = parse_args(is_testing=True, pl_opt_dir=outdir)
239
- os.makedirs('teed_tmp', exist_ok=True)
240
- if os.path.isfile(img_path):
241
- name, extension = os.path.splitext(img_path)
242
- if extension:
243
- if extension.lower() == 'txt':
244
- with open(img_path, 'r',encoding= 'utf-8') as f:
245
- filenames = f.read().splitlines()
246
- elif extension.lower() in ['.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp','tif']:
247
- filenames = [img_path]
248
- else:
249
- filenames = os.listdir(img_path)
250
- filenames = [os.path.join(img_path, filename) for filename in filenames if not filename.startswith('.') and filename.lower().endswith(('.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp','tif'))]
251
- filenames.sort()
252
- for filename in tqdm(filenames):
253
- if gaussianBlur[0] != 0:
254
- img = cv2.imread(filename)
255
- img = cv2.GaussianBlur(img, (gaussianBlur[1],gaussianBlur[1]), gaussianBlur[2])
256
- cv2.imwrite(os.path.join('teed_tmp', os.path.basename(filename)), img)
257
- else:
258
- copy_file(filename, 'teed_tmp')
259
- teed.main(args, train_info)
260
- shutil.rmtree('teed_tmp')
261
-
262
- def merge_2_images(img1, img2, mode, erosion_para = [[0,0],[0,0]], dilate = [0,0]): #将 img1 合并至 img2,调整大小与 img2 相同
263
- img1 = cv2.imread(img1)
264
- img2 = cv2.imread(img2)
265
- img1 = cv2.resize(img1, (img2.shape[1], img2.shape[0]))
266
- if erosion_para[0][1] != 0:
267
- img1 = erosion(img1, erosion_para[0][0], erosion_para[0][1], dilate[0])
268
- if erosion_para[1][1] != 0:
269
- img2 = erosion(img2, erosion_para[1][0], erosion_para[1][1], dilate[1])
270
- if mode == 'multiply':
271
- return multiply_blend(img1, img2)
272
- elif mode == 'screen':
273
- return screen_blend(img1, img2)
274
-
275
- def merge_images_in_2_folder(folder1, folder2, outdir, suffix_need_remove = None, suffix_floder = 0 , mode = 'multiply', erosion_para = [[0,0],[0,0]], dilate = [0,0]): #将 folder1 和 folder2 中的图片合并,可选是否移除某文件夹后缀,可选腐蚀参数[kernel_size,iterations]
276
- os.makedirs(outdir, exist_ok=True)
277
- name_extension_pairs_folder1 = [os.path.splitext(filename) for filename in os.listdir(folder1) if filename.endswith(('.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp','tif'))]
278
- filenames_noext_folder1, extensions_folder1 = zip(*name_extension_pairs_folder1)
279
- name_extension_pairs_folder2 = [os.path.splitext(filename) for filename in os.listdir(folder2) if filename.endswith(('.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp','tif'))]
280
- filenames_noext_folder2, extensions_folder2 = zip(*name_extension_pairs_folder2)
281
- if suffix_need_remove:
282
- if suffix_floder == 0:
283
- filenames_raw = list(filenames_noext_folder1).copy()
284
- filenames_noext_folder1 = [filename[:-len(suffix_need_remove)] + filename[-len(suffix_need_remove):].replace(suffix_need_remove, '') for filename in filenames_noext_folder1]
285
- if suffix_floder == 1:
286
- filenames_raw = list(filenames_noext_folder2).copy()
287
- filenames_noext_folder2 = [filename[:-len(suffix_need_remove)] + filename[-len(suffix_need_remove):].replace(suffix_need_remove, '') for filename in filenames_noext_folder2]
288
-
289
- for index, filename in enumerate(filenames_noext_folder1):
290
- if filename in filenames_noext_folder2:
291
- print(filename)
292
- if suffix_need_remove:
293
- if suffix_floder == 0:
294
- img1 = os.path.join(folder1, filenames_raw[index] + extensions_folder1[index])
295
- img2 = os.path.join(folder2, filename + extensions_folder2[filenames_noext_folder2.index(filename)])
296
- if suffix_floder == 1:
297
- img1 = os.path.join(folder1, filename + extensions_folder1[index])
298
- img2 = os.path.join(folder2, filenames_raw[filenames_noext_folder2.index(filename)] + extensions_folder2[filenames_noext_folder2.index(filename)])
299
- else:
300
- img1 = os.path.join(folder1, filename + extensions_folder1[index])
301
- img2 = os.path.join(folder2, filename + extensions_folder2[filenames_noext_folder2.index(filename)])
302
- result = merge_2_images(img1, img2, mode, erosion_para, dilate)
303
- cv2.imwrite(os.path.join(outdir, filename + extensions_folder1[index]), result)
304
-
305
- def process_line(img_path = './input', outdir = './output'):
306
- depth_anything(img_path, os.path.join(outdir,"depth_anything"))
307
- teed_imgs(img_path, os.path.join(outdir,"teed_imgs"), [1,7,2])
308
- teed_imgs(os.path.join(outdir,"depth_anything"), os.path.join(outdir,"dp_teed_imgs"), [0,7,2])
309
- merge_images_in_2_folder(os.path.join(outdir,"teed_imgs"), os.path.join(outdir,"dp_teed_imgs"), os.path.join(outdir,"merged_imgs"),'_depth', 1, 'multiply', [[2,0],[2,1]],[1,0])
310
-
311
- def invert_image(image):
312
- # 将图片从BGR转为灰度图
313
- gray_image = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
314
- # 对灰度图进行反转
315
- inverted_image = cv2.bitwise_not(gray_image)
316
- # 将反转后的灰度图转换回BGR格式
317
- inverted_image_bgr = cv2.cvtColor(inverted_image, cv2.COLOR_GRAY2BGR)
318
- return inverted_image_bgr
319
-
320
- def process_images(input_folder='./output/merged_imgs'):
321
- output_folder = os.path.join(os.path.dirname(input_folder), 'output_invert')
322
- os.makedirs(output_folder, exist_ok=True)
323
-
324
- # 获取输入文件夹中的所有图片文件
325
- image_files = [f for f in os.listdir(input_folder) if
326
- f.lower().endswith(('.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp', 'tif'))]
327
-
328
- for image_file in tqdm(image_files):
329
- image_path = os.path.join(input_folder, image_file)
330
- try:
331
- # 使用PIL库读取图像
332
- with Image.open(image_path) as img:
333
- image = np.array(img.convert('RGB'))[:, :, ::-1].copy()
334
- if image is not None:
335
- # 翻转图片
336
- inverted_image = invert_image(image)
337
- # 保存翻转后的图片到输出文件夹
338
- output_path = os.path.join(output_folder, image_file)
339
- cv2.imwrite(output_path, inverted_image)
340
- else:
341
- raise ValueError(f"Failed to read image: {image_file}")
342
- except Exception as e:
343
- print(f"Error processing file {image_file}: {e}")
344
-
345
- def process_line(img_path='./input', outdir='./output'):
346
- depth_anything(img_path, os.path.join(outdir, "depth_anything"))
347
- teed_imgs(img_path, os.path.join(outdir, "teed_imgs"), [1, 7, 2])
348
- teed_imgs(os.path.join(outdir, "depth_anything"), os.path.join(outdir, "dp_teed_imgs"), [0, 7, 2])
349
- merge_images_in_2_folder(os.path.join(outdir, "teed_imgs"), os.path.join(outdir, "dp_teed_imgs"), os.path.join(outdir, "merged_imgs"), '_depth', 1, 'multiply', [[2, 0], [2, 1]], [1, 0])
350
- process_images(os.path.join(outdir, "merged_imgs")) # 处理merged_imgs文件夹中的图片
351
-
352
-
353
-
354
-
355
- if __name__ == '__main__':
356
- # depth_anything()
357
- # teed_imgs('./input', './output/teed_imgs', [1,7,2])
358
- # teed_imgs('./output/depth_anything', './output/dp_teed_imgs', [0,7,2])
359
- # merge_images_in_2_folder('./output/teed_imgs', './output/dp_teed_imgs', './output/merged_imgs','_depth', 1, 'multiply', [[2,0],[2,1]],[1,0])
360
-
361
- # erosion_img_from_path('./output/teed_imgs', './output/erosion_imgs', 2, 1, True)
362
- # guassian_blur_path('./input', './output/guassian_blur', 7, 2)
363
- # erosion_img_from_path('./output/merged_imgs', './output/erosion_merged_imgs', 2, 1, False)
364
- # erosion_img_from_path('./output/erosion_merged_imgs', './output/erosion2_merged_imgs', 2, 1, True)
365
- img_path = "input"
366
- outdir = "output"
367
- process_line(img_path,outdir)