cv_unet_image-colorization生产环境部署：支持批量处理+日志记录+错误重试机制

cv_unet_image-colorization生产环境部署支持批量处理日志记录错误重试机制你是不是遇到过这样的场景手里有一堆珍贵的黑白老照片想给它们上色但一张张手动处理太费时用在线工具又担心隐私泄露而且网络一卡顿整个流程就中断了。今天我要分享的正是解决这些痛点的方案一个基于cv_unet_image-colorization模型的生产级本地部署工具。它不仅修复了新版PyTorch的兼容性问题还内置了批量处理、日志记录和错误重试三大核心机制让你能稳定、高效、安全地处理海量照片。想象一下把几百张家庭老照片拖进一个文件夹点一下按钮泡杯咖啡的功夫所有照片都自动完成了上色整个过程都有日志可查即使某张图片出错也不会影响其他任务。这就是生产环境部署的价值。接下来我将带你从零开始搭建这个强大的本地照片上色流水线。1. 项目核心不只是上色更是稳定生产工具在深入部署细节前我们先搞清楚这个工具和普通Demo的区别。它基于ModelScope的cv_unet_image-colorization模型这个模型采用ResNet编码器提取特征再通过UNet结构的生成对抗网络GAN来“想象”并填充合理的色彩效果非常自然。但我们的重点不是模型原理而是如何让它从一个“玩具”变成“工具”。原始项目通常只支持单张图片交互式处理这远远达不到生产要求。我们为其注入三大生产级特性批量处理支持输入一个包含无数图片的文件夹自动遍历、处理、保存解放双手。日志记录每一步操作从读取图片、模型推理到保存结果都有详细日志输出到文件和控制台方便追溯和调试。错误重试机制网络波动、图片格式异常、显存偶尔不足没关系工具会捕获异常记录错误并可以选择重试或跳过确保批量任务不会因为单点故障而整体失败。此外我们还修复了一个关键问题PyTorch 2.6及以上版本为了安全默认以weights_onlyTrue模式加载模型这会直接导致加载旧格式的.pth模型文件时失败。我们的方案重写了加载逻辑确保兼容性。2. 环境搭建与一键部署理论说再多不如动手跑起来。我们先准备好战场。2.1 系统与硬件要求操作系统Linux (Ubuntu 20.04/22.04 推荐), Windows 10/11, macOS (M系列芯片注意ARM架构适配)Python3.8 - 3.10 (推荐3.8兼容性最佳)CUDA11.7 或 11.8 (如果你有NVIDIA GPU并希望GPU加速)内存建议16GB以上。处理高分辨率图片时消耗较大。显卡可选。有NVIDIA GPU显存4GB以上速度会快很多。CPU也能运行只是慢一些。2.2 一步到位的环境安装最省心的方式是使用我们准备好的环境配置文件。创建一个新的工作目录然后新建一个名为environment.yaml的文件。name: image-colorization-prod channels: - pytorch - conda-forge - defaults dependencies: - python3.8 - pip - cudatoolkit11.7 # 如果使用CPU注释掉这一行 - pip: - torch1.13.1cu117 # 对应CUDA 11.7CPU版请安装 torch1.13.1cpu - torchvision0.14.1 - modelscope1.9.5 - streamlit1.28.0 - opencv-python-headless4.8.1 - Pillow10.0.0 - tqdm4.66.1然后打开终端或Anaconda Prompt进入该目录执行以下命令来创建并激活环境# 使用conda创建环境如果没安装conda请先安装Miniconda conda env create -f environment.yaml # 激活环境 conda activate image-colorization-prod给新手的提示如果你没有GPU或者不想配置CUDA可以把上面yaml文件里cudatoolkit和torch那两行注释掉然后安装CPU版本的PyTorch。你可以在 PyTorch官网 找到对应的安装命令。2.3 获取并修复核心模型代码环境好了接下来是核心代码。我们需要从ModelScope获取模型并对其打上关键的兼容性补丁。首先直接通过Python代码加载模型这会让ModelScope自动帮我们下载模型文件到本地。# download_model.py from modelscope.pipelines import pipeline from modelscope.utils.constant import Tasks # 这行代码会触发模型下载。模型默认会保存在 ~/.cache/modelscope/hub 目录下。 # 如果网络较慢可能需要等待一段时间。 colorizer pipeline(Tasks.image_colorization, modeldamo/cv_unet_image-colorization) print(模型下载完成)运行这个脚本后模型就准备好了。接下来是关键修复。由于PyTorch 2.6的安全限制直接加载这个模型会报错。我们需要找到下载的模型文件通常路径是~/.cache/modelscope/hub/damo/cv_unet_image-colorization。我们需要修改模型加载方式。创建一个新的文件比如叫colorization_prod.py在这里面实现我们修复后的、支持生产功能的管道。# colorization_prod.py import os import sys import torch import logging import traceback from pathlib import Path from typing import List, Optional from PIL import Image import cv2 import numpy as np from modelscope.pipelines import pipeline from modelscope.utils.constant import Tasks from tqdm import tqdm import streamlit as st # 1. 修复PyTorch 2.6 模型加载问题 # 猴子补丁重写 torch.load 的安全检查行为 _original_torch_load torch.load def _patched_torch_load(f, *args, **kwargs): # 强制设置 weights_onlyFalse 以加载旧式模型 kwargs[weights_only] False return _original_torch_load(f, *args, **kwargs) torch.load _patched_torch_load # 2. 配置生产级日志 logging.basicConfig( levellogging.INFO, format%(asctime)s - %(name)s - %(levelname)s - %(message)s, handlers[ logging.FileHandler(colorization_batch.log), logging.StreamHandler(sys.stdout) ] ) logger logging.getLogger(__name__) class ImageColorizationProducer: 生产级图像上色处理器 def __init__(self, device: str None): 初始化上色管道 Args: device: 指定设备如 cuda:0 或 cpu。为None则自动检测。 self.device device if device else (cuda:0 if torch.cuda.is_available() else cpu) logger.info(f初始化模型运行设备: {self.device}) try: # 加载修复后的模型管道 self.colorizer pipeline( Tasks.image_colorization, modeldamo/cv_unet_image-colorization, deviceself.device ) logger.info(模型管道加载成功) except Exception as e: logger.error(f模型加载失败: {e}) raise def colorize_single(self, image_path: Path, output_dir: Path, max_retries: int 3): 处理单张图片包含错误重试机制 Args: image_path: 输入图片路径 output_dir: 输出目录 max_retries: 最大重试次数 Returns: bool: 是否成功 str: 输出文件路径或错误信息 for attempt in range(max_retries): try: logger.info(f尝试处理图片: {image_path.name} (第{attempt1}次尝试)) # 读取图片 img Image.open(image_path).convert(RGB) img_np np.array(img) # 执行上色 result self.colorizer(img_np) colorized_np result[output_img] # 模型输出是numpy数组 # 确保输出目录存在 output_dir.mkdir(parentsTrue, exist_okTrue) # 保存结果 output_path output_dir / fcolorized_{image_path.stem}.jpg # 使用OpenCV保存BGR顺序或PIL colorized_bgr cv2.cvtColor(colorized_np, cv2.COLOR_RGB2BGR) cv2.imwrite(str(output_path), colorized_bgr) logger.info(f图片处理成功: {output_path}) return True, str(output_path) except Exception as e: logger.warning(f处理图片 {image_path.name} 失败 (尝试 {attempt1}/{max_retries}): {e}) if attempt max_retries - 1: # 最后一次尝试也失败 error_msg f处理失败: {image_path.name} - {str(e)} logger.error(error_msg) return False, error_msg # 非最后一次失败等待后重试 import time time.sleep(2) # 等待2秒后重试 def process_batch(self, input_dir: Path, output_dir: Path, max_retries: int 3): 批量处理目录中的所有图片 Args: input_dir: 输入图片目录 output_dir: 输出目录 max_retries: 每张图最大重试次数 Returns: dict: 处理结果统计 if not input_dir.exists(): logger.error(f输入目录不存在: {input_dir}) return {error: 输入目录不存在} # 支持的图片格式 valid_extensions {.jpg, .jpeg, .png, .bmp, .tiff} image_files [f for f in input_dir.iterdir() if f.suffix.lower() in valid_extensions] if not image_files: logger.warning(f在目录 {input_dir} 中未找到支持的图片文件) return {error: 未找到图片文件} logger.info(f开始批量处理共 {len(image_files)} 张图片) success_count 0 fail_count 0 fail_details [] # 使用tqdm显示进度条 for img_path in tqdm(image_files, desc批量上色进度): success, result self.colorize_single(img_path, output_dir, max_retries) if success: success_count 1 else: fail_count 1 fail_details.append(result) stats { total: len(image_files), success: success_count, failed: fail_count, fail_details: fail_details } logger.info(f批量处理完成。成功: {success_count}, 失败: {fail_count}) return stats # 后续将在这里添加Streamlit界面代码 if __name__ __main__: # 这里是测试代码 producer ImageColorizationProducer() # 测试单张图片 # producer.colorize_single(Path(test.jpg), Path(./output))这段代码就是我们生产工具的核心引擎。它定义了一个ImageColorizationProducer类封装了模型加载、单张处理、批量处理和日志错误处理的所有逻辑。3. 构建Streamlit可视化界面有了强大的后端引擎我们还需要一个友好易用的前端界面。Streamlit非常适合快速构建数据科学应用。我们在同一个文件中继续添加界面代码。# 接 colorization_prod.py 的 __main__ 部分之前 def main(): Streamlit主应用界面 st.set_page_config( page_title生产级老照片上色工具, page_icon, layoutwide ) st.title( 生产级老照片上色工具) st.markdown( 基于 **cv_unet_image-colorization** 模型支持批量处理、日志记录与错误重试。 纯本地运行无需上传至云端保护您的隐私。 ) # 初始化处理器单例模式利用Streamlit的缓存 st.cache_resource def get_colorizer(): return ImageColorizationProducer() colorizer_producer get_colorizer() # 侧边栏模式选择 st.sidebar.header(⚙️ 工作模式) mode st.sidebar.radio( 选择处理模式, [单张图片上色, 批量图片上色] ) if mode 单张图片上色: st.header(单张图片上色) col1, col2 st.columns(2) with col1: st.subheader(1. 上传黑白照片) uploaded_file st.file_uploader( 选择一张图片, type[jpg, jpeg, png, bmp, tiff] ) if uploaded_file is not None: # 显示原图 image Image.open(uploaded_file).convert(RGB) st.image(image, caption原始黑白图片, use_column_widthTrue) # 临时保存上传的文件 input_path Path(./temp_input) input_path.mkdir(exist_okTrue) temp_file_path input_path / uploaded_file.name with open(temp_file_path, wb) as f: f.write(uploaded_file.getbuffer()) with col2: st.subheader(2. 上色结果) if st.button( 开始上色, typeprimary): with st.spinner(AI正在为照片上色请稍候...): output_dir Path(./output_single) success, result colorizer_producer.colorize_single( temp_file_path, output_dir ) if success: st.success(处理完成) colorized_image Image.open(result) st.image(colorized_image, captionAI上色结果, use_column_widthTrue) # 提供下载链接 with open(result, rb) as file: st.download_button( label下载上色图片, datafile, file_namefcolorized_{uploaded_file.name}, mimeimage/jpeg ) else: st.error(f处理失败: {result}) else: # 批量图片上色模式 st.header(批量图片上色) st.info(此模式将处理指定文件夹内的所有图片适合大量老照片修复。) col1, col2 st.columns(2) with col1: st.subheader(1. 设置输入输出目录) # 注意Streamlit Cloud等在线环境可能限制文件系统访问。 # 本地运行时可以手动输入路径或使用更高级的文件选择器。 input_dir_str st.text_input(输入目录路径包含黑白图片, value./batch_input) output_dir_str st.text_input(输出目录路径, value./batch_output) max_retries st.slider(每张图片最大重试次数, 1, 5, 3) with col2: st.subheader(2. 执行批量处理) st.markdown(f **配置摘要**: - 输入目录: {input_dir_str} - 输出目录: {output_dir_str} - 最大重试次数: {max_retries} ) if st.button( 开始批量处理, typeprimary): input_dir Path(input_dir_str) output_dir Path(output_dir_str) if not input_dir.exists(): st.error(f输入目录不存在: {input_dir}) else: with st.spinner(正在批量处理中请查看控制台日志...): stats colorizer_producer.process_batch(input_dir, output_dir, max_retries) if error in stats: st.error(stats[error]) else: st.success(f批量处理完成) st.metric(label成功数量, valuestats[success]) st.metric(label失败数量, valuestats[failed]) if stats[failed] 0: with st.expander(查看失败详情): for detail in stats[fail_details]: st.text(detail) # 显示部分处理结果预览 output_dir_path Path(output_dir_str) output_images list(output_dir_path.glob(colorized_*.jpg))[:5] # 预览前5张 if output_images: st.subheader(结果预览前5张) preview_cols st.columns(min(len(output_images), 5)) for idx, img_path in enumerate(output_images): with preview_cols[idx]: st.image(Image.open(img_path), captionimg_path.name, use_column_widthTrue) # 侧边栏日志查看器 st.sidebar.header( 系统日志) if st.sidebar.button(查看最新日志): log_file Path(colorization_batch.log) if log_file.exists(): with open(log_file, r, encodingutf-8) as f: log_content f.read() st.sidebar.text_area(日志内容, log_content[-5000:], height300) # 显示最后5000字符 else: st.sidebar.warning(日志文件尚未生成。) st.sidebar.markdown(---) st.sidebar.info( **使用提示**: - 首次运行会下载模型约1.2GB请保持网络通畅。 - 批量处理时请确保输入目录路径正确。 - 所有日志保存在 colorization_batch.log 文件中。 ) if __name__ __main__: # 如果是直接运行这个脚本启动Streamlit界面实际部署时通过命令启动 # 这里我们直接调用main函数但实际部署是通过 streamlit run colorization_prod.py 命令 # 为了脚本也能独立测试我们做个判断 if len(sys.argv) 1 and sys.argv[1] --cli: # 命令行批处理测试模式 producer ImageColorizationProducer() stats producer.process_batch(Path(./test_input), Path(./test_output)) print(stats) else: # 正常启动Streamlit实际部署时不从这里启动 # 提示用户使用正确的命令 print(请使用命令启动应用: streamlit run colorization_prod.py) # 但我们也可以直接调用main()来测试界面功能非Streamlit服务器环境可能报错 # main()现在我们的完整生产工具就编码完成了。它包含了修复后的模型加载、健壮的错误处理、批量逻辑和一个美观的Web界面。4. 生产环境部署与运行代码写好了怎么把它跑起来尤其是长期稳定地运行呢4.1 启动应用确保你在之前创建的image-colorization-prod环境中并且终端位于脚本所在目录。# 启动Streamlit应用服务器 streamlit run colorization_prod.py执行后终端会输出一个本地URL通常是http://localhost:8501。用浏览器打开它你就能看到我们刚刚构建的工具界面了。4.2 使用指南单张与批量处理单张处理在左侧选择“单张图片上色”模式。在左侧面板上传你的黑白照片。点击右侧的“开始上色”按钮。等待处理完成预览并下载结果。批量处理在左侧选择“批量图片上色”模式。准备一个文件夹例如batch_input把所有要处理的图片放进去。在界面输入输入目录和输出目录的路径。设置最大重试次数通常3次足够。点击“开始批量处理”。处理过程中可以随时点击侧边栏的“查看最新日志”按钮监控进度和错误。处理完成后查看统计信息和结果预览。4.3 进阶后台服务与自动化对于真正的生产环境你可能希望它作为一个后台服务运行或者定时自动处理某个文件夹里的新照片。方案一使用系统服务Linux创建一个系统服务文件比如/etc/systemd/system/colorization.service[Unit] DescriptionImage Colorization Production Tool Afternetwork.target [Service] Typesimple Useryour_username WorkingDirectory/path/to/your/project EnvironmentPATH/home/your_username/miniconda3/envs/image-colorization-prod/bin ExecStart/home/your_username/miniconda3/envs/image-colorization-prod/bin/streamlit run colorization_prod.py --server.port 8501 --server.headless true Restarton-failure RestartSec10 [Install] WantedBymulti-user.target然后启用并启动服务sudo systemctl daemon-reload sudo systemctl enable colorization.service sudo systemctl start colorization.service # 查看状态 sudo systemctl status colorization.service方案二编写自动化脚本你可以编写一个Python脚本定时扫描特定文件夹并调用我们写好的ImageColorizationProducer类进行自动处理。# auto_process.py import schedule import time from pathlib import Path from colorization_prod import ImageColorizationProducer def job(): 定时任务处理watch_folder中的新图片 input_folder Path(/path/to/watch_folder) output_folder Path(/path/to/output_folder) processed_log Path(/path/to/processed.log) # 读取已处理过的文件列表 processed set() if processed_log.exists(): with open(processed_log, r) as f: processed set(f.read().splitlines()) # 查找新文件 valid_ext {.jpg, .jpeg, .png} new_files [f for f in input_folder.iterdir() if f.suffix.lower() in valid_ext and f.name not in processed] if new_files: producer ImageColorizationProducer() for file in new_files: success, _ producer.colorize_single(file, output_folder) if success: # 记录已处理 with open(processed_log, a) as f: f.write(f{file.name}\n) print(f已处理并记录: {file.name}) # 每10分钟运行一次 schedule.every(10).minutes.do(job) print(自动化批处理服务已启动...) while True: schedule.run_pending() time.sleep(1)5. 总结通过以上步骤我们成功将一个基础的AI模型部署成了一个具备生产级能力的工具。我们来回顾一下关键点核心修复通过重写torch.load解决了PyTorch 2.6的模型兼容性问题这是项目能跑起来的前提。生产化改造我们为工具添加了三大支柱——批量处理能力让你能一次性处理整个照片库日志记录系统让每一次运行都有迹可循错误重试机制确保了长时间批量任务的稳定性。友好界面利用Streamlit快速构建了直观的Web界面同时支持单张体验和批量作业满足了不同场景的需求。部署灵活性工具可以以交互式Web应用、后台系统服务或定时自动化脚本等多种方式运行适应从个人到轻量级生产的不同环境。这个项目的价值在于它把前沿的AI图像上色能力封装成了一个可靠、易用且私密的本地化工具。无论是修复家族相册还是处理历史档案你都可以在自己的电脑上完成数据不出本地安全又高效。获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。