SUPER COLORIZER数据预处理教程：构建高质量图像训练集

张

张建站

2026/7/13 3:33:36

10分钟阅读

SUPER COLORIZER数据预处理教程构建高质量图像训练集想自己训练一个给线稿上色的模型结果第一步就被数据难住了网上找的图片尺寸五花八门有的带水印有的背景杂乱直接丢给模型训练效果肯定好不了。别急这篇文章就是来解决这个问题的。我会手把手带你走一遍SUPER COLORIZER模型的数据预处理全流程。你不用是数据科学专家只要会一点基础的Python就能跟着我把一堆原始图片变成模型能“消化”的高质量训练集。整个过程就像给食材做预处理洗菜、切配、腌制每一步都做好了最后炒出来的菜才香。我们的目标很明确从零开始构建一个格式规范、质量上乘的线稿-彩色图对数据集为后续的模型微调打下坚实基础。1. 准备工作理清思路与备好工具在开始动手处理图片之前我们得先想清楚要做什么以及需要哪些工具。盲目开始很容易中途卡壳。1.1 理解数据需求模型要“吃”什么SUPER COLORIZER这类图像上色模型训练时通常需要“成对”的数据。简单说就是一张线稿图输入和它对应的彩色完成图目标输出。模型的任务是学会从线稿推断出合理的颜色。因此我们最终要准备的数据集应该是许多这样的“线稿-彩图”对。每一对图片中的两张图在内容上必须严格对齐也就是说线稿的线条和彩图的轮廓得对得上。1.2 搭建你的工作环境你需要准备两样东西一个能运行Python的环境和存放图片的文件夹。Python环境我推荐使用Anaconda来管理环境这样库的版本不会冲突。如果你还没安装去Anaconda官网下载一个安装就行。然后我们创建一个专门用于这个项目的新环境conda create -n colorizer-data python3.9 conda activate colorizer-data安装必要的库处理图片我们主要用PILPillow库和OpenCV。在激活的环境里用pip安装它们pip install pillow opencv-python numpy组织文件夹在你电脑上找个地方新建一个项目文件夹比如叫super_colorizer_data。在里面再创建几个子文件夹先搭好架子super_colorizer_data/ ├── raw_images/ # 第一步存放你从各处收集来的原始图片 ├── processed/ # 第二步存放裁剪、标准化后的图片 ├── line_art/ # 第三步存放生成的线稿图 ├── color/ # 第三步存放对应的彩色图或原始彩图 └── final_pairs/ # 第四步存放最终整理好的配对数据集文件夹先建好思路就清晰了一半。2. 第一步原始图像收集与“避坑”指南数据质量的天花板在收集阶段就定下了。这里有些实操方法和重要注意事项。2.1 去哪里找图片开源数据集这是最省心、版权最清晰的选择。比如COCO、ADE20K等大型通用数据集或者一些动漫、艺术风格的专门数据集。自己拍摄或绘制如果你有特定需求这是最好的方式版权完全自主。版权明确的图库一些注明可免费用于商业或研究用途的网站如Unsplash, Pixabay。务必仔细阅读每张图片的授权协议。2.2 必须牢记的版权注意事项这是严肃的一环处理不好可能会有法律风险。明确授权只使用明确标注为“公共领域”、“知识共享CC协议”或允许用于“非商业/商业研究”的图片。不要想当然。保留记录最好用一个简单的表格记录每张图片的来源、授权类型和作者信息。养成好习惯。警惕“搬运”内容从社交媒体、壁纸网站随意下载的图片绝大多数都不具备清晰的二次授权风险很高不建议使用。内部使用如果你的研究项目完全在内部进行不公开发布数据集或模型风险相对较低但仍建议使用合规素材。2.3 初步筛选原则收集时就可以做第一次筛选内容尽量选择主体清晰、背景不太复杂的图片。过于杂乱或抽象的图片模型学习起来会更困难。格式优先选择.png或.jpg格式。.png通常无损.jpg更常见。尺寸暂时不要求统一但尽量不要用手机拍的几十兆超大图也不要用尺寸过小如宽度小于256像素的图后续处理会麻烦。把收集来的所有原始图片都先扔进我们之前建好的raw_images文件夹里。3. 第二步批量裁剪与分辨率标准化现在raw_images里的图片可能大小不一、比例各异。我们需要把它们变成统一的“规格”。3.1 为什么要统一尺寸神经网络通常要求输入尺寸是固定的。统一的尺寸能保证模型训练时批次处理batch processing的高效性。所有图像数据在数值分布上处于相近的范围有利于模型稳定学习。对于图像上色任务常见的输入尺寸是256x256、512x512或更大的方形图。我们这里以512x512为例。3.2 编写批量处理脚本我们写一个Python脚本一次性处理raw_images里的所有图片。核心思路是中心裁剪或缩放后填充以保证内容不变形。创建一个名为01_crop_and_resize.py的文件放在项目根目录下。import os from PIL import Image import argparse def process_image(image_path, output_dir, target_size(512, 512), modecrop): 处理单张图片裁剪或缩放至目标尺寸。参数: image_path: 输入图片路径 output_dir: 输出目录 target_size: 目标尺寸 (宽高) mode: 处理模式crop为中心裁剪pad为缩放后填充 try: img Image.open(image_path) img img.convert(RGB) # 确保是RGB三通道去掉Alpha通道 original_width, original_height img.size target_width, target_height target_size if mode crop: # 中心裁剪计算裁剪区域 left (original_width - target_width) / 2 top (original_height - target_height) / 2 right (original_width target_width) / 2 bottom (original_height target_height) / 2 # PIL的crop需要整数坐标 img_cropped img.crop((int(left), int(top), int(right), int(bottom))) img_resized img_cropped.resize(target_size, Image.Resampling.LANCZOS) elif mode pad: # 计算缩放比例使长边匹配目标尺寸短边按比例缩放 ratio min(target_width / original_width, target_height / original_height) new_size (int(original_width * ratio), int(original_height * ratio)) img_resized img.resize(new_size, Image.Resampling.LANCZOS) # 创建新画布将缩放后的图粘贴到中心 new_img Image.new(RGB, target_size, (255, 255, 255)) # 白色填充 paste_position ((target_width - new_size[0]) // 2, (target_height - new_size[1]) // 2) new_img.paste(img_resized, paste_position) img_resized new_img else: print(f未知模式: {mode}) return # 保存图片 base_name os.path.basename(image_path) name, ext os.path.splitext(base_name) # 统一保存为.jpg格式也可改为.png output_path os.path.join(output_dir, f{name}_processed.jpg) img_resized.save(output_path, JPEG, quality95) print(f已处理: {base_name} - {output_path}) except Exception as e: print(f处理图片 {image_path} 时出错: {e}) def batch_process(input_dir, output_dir, target_size(512, 512), modecrop): 批量处理目录下所有图片 if not os.path.exists(output_dir): os.makedirs(output_dir) supported_formats (.jpg, .jpeg, .png, .bmp, .tiff) for filename in os.listdir(input_dir): if filename.lower().endswith(supported_formats): image_path os.path.join(input_dir, filename) process_image(image_path, output_dir, target_size, mode) if __name__ __main__: parser argparse.ArgumentParser(description批量裁剪/缩放图片) parser.add_argument(--input, typestr, default./raw_images, help原始图片目录) parser.add_argument(--output, typestr, default./processed, help处理后的图片目录) parser.add_argument(--size, typeint, default512, help目标尺寸正方形边长) parser.add_argument(--mode, typestr, defaultcrop, choices[crop, pad], help处理模式crop(裁剪) 或 pad(填充)) args parser.parse_args() print(f开始批量处理... 模式: {args.mode}, 目标尺寸: {args.size}x{args.size}) batch_process(args.input, args.output, (args.size, args.size), args.mode) print(批量处理完成)如何使用这个脚本打开终端或Anaconda Prompt切换到你的项目目录然后运行# 使用默认的中心裁剪模式输出到processed文件夹 python 01_crop_and_resize.py # 或者如果你想用填充模式并指定尺寸为256 python 01_crop_and_resize.py --mode pad --size 256运行后所有处理好的、尺寸统一的图片就会出现在processed文件夹里。建议你先用crop模式试试因为它能保证图片内容不变形只是可能裁掉一些边缘。对于主体在中心的图片效果很好。4. 第三步生成线稿-彩色图对这是最关键的一步。我们需要从一张彩色图生成它对应的线稿。这里介绍两种主流方法。4.1 方法一使用OpenCV进行边缘检测这是一种比较通用、快速的方法适合多种风格的图片。我们写第二个脚本02_generate_lineart_cv.py。import os import cv2 import numpy as np from PIL import Image import argparse def generate_lineart_opencv(color_img_path, output_line_dir, output_color_dir): 使用OpenCV生成线稿。步骤灰度化 - 高斯模糊去噪 - Canny边缘检测 - 反相黑底白线 try: # 读取彩色图片 color_img cv2.imread(color_img_path) if color_img is None: print(f无法读取图片: {color_img_path}) return # 1. 转为灰度图 gray cv2.cvtColor(color_img, cv2.COLOR_BGR2GRAY) # 2. 使用高斯模糊减少噪声让边缘更干净 blurred cv2.GaussianBlur(gray, (5, 5), 0) # 3. Canny边缘检测 # 这两个阈值很重要可以调整。比例通常为2:1或3:1 edges cv2.Canny(blurred, threshold150, threshold2150) # 4. 反相将黑底白边变为白底黑边更符合常见线稿 lineart 255 - edges # 保存线稿图 base_name os.path.basename(color_img_path) name, _ os.path.splitext(base_name) lineart_path os.path.join(output_line_dir, f{name}_line.jpg) color_save_path os.path.join(output_color_dir, f{name}_color.jpg) cv2.imwrite(lineart_path, lineart) # 也保存一份彩色图到color目录确保文件名对应 cv2.imwrite(color_save_path, color_img) print(f已生成: {base_name} - 线稿和彩图对) except Exception as e: print(f处理图片 {color_img_path} 时出错: {e}) def batch_generate_cv(input_dir, line_dir, color_dir): 批量生成线稿 if not os.path.exists(line_dir): os.makedirs(line_dir) if not os.path.exists(color_dir): os.makedirs(color_dir) supported_formats (.jpg, .jpeg, .png) for filename in os.listdir(input_dir): if filename.lower().endswith(supported_formats): img_path os.path.join(input_dir, filename) generate_lineart_opencv(img_path, line_dir, color_dir) if __name__ __main__: parser argparse.ArgumentParser(description使用OpenCV批量生成线稿) parser.add_argument(--input, typestr, default./processed, help处理后的彩色图片目录) parser.add_argument(--line_dir, typestr, default./line_art, help线稿输出目录) parser.add_argument(--color_dir, typestr, default./color, help彩色图输出目录与线稿配对) args parser.parse_args() print(开始使用OpenCV生成线稿...) batch_generate_cv(args.input, args.line_dir, args.color_dir) print(线稿生成完成)运行它python 02_generate_lineart_cv.py4.2 方法二使用XDoG算法更艺术化XDoGExtended Difference of Gaussians能产生更粗、更有绘画感的线条非常适合动漫、插画风格。我们需要先安装scikit-image库。pip install scikit-image然后创建脚本03_generate_lineart_xdog.py。import os import numpy as np from PIL import Image import argparse from skimage import color, filters def xdog_filter(image, gamma0.98, phi200, epsilon0.01, k1.6, sigma0.5): 应用XDoG滤镜生成艺术化线稿。参数可以调整以获得不同粗细和风格的线条。 # 转换为灰度 if len(image.shape) 3: gray color.rgb2gray(image) else: gray image # 应用两个不同标准差的高斯滤波 g1 filters.gaussian(gray, sigma) g2 filters.gaussian(gray, sigma * k) # 计算差值并应用阈值 d g1 - gamma * g2 d d / d.max() # 归一化 # 软阈值处理让线条过渡更自然 for i in range(d.shape[0]): for j in range(d.shape[1]): if d[i,j] epsilon: d[i,j] 1 else: # 这个公式能产生平滑的过渡 d[i,j] 1 np.tanh(phi * (d[i,j] - epsilon)) # 反相黑线白底 lineart 1 - d # 缩放到0-255范围 lineart (lineart * 255).astype(np.uint8) return lineart def generate_lineart_xdog(color_img_path, output_line_dir, output_color_dir): 使用XDoG生成线稿 try: # 用PIL打开转为numpy数组 color_img np.array(Image.open(color_img_path).convert(RGB)) lineart xdog_filter(color_img) # 保存 base_name os.path.basename(color_img_path) name, _ os.path.splitext(base_name) lineart_path os.path.join(output_line_dir, f{name}_line_xdog.jpg) color_save_path os.path.join(output_color_dir, f{name}_color.jpg) Image.fromarray(lineart).save(lineart_path) Image.fromarray(color_img).save(color_save_path) print(f已生成(XDoG): {base_name}) except Exception as e: print(f处理图片 {color_img_path} 时出错: {e}) def batch_generate_xdog(input_dir, line_dir, color_dir): if not os.path.exists(line_dir): os.makedirs(line_dir) if not os.path.exists(color_dir): os.makedirs(color_dir) supported_formats (.jpg, .jpeg, .png) for filename in os.listdir(input_dir): if filename.lower().endswith(supported_formats): img_path os.path.join(input_dir, filename) generate_lineart_xdog(img_path, line_dir, color_dir) if __name__ __main__: parser argparse.ArgumentParser(description使用XDoG批量生成艺术化线稿) parser.add_argument(--input, typestr, default./processed, help处理后的彩色图片目录) parser.add_argument(--line_dir, typestr, default./line_art_xdog, helpXDoG线稿输出目录) parser.add_argument(--color_dir, typestr, default./color_xdog, help彩色图输出目录) args parser.parse_args() print(开始使用XDoG生成艺术化线稿...) batch_generate_xdog(args.input, args.line_dir, args.color_dir) print(XDoG线稿生成完成)运行它python 03_generate_lineart_xdog.py两种方法怎么选OpenCV Canny线条清晰、精准适合写实风格、建筑、物体等。速度快。XDoG线条有笔触感粗细有变化更像手绘非常适合动漫、卡通、艺术插画风格。速度稍慢。你可以两种都试试看看哪种生成的线稿更符合你数据集的风格。生成的线稿和对应的彩色图现在分别存放在line_art和color或line_art_xdog和color_xdog文件夹里并且文件名是一一对应的这很重要。5. 第四步数据增强与最终整理有了配对数据我们还可以通过一些简单的变换让数据集变得更“大”、更“鲁棒”。5.1 简单的数据增强数据增强是指在不改变图片语义的前提下对其进行一些随机变换如翻转、旋转、微调颜色从而创造出“新”的训练样本。关键点是必须对线稿和彩图进行完全相同的变换否则配对关系就被破坏了。我们写一个简单的增强脚本作为示例04_augment_pairs.py这里只实现水平翻转import os from PIL import Image import random import argparse def augment_pair(line_path, color_path, output_dir, augment_name): 对一对线稿和彩图进行相同的增强操作 try: line_img Image.open(line_path) color_img Image.open(color_path) # 示例随机水平翻转 (50%概率) if random.random() 0.5: line_img line_img.transpose(Image.FLIP_LEFT_RIGHT) color_img color_img.transpose(Image.FLIP_LEFT_RIGHT) aug_suffix _flip else: aug_suffix _orig # 可以在这里添加其他增强如小角度随机旋转、亮度微调等 # 但切记所有操作必须同时应用于line_img和color_img base_name os.path.basename(line_path).replace(_line, ) name, ext os.path.splitext(base_name) new_line_name f{name}_line{aug_suffix}{ext} new_color_name f{name}_color{aug_suffix}{ext} line_img.save(os.path.join(output_dir, line, new_line_name)) color_img.save(os.path.join(output_dir, color, new_color_name)) except Exception as e: print(f增强配对 {line_path} 时出错: {e}) def batch_augment(line_dir, color_dir, output_dir): 批量增强所有配对 line_output os.path.join(output_dir, line) color_output os.path.join(output_dir, color) os.makedirs(line_output, exist_okTrue) os.makedirs(color_output, exist_okTrue) # 首先把原始未增强的复制一份过去 for f in os.listdir(line_dir): if f.endswith((.jpg, .png)): src os.path.join(line_dir, f) dst os.path.join(line_output, f) Image.open(src).save(dst) for f in os.listdir(color_dir): if f.endswith((.jpg, .png)): src os.path.join(color_dir, f) dst os.path.join(color_output, f) Image.open(src).save(dst) # 然后对每一对进行增强生成新文件 line_files [f for f in os.listdir(line_dir) if f.endswith((.jpg, .png))] for line_file in line_files: # 根据命名规则找到对应的彩图文件 color_file line_file.replace(_line, _color) line_path os.path.join(line_dir, line_file) color_path os.path.join(color_dir, color_file) if os.path.exists(color_path): augment_pair(line_path, color_path, output_dir, aug1) else: print(f未找到配对彩图: {color_file}) print(f数据增强完成。原始增强数据已保存至: {output_dir}) if __name__ __main__: parser argparse.ArgumentParser(description对线稿-彩图对进行数据增强) parser.add_argument(--line_dir, typestr, default./line_art, help线稿目录) parser.add_argument(--color_dir, typestr, default./color, help彩色图目录) parser.add_argument(--output, typestr, default./final_pairs, help最终输出目录) args parser.parse_args() batch_augment(args.line_dir, args.color_dir, args.output)运行它将增强后的数据放到final_pairs文件夹python 04_augment_pairs.py5.2 组织最终数据集格式现在final_pairs文件夹里应该有两个子文件夹line和color里面存放着严格配对的图片。这是最常见的数据集组织形式之一。另一种模型训练时常用的格式是准备一个清单文件如.txt或.csv里面每一行记录一对图片的路径。我们可以用一个小脚本生成它import os final_dir ./final_pairs line_dir os.path.join(final_dir, line) color_dir os.path.join(final_dir, color) pair_list [] for f in os.listdir(line_dir): if f.endswith((.jpg, .png)): line_path os.path.join(line, f) color_file f.replace(_line, _color) color_path os.path.join(color, color_file) # 检查配对是否存在 if os.path.exists(os.path.join(final_dir, color_path)): pair_list.append(f{line_path},{color_path}) # 将清单写入文件 with open(os.path.join(final_dir, train_pairs.txt), w) as f: f.write(\n.join(pair_list)) print(f已生成训练配对清单共 {len(pair_list)} 对数据。)运行后你会得到一个train_pairs.txt文件内容像这样line/img1_line.jpg,color/img1_color.jpg line/img1_line_flip.jpg,color/img1_color_flip.jpg line/img2_line.jpg,color/img2_color.jpg ...这个文件在后续训练配置中会非常有用。6. 总结与下一步建议走完这一整套流程你应该已经拥有了一个属于自己的、规整的线稿上色训练数据集。回顾一下我们经历了收集、裁剪、生成线稿、增强、整理这几个关键步骤。每一步的脚本都尽量做到了简单明了你可以根据自己图片的特点去调整里面的参数比如裁剪模式、边缘检测的阈值、XDoG的线条粗细等。用下来感觉数据预处理确实是个需要耐心的活儿但它的重要性怎么强调都不为过。你喂给模型的数据质量直接决定了它能学到多好。现在数据集准备好了下一步就是选择具体的SUPER COLORIZER模型框架比如基于PyTorch或TensorFlow的实现然后按照其文档要求将我们准备好的final_pairs文件夹或者train_pairs.txt清单配置到训练脚本中就可以开始你的模型微调之旅了。遇到问题很正常多调整参数多看看生成的中间结果比如线稿是否清晰、配对是否准确慢慢你就能找到感觉。祝你训练顺利获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。