cv_unet_image-colorization多图批量处理扩展教程：Python脚本自动化上色

cv_unet_image-colorization多图批量处理扩展教程Python脚本自动化上色1. 引言从手动到自动解放你的生产力你是不是也遇到过这样的场景手头有一堆黑白老照片想用AI工具给它们上色但每次只能上传一张点一下按钮等一会儿再下载保存。一张两张还好要是有几十张、上百张这个重复的过程就变得非常枯燥和耗时。今天我们就来解决这个问题。本文将带你深入cv_unet_image-colorization工具并教你如何编写一个Python脚本实现多图批量自动化上色。你将学会如何理解并调用工具背后的核心上色函数。编写一个脚本让它自动遍历文件夹里的所有黑白图片。让脚本批量处理完成后自动保存你只需要泡杯咖啡等着就行。无论你是想批量修复家庭相册还是需要处理大量历史影像资料这个自动化方案都能让你的效率提升十倍不止。我们不需要修改原始工具的Streamlit界面而是直接利用它的“引擎”用代码来驱动它。2. 核心原理找到批量处理的“钥匙”在开始写代码之前我们需要先搞清楚cv_unet_image-colorization这个工具是怎么工作的。它的Streamlit界面很友好但对我们批量处理来说关键不在于界面而在于界面背后那个真正执行上色任务的“核心函数”。通过分析工具代码通常是app.py或类似的主文件我们可以找到这个核心。它一般长这样# 这是工具内部的核心上色函数示例具体名称可能不同 def colorize_image(input_image): 核心上色函数 参数: input_image (PIL.Image 或 numpy数组) - 输入的黑白/灰度图像 返回: colored_image (PIL.Image) - 上色后的彩色图像 # 1. 初始化模型通常用st.cache_resource缓存了 model get_colorization_model() # 2. 预处理图像转换为模型需要的格式如Tensor processed_img preprocess(input_image) # 3. 模型推理UNet进行上色预测 output model(processed_img) # 4. 后处理将模型输出转换回PIL图像 result_img postprocess(output) return result_img批量处理的思路就变得非常清晰了剥离界面我们不需要Streamlit的st.file_uploader和st.button只需要导入这个colorize_image函数。文件遍历用Python的os和PIL库读取一个文件夹下的所有图片文件。循环调用对每一张图片调用colorize_image函数。保存结果将上色后的图片保存到另一个文件夹。接下来我们就一步步实现这个思路。3. 环境准备与项目结构在动手写脚本前请确保你的环境已经准备好。3.1 确认基础环境你的电脑或服务器上应该已经能正常运行原始的cv_unet_image-colorization的Streamlit应用。这意味着所有依赖都已经安装# 通常需要这些库如果你的环境已经能运行原工具则无需重复安装 # pip install modelscope opencv-python torch streamlit Pillow numpy关键点确保模型权重文件存在于正确的路径如/root/ai-models/iic/cv_unet_image-colorization这与原工具要求一致。3.2 规划项目目录为了让脚本运行得更清晰我建议你按以下方式组织文件夹your_project_folder/ ├── batch_colorizer.py # 我们将要编写的批量处理脚本 ├── input_images/ # 存放所有待上色的黑白图片 │ ├── photo1.jpg │ ├── photo2.png │ └── ... ├── output_images/ # 脚本运行后彩色图片将自动保存到这里 └── original_app.py # 原始的cv_unet_image-colorization应用文件用于导入函数把你要处理的所有黑白照片都丢进input_images文件夹。output_images文件夹可以空着脚本会自动创建。4. 编写批量上色Python脚本现在我们开始编写核心脚本batch_colorizer.py。我会逐部分解释并提供完整代码。4.1 导入必要的库首先我们需要导入一些“帮手”。import os import sys from pathlib import Path from PIL import Image import traceback # 将原始应用所在目录加入Python路径以便导入其中的函数 # 假设 original_app.py 在当前目录的上一级或者你知道它的具体位置 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))4.2 导入核心上色函数这是最关键的一步从原始工具中“借用”其上色引擎。try: # 尝试从原始应用文件中导入核心函数和模型 # 注意你需要根据原始文件的实际内容调整导入对象 from original_app import colorize_image, get_colorization_model # 或者如果模型初始化是全局的你可能需要这样 # from original_app import pipeline, model print(成功导入核心上色函数。) except ImportError as e: print(f导入核心函数失败: {e}) print(请确保‘original_app.py’文件存在且包含了colorize_image函数。) print(你可能需要根据原始代码调整导入语句。) sys.exit(1)重要提示original_app应替换为你的原始应用文件名如app。colorize_image和get_colorization_model也需要替换为原始代码中实际的函数名。你可能需要查看原代码来确定。4.3 核心批量处理函数我们来编写一个函数它负责处理单张图片读取、上色、保存。def process_single_image(input_path, output_path, model_pipeline): 处理单张图片上色并保存。 参数: input_path (str): 输入图片的完整路径。 output_path (str): 输出图片的完整路径。 model_pipeline: 已初始化的上色模型管道。 try: print(f正在处理: {os.path.basename(input_path)}) # 1. 使用PIL打开图片 with Image.open(input_path) as img: # 2. 确保图片是RGB或L灰度模式必要时转换 if img.mode not in (RGB, L): img img.convert(RGB) # 如果是彩色图但可能是黑白内容可先转为灰度再让模型处理这里假设输入就是黑白/灰度图 # 为了兼容我们传递PIL Image对象 # 3. 调用核心上色函数 # 注意这里需要根据你导入的函数实际调用方式调整 colored_img colorize_image(img) # 假设colorize_image接受PIL Image # 4. 保存结果 colored_img.save(output_path) print(f 处理完成已保存至: {os.path.basename(output_path)}) except Exception as e: print(f 处理图片 {input_path} 时出错: {e}) traceback.print_exc()4.4 主函数遍历文件夹控制流程主函数将串联所有步骤准备文件夹遍历图片调用处理函数。def main(input_dirinput_images, output_diroutput_images): 批量处理主函数。 参数: input_dir (str): 输入图片目录名。 output_dir (str): 输出图片目录名。 # 将相对路径转换为绝对路径 current_dir Path(__file__).parent input_dir_path current_dir / input_dir output_dir_path current_dir / output_dir # 1. 检查输入目录是否存在 if not input_dir_path.exists(): print(f错误输入目录 {input_dir} 不存在。) return # 2. 创建输出目录如果不存在 output_dir_path.mkdir(parentsTrue, exist_okTrue) # 3. 初始化模型全局一次避免每次处理都重复加载 print(正在初始化AI上色模型这可能需要几秒钟...) try: # 假设我们导入的get_colorization_model()返回初始化好的pipeline model_pipeline get_colorization_model() print(模型初始化成功) except Exception as e: print(f模型初始化失败: {e}) return # 4. 遍历输入目录中的所有图片文件 # 支持常见图片格式 supported_extensions (.jpg, .jpeg, .png, .bmp, .tiff, .webp) image_files [f for f in input_dir_path.iterdir() if f.suffix.lower() in supported_extensions] if not image_files: print(f在 {input_dir} 目录中没有找到支持的图片文件。) return print(f找到 {len(image_files)} 张待处理图片。) print(开始批量上色处理...) print(- * 50) # 5. 循环处理每张图片 processed_count 0 for img_path in image_files: # 构建输出文件路径保持原名也可以添加后缀如‘_colored’ output_filename f{img_path.stem}_colored{img_path.suffix} output_path output_dir_path / output_filename process_single_image(str(img_path), str(output_path), model_pipeline) processed_count 1 print(- * 50) print(f批量处理完成共处理 {processed_count} 张图片。) print(f彩色结果已保存至: {output_dir_path})4.5 脚本入口最后添加脚本的启动入口。if __name__ __main__: # 你可以在这里直接指定输入输出文件夹名 main(input_dirinput_images, output_diroutput_images)5. 运行与使用指南脚本写好了我们来试试怎么用它。5.1 第一步放置文件与图片将完整的batch_colorizer.py脚本保存到你的项目文件夹。确保原始的cv_unet_image-colorization应用文件比如app.py也在同一目录或者你知道如何正确导入其中的函数可能需要调整sys.path。把所有黑白图片放入input_images文件夹。5.2 第二步调整导入关键打开你原始的app.py文件找到执行上色操作的那个函数。它可能不叫colorize_image而是被包裹在Streamlit的按钮点击事件里。你需要将这个核心逻辑提取或封装成一个独立的函数以便我们的脚本调用。例如在app.py中添加或确认这样一个函数# 在app.py中确保有这个函数 def colorize_image(uploaded_image): # 这里是原始的上色逻辑通常包括 # 1. 将上传的文件转换为PIL Image或numpy数组 # 2. 调用model_pipeline进行预测 # 3. 返回上色后的PIL Image # ... 具体代码取决于原app的实现 return colored_pil_image然后在我们的batch_colorizer.py中确保from original_app import colorize_image这行代码能正确导入这个函数。5.3 第三步执行脚本打开终端命令行进入你的项目目录运行脚本python batch_colorizer.py你会看到类似这样的输出成功导入核心上色函数。 正在初始化AI上色模型这可能需要几秒钟... 模型初始化成功 找到 15 张待处理图片。 开始批量上色处理... -------------------------------------------------- 正在处理: photo1.jpg 处理完成已保存至: photo1_colored.jpg 正在处理: photo2.png 处理完成已保存至: photo2_colored.png ... -------------------------------------------------- 批量处理完成共处理 15 张图片。 彩色结果已保存至: /your_project_folder/output_images现在去output_images文件夹看看你的成果吧6. 脚本优化与进阶技巧基本的批量功能已经实现但我们可以让它更强大、更灵活。6.1 处理进度提示处理大量图片时有个进度条会更友好。我们可以使用tqdm库。# 首先安装pip install tqdm from tqdm import tqdm # 在主函数的循环处替换 print(开始批量上色处理...) for img_path in tqdm(image_files, desc上色进度): # ... 处理逻辑 ...6.2 错误处理与日志为了更稳定我们可以将错误信息记录到日志文件而不是仅仅打印出来。import logging logging.basicConfig(filenamebatch_colorizer.log, levellogging.ERROR, format%(asctime)s - %(message)s) # 在process_single_image的except块中 except Exception as e: error_msg f处理图片 {input_path} 时出错: {e} print(error_msg) logging.error(error_msg) # 可以选择跳过这张图继续处理下一张6.3 支持命令行参数让脚本可以通过命令行指定输入输出目录更加灵活。import argparse if __name__ __main__: parser argparse.ArgumentParser(description批量黑白图片上色脚本) parser.add_argument(-i, --input, defaultinput_images, help输入图片目录路径) parser.add_argument(-o, --output, defaultoutput_images, help输出图片目录路径) args parser.parse_args() main(input_dirargs.input, output_dirargs.output)这样你就可以这样运行了python batch_colorizer.py -i ./my_photos -o ./colored_results6.4 性能考虑并行处理如果你的图片非常多且电脑性能足够特别是多核CPU可以考虑使用多进程加速。from multiprocessing import Pool, cpu_count def main(): # ... 前面的代码 ... # 初始化模型等 # 创建进程池进程数一般设为CPU核心数 num_processes max(1, cpu_count() - 1) print(f使用 {num_processes} 个进程进行并行处理。) # 准备参数列表 task_list [(str(img_path), str(output_dir_path / f{img_path.stem}_colored{img_path.suffix}), model_pipeline) for img_path in image_files] with Pool(num_processes) as pool: # 注意model_pipeline可能无法直接序列化传递需要调整。 # 更稳妥的方式是在每个子进程内部初始化模型。 results pool.starmap(process_single_image_wrapper, task_list) # ... 后续统计 ...注意深度学习模型通常较大在多进程间传递可能有问题。一个更常见的模式是“生产者-消费者”队列或者让每个进程独立加载模型会占用更多显存/内存。对于GPU推理通常更推荐使用单进程因为GPU本身并行度就很高。7. 总结通过本教程我们完成了一次从交互式工具到自动化脚本的“升级”。你现在拥有了一个强大的批量图片上色工具它的核心价值在于效率倍增告别手动点击一次性处理成百上千张图片。流程固化将上色任务脚本化可以轻松集成到更大的自动化工作流中。灵活可控通过修改脚本你可以轻松调整输出格式、命名规则、错误处理策略等。核心步骤回顾理解原理找到原始工具中执行上色的核心函数。搭建框架编写脚本实现文件夹遍历、图片读取、函数调用、结果保存的闭环。连接桥梁正确导入原始工具的核心函数这是成功的关键。运行调试处理可能遇到的路径、导入、模型初始化等问题。优化扩展根据需求添加进度提示、日志、命令行参数等功能。这个思路不仅适用于cv_unet_image-colorization也适用于其他任何提供了核心处理函数的AI工具如图像修复、风格迁移、超分辨率等。掌握了这种方法你就掌握了将单点AI应用扩展为批量生产力工具的钥匙。获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。