Stable-Diffusion-V1-5 数据管道构建使用Python处理训练数据集与生成结果如果你正在尝试用LoRA这类方法微调Stable Diffusion模型那你肯定遇到过这两个让人头疼的问题一是手头有一堆图片尺寸不一、质量参差怎么把它们变成模型能“消化”的训练数据二是模型跑起来后生成了成百上千张图一张张看、一张张挑眼睛都快花了有没有什么省力的办法这两个问题本质上都是数据管理的问题。前者是数据“预处理”后者是数据“后处理”。今天我们就来聊聊如何用Python搭建一套自动化数据管道把这两个环节都管起来让你能把精力更多地花在创意和调参上而不是繁琐的重复劳动上。1. 为什么你需要一个数据管道在开始动手写代码之前我们先得想明白为什么费这个劲去搭一个管道直接手动处理不行吗想象一下你收集了500张精美的动漫角色图打算训练一个专属的LoRA模型。手动操作的话你需要一张张打开图片检查尺寸如果不是512x512就用修图软件裁剪或缩放。为每张图片想一个合适的描述标签Prompt并记录在文本文件里。把所有处理好的图片和标签文件按照特定格式比如每张图片配一个同名的.txt文件整理好。这还没完当模型开始训练并生成测试图后你生成了1000张不同参数下的结果图。你需要一张张浏览凭感觉把“好图”、“一般图”、“废图”分开。你可能会想“这张图是用了那个提示词生成的来着”然后在一堆日志文件里翻找。这个过程不仅极其耗时而且容易出错更谈不上可重复性。今天处理50张图和500张图工作量是天壤之别。一个自动化的数据管道就是为了解决这些痛点而生的。它能把上述所有步骤标准化、流程化你只需要准备好原始素材运行脚本剩下的脏活累活都交给代码。无论是处理10张还是10000张图对你来说只是多等几分钟而已。2. 训练数据预处理从杂乱图片到标准数据集我们的目标是输入一个装满原始图片的文件夹输出一个符合Stable Diffusion训练要求例如每张图片有对应的描述文本文件的规整数据集。这个过程主要包含三个核心步骤图片标准化、自动打标可选和文件整理。2.1 环境准备与核心工具首先确保你的Python环境里安装了必要的库。我们将主要依赖PILPillow来处理图像os和shutil来操作文件。pip install Pillow如果你的图片来自网络可能还需要requests来下载。对于自动打标我们可以使用一些轻量级的图像识别或CLIP模型这里为了简化我们先以手动/半自动打标为例但会给出扩展思路。2.2 构建图片标准化处理脚本图片标准化通常指统一尺寸和裁剪。Stable Diffusion V1-5常用512x512的分辨率进行训练。我们的脚本需要能智能地处理不同长宽比的图片。一种常见的策略是“中心裁剪”或“缩放后填充”。这里我们实现一个更通用的方法先将短边缩放到目标尺寸然后从中心裁剪出正方形。创建一个名为preprocess_images.py的文件import os from PIL import Image from pathlib import Path def process_image(image_path, output_size512, output_dir./processed): 处理单张图片缩放并中心裁剪至指定正方形尺寸。 参数: image_path: 输入图片路径 output_size: 输出图片的边长默认512 output_dir: 处理后的图片保存目录 # 创建输出目录 Path(output_dir).mkdir(parentsTrue, exist_okTrue) try: with Image.open(image_path) as img: img img.convert(RGB) # 确保为RGB模式避免Alpha通道问题 # 计算缩放比例使短边等于output_size width, height img.size if width height: new_width output_size new_height int(height * (output_size / width)) else: new_height output_size new_width int(width * (output_size / height)) # 等比缩放 img_resized img.resize((new_width, new_height), Image.Resampling.LANCZOS) # 中心裁剪 left (new_width - output_size) / 2 top (new_height - output_size) / 2 right (new_width output_size) / 2 bottom (new_height output_size) / 2 img_cropped img_resized.crop((left, top, right, bottom)) # 保存图片 output_path Path(output_dir) / Path(image_path).name img_cropped.save(output_path, quality95) # 保存质量设为95平衡质量和大小 print(f已处理: {image_path} - {output_path}) return output_path except Exception as e: print(f处理图片 {image_path} 时出错: {e}) return None def batch_process(input_dir, output_size512): 批量处理一个目录下的所有图片。 参数: input_dir: 原始图片所在目录 output_size: 输出图片尺寸 input_dir Path(input_dir) output_dir input_dir.parent / f{input_dir.name}_processed supported_formats (.png, .jpg, .jpeg, .bmp, .webp) image_files [f for f in input_dir.rglob(*) if f.suffix.lower() in supported_formats] print(f找到 {len(image_files)} 张待处理图片。) for img_path in image_files: process_image(img_path, output_size, output_dir) print(f\n所有图片处理完成已保存至: {output_dir}) if __name__ __main__: # 使用示例将 ./raw_images 文件夹内的图片处理为512x512 batch_process(./raw_images)这个脚本会读取raw_images文件夹里的所有图片将它们统一处理成512x512的中心裁剪版本并保存到raw_images_processed文件夹中。2.3 为图片添加描述标签处理完尺寸接下来是关键的一步为每张图片添加文本描述。这是训练LoRA等模型时指导模型学习内容的核心。我们可以采用几种策略手动编写质量最高但最耗时。可以创建一个CSV文件或使用脚本辅助输入。自动生成基础使用现有的图像描述模型如BLIP、CLIP Interrogator等生成初步标签然后人工修正。文件名即标签如果图片文件名已经包含了描述性信息例如a_beautiful_landscape_sunset.jpg可以将其作为标签。这里我们实现一个简单的交互式脚本辅助你快速为一批图片打标并自动生成对应的.txt文件。创建一个名为generate_captions.py的文件import os from pathlib import Path def generate_captions_from_filenames(image_dir): 将图片文件名去除后缀和下划线作为初始标签。 例如a_beautiful_landscape.jpg - a beautiful landscape image_dir Path(image_dir) supported_formats (.png, .jpg, .jpeg, .bmp, .webp) for img_path in image_dir.rglob(*): if img_path.suffix.lower() in supported_formats: # 从文件名生成标签去除后缀将下划线替换为空格 caption img_path.stem.replace(_, ) # 创建同名的.txt文件并写入标签 caption_file img_path.with_suffix(.txt) with open(caption_file, w, encodingutf-8) as f: f.write(caption) print(f已为 {img_path.name} 生成标签: {caption}) def manual_caption_assistant(image_dir, caption_filecaptions.csv): 辅助手动打标生成一个CSV文件列出所有图片供你离线填写描述。 image_dir Path(image_dir) supported_formats (.png, .jpg, .jpeg, .bmp, .webp) image_list [f for f in image_dir.rglob(*) if f.suffix.lower() in supported_formats] # 生成CSV文件包含图片路径和空的描述列 import csv with open(caption_file, w, newline, encodingutf-8) as csvfile: writer csv.writer(csvfile) writer.writerow([image_path, caption]) for img_path in image_list: # 使用相对路径方便移植 rel_path img_path.relative_to(image_dir) writer.writerow([str(rel_path), ]) print(f已生成打标辅助文件: {caption_file}) print(f请用Excel或文本编辑器打开此文件在 caption 列填写每张图片的描述。) print(f填写完成后可运行另一个脚本将CSV中的描述写入对应的.txt文件。) if __name__ __main__: # 方法一使用文件名自动生成简单标签 # generate_captions_from_filenames(./raw_images_processed) # 方法二生成CSV文件供手动填写详细描述 manual_caption_assistant(./raw_images_processed)进阶提示对于自动生成高质量标签你可以集成transformers库中的BLIP模型。这需要更多的计算资源但能极大提升效率。思路是加载预训练的BLIP模型对每张图片生成描述然后你可以将其作为初稿进行修改。2.4 整理成最终训练格式经过前两步我们有了标准尺寸的图片和对应的文本描述文件。通常Stable Diffusion训练要求每个图片文件旁边都有一个同名的文本文件例如image01.jpg和image01.txt。如果你使用了上面的manual_caption_assistant并填写了CSV现在需要将CSV中的描述写回单独的.txt文件。这里提供一个转换脚本import csv from pathlib import Path def csv_to_txt(csv_path, image_base_dir): 将填写好的CSV文件中的描述写入对应图片的.txt文件。 image_base_dir Path(image_base_dir) with open(csv_path, r, encodingutf-8) as csvfile: reader csv.DictReader(csvfile) for row in reader: img_rel_path Path(row[image_path]) caption row[caption].strip() if caption: # 只处理有描述的条目 img_full_path image_base_dir / img_rel_path txt_file_path img_full_path.with_suffix(.txt) # 确保图片存在 if img_full_path.exists(): with open(txt_file_path, w, encodingutf-8) as f: f.write(caption) print(f已写入: {txt_file_path}) else: print(f警告图片不存在 {img_full_path}) if __name__ __main__: csv_to_txt(captions_filled.csv, ./raw_images_processed)至此一个规整的训练数据集就准备好了。你可以将这个_processed文件夹直接用于像kohya_ss这样的LoRA训练工具。3. 生成结果后处理从海量输出中筛选与管理模型训练好了或者你在尝试不同的提示词和参数一下子生成了几百张图。如何高效地管理、筛选和分类这些结果手动操作是不可持续的。3.1 构建结果图自动整理脚本我们假设生成的结果图都放在一个文件夹里文件名可能包含了生成参数例如seed_1234_steps_30_cfg_7.5.jpg。我们的目标是自动分类根据某些规则如使用的模型名称、主要提示词关键词将图片移动到不同的子文件夹。信息提取从文件名或同名的文本文件中解析生成参数。快速预览与评分提供一个简单界面让你能快速浏览并给图片打分。首先实现一个基于文件名的简单分类器。例如把所有包含“portrait”提示词的图放到“portrait”文件夹。import shutil from pathlib import Path import re def organize_by_keyword(source_dir, keyword_mappings): 根据关键词将图片整理到不同文件夹。 参数: source_dir: 结果图所在目录 keyword_mappings: 字典{‘关键词’: ‘目标文件夹名’} source_dir Path(source_dir) supported_formats (.png, .jpg, .jpeg, .bmp, .webp) for img_path in source_dir.rglob(*): if img_path.suffix.lower() in supported_formats: # 尝试从同名的.txt文件读取提示词 txt_file img_path.with_suffix(.txt) prompt if txt_file.exists(): with open(txt_file, r, encodingutf-8) as f: prompt f.read().lower() # 也检查文件名本身 file_stem img_path.stem.lower() combined_text prompt file_stem moved False for keyword, target_folder in keyword_mappings.items(): if keyword.lower() in combined_text: target_dir source_dir / target_folder target_dir.mkdir(exist_okTrue) shutil.move(str(img_path), str(target_dir / img_path.name)) # 同时移动对应的.txt文件 if txt_file.exists(): shutil.move(str(txt_file), str(target_dir / txt_file.name)) print(f已移动 {img_path.name} - {target_folder}/) moved True break if not moved: # 未匹配任何关键词的放入“其他”文件夹 other_dir source_dir / 其他 other_dir.mkdir(exist_okTrue) shutil.move(str(img_path), str(other_dir / img_path.name)) if txt_file.exists(): shutil.move(str(txt_file), str(other_dir / txt_file.name)) print(f已移动 {img_path.name} - 其他/) if __name__ __main__: # 定义你的关键词映射 my_keywords { portrait: 人像, landscape: 风景, cat: 猫, cyberpunk: 赛博朋克, } organize_by_keyword(./generated_outputs, my_keywords)3.2 实现一个简单的图片评分与筛选工具分类之后我们还需要在同类图片中挑选出最好的。我们可以写一个脚本顺序展示图片并记录你的评分。import tkinter as tk from tkinter import ttk from PIL import Image, ImageTk from pathlib import Path import json class ImageRater: def __init__(self, image_folder): self.image_folder Path(image_folder) self.images list(self.image_folder.glob(*.jpg)) list(self.image_folder.glob(*.png)) self.current_index 0 self.ratings {} # 存储评分 {‘文件名’: 分数} # 尝试加载已保存的评分 self.rating_file self.image_folder / ratings.json if self.rating_file.exists(): with open(self.rating_file, r) as f: self.ratings json.load(f) self.setup_ui() self.load_image() def setup_ui(self): self.root tk.Tk() self.root.title(图片评分器) # 图片显示区域 self.image_label ttk.Label(self.root) self.image_label.pack() # 提示词显示 self.prompt_label ttk.Label(self.root, text, wraplength600) self.prompt_label.pack() # 评分按钮 btn_frame ttk.Frame(self.root) btn_frame.pack(pady10) for score in [1, 2, 3, 4, 5]: btn ttk.Button(btn_frame, textf{score}分, commandlambda sscore: self.rate(s)) btn.pack(sidetk.LEFT, padx5) # 跳过按钮 ttk.Button(btn_frame, text跳过, commandself.next_image).pack(sidetk.LEFT, padx5) # 进度显示 self.progress_label ttk.Label(self.root, text) self.progress_label.pack() self.root.bind(Right, lambda e: self.next_image()) self.root.bind(Left, lambda e: self.prev_image()) for i in range(1,6): self.root.bind(str(i), lambda e, si: self.rate(s)) def load_image(self): if self.current_index len(self.images): self.save_and_exit() return img_path self.images[self.current_index] # 加载图片 img Image.open(img_path) img.thumbnail((800, 800)) # 缩放到适合窗口的大小 photo ImageTk.PhotoImage(img) self.image_label.config(imagephoto) self.image_label.image photo # 保持引用 # 加载提示词 txt_file img_path.with_suffix(.txt) prompt if txt_file.exists(): with open(txt_file, r, encodingutf-8) as f: prompt f.read() self.prompt_label.config(textf提示词: {prompt}) # 更新进度 self.progress_label.config(textf进度: {self.current_index 1}/{len(self.images)}) def rate(self, score): img_name self.images[self.current_index].name self.ratings[img_name] score print(f为 {img_name} 评分: {score}) self.next_image() def next_image(self): self.current_index 1 if self.current_index len(self.images): self.load_image() else: self.save_and_exit() def prev_image(self): if self.current_index 0: self.current_index - 1 self.load_image() def save_and_exit(self): # 保存评分到文件 with open(self.rating_file, w) as f: json.dump(self.ratings, f, indent2) print(f评分已保存至 {self.rating_file}) self.root.quit() self.root.destroy() def run(self): self.root.mainloop() if __name__ __main__: # 指定需要评分的图片文件夹 rater ImageRater(./generated_outputs/人像) rater.run()运行这个脚本它会弹出一个窗口依次展示“人像”文件夹里的图片和对应的提示词。你可以按1-5数字键评分或按方向键切换图片。所有评分会自动保存到ratings.json文件里。之后你就可以根据评分轻松筛选出高分作品了。4. 总结搭建数据管道听起来有点工程化但一旦搭建好它为你节省的时间和带来的便利是巨大的。对于训练数据预处理我们通过标准化图片尺寸和半自动化打标将杂乱无章的原始素材变成了模型“爱吃”的标准餐。对于生成结果后处理我们通过自动分类和交互式评分将你从海量图片中手动筛选的苦差事中解放出来。这套流程是可扩展的。你可以根据需求加入更复杂的自动打标模型如BLIP、基于图像特征的相似度去重、或者将评分结果直接反馈到提示词优化循环中。核心思想是将重复性工作自动化让你能聚焦在创造性和决策性的事情上。开始可能会花点时间编写和调试脚本但这是一劳永逸的投资。下次当你再面对成千上万的图片时你会庆幸自己已经准备好了这套“流水线”。获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。