cv_unet_image-colorization从零开始：PyTorch 2.6+兼容性修复代码详解

cv_unet_image-colorization从零开始PyTorch 2.6兼容性修复代码详解你有没有遇到过这样的情况翻出家里的老相册看到那些珍贵的黑白照片总想着如果它们是彩色的该多好。或者在网上找到一张历史感十足的黑白图片却因为缺乏色彩而显得单调。现在借助AI技术我们可以轻松地为这些照片“上色”让历史重现光彩。今天我要介绍的就是一个能帮你实现这个愿望的本地工具——基于ModelScope的cv_unet_image-colorization模型。但更重要的是我会带你从零开始一步步解决一个关键的技术难题如何在PyTorch 2.6版本上顺利运行这个模型。如果你之前尝试过运行一些旧的AI模型很可能遇到过这样的报错“RuntimeError: PytorchStreamReader failed reading zip archive: failed finding central directory”。这就是PyTorch 2.6版本引入的安全限制导致的。别担心读完这篇文章你不仅能学会如何修复这个问题还能亲手搭建一个功能完整的黑白照片上色工具。1. 项目核心不只是上色更是兼容性修复这个工具的核心价值其实有两层。第一层当然是它的主要功能——为黑白照片自动上色。但更深一层的价值是它解决了新旧技术栈之间的兼容性问题。1.1 为什么需要修复PyTorch在2.6版本引入了一个重要的安全特性默认情况下torch.load()函数只允许加载“安全”的模型文件即weights_onlyTrue。这个改动本意是好的能防止恶意代码通过模型文件执行。但问题在于很多旧的模型文件包括我们今天要用的这个在保存时包含了完整的Python对象而不仅仅是权重参数。当PyTorch 2.6尝试加载这些旧模型时就会抛出错误因为它检测到文件里包含了“不安全”的内容。我们的修复方案就是告诉PyTorch“我知道这个文件可能包含代码但我信任它请继续加载。”1.2 工具能做什么简单来说这个工具能帮你自动上色上传一张黑白照片AI会自动分析图像内容为不同的物体填充合理的颜色本地运行所有处理都在你的电脑上完成照片不会上传到任何服务器保护你的隐私快速处理利用GPU加速即使是高清图片也能在短时间内完成上色直观对比处理完成后你可以并排查看原图和上色后的效果2. 环境准备搭建你的AI上色工作站在开始写代码之前我们需要准备好运行环境。别担心步骤很简单跟着我做就行。2.1 检查你的硬件和软件首先确保你的电脑满足以下要求操作系统Windows 10/11macOS或者Linux都可以Python版本3.8或更高版本显卡有NVIDIA显卡最好能大幅加速没有的话用CPU也能运行只是会慢一些内存至少8GB处理大图片时建议16GB以上2.2 安装必要的软件包打开你的命令行工具Windows上是CMD或PowerShellmacOS/Linux上是Terminal依次执行以下命令# 首先创建一个新的Python环境可选但推荐 python -m venv deoldify_env # 激活环境 # Windows: deoldify_env\Scripts\activate # macOS/Linux: source deoldify_env/bin/activate # 安装PyTorch根据你的CUDA版本选择如果没有GPU就选CPU版本 # 有CUDA 11.8的GPU pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 # 只有CPU pip install torch torchvision torchaudio # 安装其他必要的包 pip install modelscope1.11.0 pip install streamlit1.31.0 pip install Pillow10.2.0 pip install opencv-python4.9.0.80 pip install numpy1.24.4安装过程可能需要几分钟取决于你的网速。如果遇到网络问题可以尝试使用国内的镜像源pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118 -i https://pypi.tuna.tsinghua.edu.cn/simple3. 核心修复让旧模型在新PyTorch上运行这是本文最核心的部分。我们要解决PyTorch 2.6加载旧模型的问题。原理其实很简单重写torch.load函数的行为。3.1 理解问题根源在PyTorch 2.6之前加载模型通常是这样的import torch model torch.load(old_model.pth)从PyTorch 2.6开始默认行为变成了model torch.load(old_model.pth, weights_onlyTrue)这个weights_onlyTrue参数就是问题的关键。当它为True时PyTorch会严格检查模型文件如果发现里面包含了Python代码比如自定义的类定义就会拒绝加载。3.2 实现兼容性修复我们需要创建一个补丁文件在加载模型之前修改PyTorch的行为。创建一个新文件命名为compatibility_patch.py PyTorch 2.6 兼容性修复补丁 解决旧模型文件加载时的安全性错误 import torch import warnings from typing import Any, Optional import os def safe_torch_load( filepath: str, map_location: Any None, pickle_module: Any None, **kwargs: Any ) - Any: 安全的torch.load替代函数 强制设置weights_onlyFalse以兼容旧模型 Args: filepath: 模型文件路径 map_location: 设备映射如cuda:0或cpu pickle_module: pickle模块通常不需要指定 **kwargs: 其他torch.load参数 Returns: 加载的模型或数据 # 检查文件是否存在 if not os.path.exists(filepath): raise FileNotFoundError(f模型文件不存在: {filepath}) # 显示警告信息让用户知道我们正在使用非安全模式 warnings.warn( f正在以非安全模式加载模型文件: {filepath}\n 此模式仅用于加载可信的旧版模型文件。\n 请确保文件来源可靠避免加载未知来源的模型。, UserWarning, stacklevel2 ) try: # 关键修复强制设置weights_onlyFalse # 这样PyTorch就不会检查模型文件中的Python代码 result torch.load( filepath, map_locationmap_location, pickle_modulepickle_module, weights_onlyFalse, # 这是修复的核心 **kwargs ) return result except Exception as e: # 如果还是出错提供更详细的错误信息 error_msg ( f加载模型文件失败: {filepath}\n f错误类型: {type(e).__name__}\n f错误信息: {str(e)}\n 可能的原因:\n 1. 模型文件已损坏\n 2. PyTorch版本不兼容\n 3. 文件格式不正确 ) raise RuntimeError(error_msg) from e def apply_patch(): 应用兼容性补丁 用我们的安全加载函数替换torch.load # 保存原始的torch.load函数 original_torch_load torch.load def patched_load(*args, **kwargs): 修补后的加载函数 自动处理weights_only参数 # 如果用户没有显式指定weights_only我们设置为False # 如果用户显式指定了尊重用户的选择 if weights_only not in kwargs: kwargs[weights_only] False warnings.warn( 自动启用兼容模式 (weights_onlyFalse) 以加载旧模型, UserWarning, stacklevel2 ) return original_torch_load(*args, **kwargs) # 替换torch.load torch.load patched_load print(✅ 兼容性补丁已应用) print( 现在可以加载PyTorch 2.6之前保存的模型文件了) return original_torch_load def restore_patch(original_func): 恢复原始torch.load函数 Args: original_func: 原始的torch.load函数 torch.load original_func print(✅ 兼容性补丁已移除恢复原始torch.load行为) # 使用示例 if __name__ __main__: # 应用补丁 original_load apply_patch() # 现在可以安全加载旧模型了 try: # 这里替换成你的模型文件路径 # model torch.load(your_old_model.pth) print(补丁测试完成) finally: # 恢复原始函数可选 restore_patch(original_load)这个修复代码做了几件重要的事情自动处理兼容性当加载模型时如果没有指定weights_only参数自动设置为False提供明确警告让用户知道正在以非安全模式加载模型错误处理如果加载失败提供详细的错误信息帮助调试可逆操作可以随时恢复原始的torch.load行为3.3 在项目中应用修复在实际项目中我们只需要在程序开始时应用这个补丁。创建一个主程序文件main.py# 首先应用兼容性补丁 from compatibility_patch import apply_patch original_torch_load apply_patch() # 现在导入其他依赖 import torch from modelscope.pipelines import pipeline from modelscope.utils.constant import Tasks import streamlit as st from PIL import Image import numpy as np import cv2 import tempfile import os # ... 其他代码 ... # 程序结束时可以恢复原始函数可选 # restore_patch(original_torch_load)4. 构建完整的照片上色工具修复了兼容性问题后我们就可以专注于构建上色工具本身了。这个工具的核心是ModelScope的cv_unet_image-colorization模型它基于一个聪明的技术组合ResNet编码器 UNet生成对抗网络。4.1 理解模型的工作原理简单来说这个模型的工作流程是这样的编码阶段使用ResNet一个很厉害的图像识别网络分析黑白图片理解里面有什么物体人、树、房子等理解上下文模型根据学到的知识推断每个物体应该是什么颜色天空是蓝色的树叶是绿色的等生成阶段使用UNet网络擅长生成细节为图片填充颜色对抗训练还有一个“判别器”网络在训练时不断判断生成的颜色是否真实让模型越来越聪明4.2 创建完整的应用程序让我们创建一个完整的Streamlit应用。新建一个文件deoldify_app.py 黑白照片上色工具 - 完整实现 基于ModelScope cv_unet_image-colorization模型 import streamlit as st import torch from PIL import Image import numpy as np import cv2 import tempfile import os from datetime import datetime # 应用兼容性补丁 try: from compatibility_patch import apply_patch original_torch_load apply_patch() except ImportError: st.warning(兼容性补丁未找到如果使用PyTorch 2.6可能会出错) original_torch_load None # 设置页面配置 st.set_page_config( page_title黑白照片上色工具, page_icon, layoutwide, initial_sidebar_stateexpanded ) # 初始化模型使用缓存避免重复加载 st.cache_resource def load_colorization_model(): 加载上色模型 使用缓存只在第一次运行时加载 try: from modelscope.pipelines import pipeline from modelscope.utils.constant import Tasks # 显示加载状态 with st.spinner(正在加载AI上色模型首次加载可能需要1-2分钟...): # 创建上色pipeline colorizer pipeline( Tasks.image_colorization, modeldamo/cv_unet_image-colorization, devicecuda if torch.cuda.is_available() else cpu ) st.success(✅ 模型加载完成) if torch.cuda.is_available(): st.info(f使用GPU加速: {torch.cuda.get_device_name(0)}) else: st.info(使用CPU运行速度较慢建议使用GPU) return colorizer except Exception as e: st.error(f模型加载失败: {str(e)}) st.info(请确保已安装modelscope: pip install modelscope) return None def preprocess_image(image): 预处理上传的图片 确保图片格式正确调整大小可选 # 转换为RGB处理可能的上传格式问题 if image.mode ! RGB: image image.convert(RGB) # 转换为numpy数组 img_array np.array(image) # 可选调整图片大小避免处理过大图片 max_size 1024 # 最大边长 height, width img_array.shape[:2] if max(height, width) max_size: # 计算缩放比例 scale max_size / max(height, width) new_width int(width * scale) new_height int(height * scale) # 使用高质量缩放 img_array cv2.resize(img_array, (new_width, new_height), interpolationcv2.INTER_LANCZOS4) return img_array def colorize_image(model, image_array): 使用模型为图片上色 if model is None: st.error(模型未加载无法进行上色) return None try: with st.spinner(AI正在为照片上色请稍候...): # 执行上色 result model(image_array) # 提取结果 if isinstance(result, dict) and output_img in result: output_img result[output_img] else: output_img result # 确保输出是numpy数组 if isinstance(output_img, Image.Image): output_img np.array(output_img) return output_img except Exception as e: st.error(f上色过程出错: {str(e)}) return None def main(): 主应用程序 # 标题和介绍 st.title( 黑白照片上色工具) st.markdown( 使用AI技术为黑白照片自动上色让历史重现光彩。 所有处理均在本地完成保护您的隐私。 ) # 侧边栏 with st.sidebar: st.header( 上传照片) # 文件上传器 uploaded_file st.file_uploader( 选择一张黑白/老照片, type[jpg, jpeg, png, bmp, tiff], help支持JPG、PNG、BMP、TIFF格式建议图片大小不超过10MB ) st.markdown(---) st.header(⚙️ 设置) # 性能选项 use_gpu st.checkbox(使用GPU加速如果可用, valueTrue) # 图片处理选项 preserve_size st.checkbox(保持原图尺寸, valueTrue) st.markdown(---) st.header(ℹ️ 使用说明) st.markdown( 1. 在左侧上传黑白照片 2. 点击「开始上色」按钮 3. 等待AI处理完成 4. 查看并下载彩色结果 **注意** - 首次使用需要加载模型约1-2分钟 - 处理时间取决于图片大小和硬件性能 - 建议使用清晰的黑白照片以获得最佳效果 ) # 主内容区 col1, col2 st.columns(2) # 加载模型全局一次 if colorizer not in st.session_state: st.session_state.colorizer load_colorization_model() colorizer st.session_state.colorizer # 左侧列原图 with col1: st.subheader( 原始照片) if uploaded_file is not None: # 显示原图 try: image Image.open(uploaded_file) st.image(image, caption上传的原始照片, use_column_widthTrue) # 保存到session state st.session_state.original_image image st.session_state.image_array preprocess_image(image) # 显示图片信息 st.info(f 图片信息 - 格式{image.format} - 尺寸{image.size[0]} × {image.size[1]} 像素 - 模式{image.mode} ) except Exception as e: st.error(f无法读取图片: {str(e)}) else: st.info(请从左侧上传一张黑白照片) # 显示示例图片 st.image(https://via.placeholder.com/600x400/CCCCCC/000000?text示例黑白照片, caption示例上传后这里会显示你的照片, use_column_widthTrue) # 右侧列上色结果 with col2: st.subheader( 上色结果) # 上色按钮 if uploaded_file is not None and image_array in st.session_state: if st.button( 开始上色, typeprimary, use_container_widthTrue): if colorizer is None: st.error(模型加载失败请检查控制台输出) else: # 执行上色 colored_image colorize_image(colorizer, st.session_state.image_array) if colored_image is not None: # 显示结果 st.image(colored_image, captionAI上色后的照片, use_column_widthTrue) st.success(✅ 处理完成) # 保存结果到session state st.session_state.colored_image colored_image # 下载按钮 colored_pil Image.fromarray(colored_image) # 创建临时文件供下载 with tempfile.NamedTemporaryFile(deleteFalse, suffix.jpg) as tmp_file: colored_pil.save(tmp_file.name, JPEG, quality95) tmp_file_path tmp_file.name with open(tmp_file_path, rb) as file: st.download_button( label 下载彩色照片, datafile, file_namefcolored_{uploaded_file.name}, mimeimage/jpeg, use_container_widthTrue ) # 清理临时文件 os.unlink(tmp_file_path) # 显示处理信息 st.info( **上色完成** - AI已根据图像内容自动填充合理颜色 - 您可以通过下载按钮保存结果 - 可以尝试上传其他照片继续体验 ) else: st.info(点击「开始上色」按钮AI将为照片添加色彩) else: st.info(上传照片后这里将显示上色结果) # 底部信息 st.markdown(---) st.markdown( ### 使用技巧 1. **最佳效果**使用清晰、高对比度的黑白照片 2. **人物照片**AI能较好识别人脸和肤色 3. **风景照片**天空、树木、建筑等通常上色效果较好 4. **历史照片**老照片可能有划痕或噪点效果可能受影响 ### 技术信息 - **模型架构**ResNet编码器 UNet生成对抗网络 - **兼容性**已修复PyTorch 2.6加载问题 - **运行方式**纯本地处理无需网络连接 - **隐私保护**所有图片处理均在本地完成 ) if __name__ __main__: # 检查必要的包 try: import modelscope import streamlit main() except ImportError as e: st.error(f缺少必要的包: {str(e)}) st.info(请运行: pip install modelscope streamlit Pillow opencv-python)4.3 运行你的上色工具保存所有文件后在命令行中运行streamlit run deoldify_app.pyStreamlit会自动打开浏览器显示你的上色工具界面。第一次运行时会下载模型文件大约几百MB需要一些时间。5. 实际效果展示与使用技巧现在工具已经可以运行了让我们看看它能做什么以及如何获得最佳效果。5.1 看看实际效果我测试了几种不同类型的照片以下是观察到的效果人物照片AI能很好地识别人脸为皮肤添加自然的肤色为头发添加适当的颜色。衣服的颜色通常也比较合理。风景照片天空会被上色为蓝色如果是白天树木和草地是绿色建筑物会根据材质上色。建筑与街景砖墙、窗户、道路等元素的上色通常比较准确能还原出历史场景的色彩。老照片修复对于有划痕或噪点的老照片上色后视觉效果会有明显提升但原有的瑕疵可能仍然可见。5.2 获得最佳效果的建议根据我的测试经验这里有一些实用建议选择高质量源图片清晰、高对比度的黑白照片效果最好适当调整亮度如果原图太暗或太亮可以先简单调整一下理解AI的局限AI根据常见物体颜色进行上色可能不完全符合历史事实非常规颜色的物体比如紫色的猫可能被上色为常见颜色极低分辨率或严重损坏的照片效果可能不理想多次尝试有时同样的照片重新处理一次可能会有不同的颜色分布5.3 处理不同尺寸的图片工具会自动处理大图片但如果你需要手动控制可以修改预处理函数def preprocess_image(image, max_size1024, force_sizeNone): 增强的图片预处理 Args: image: PIL Image对象 max_size: 最大边长保持比例 force_size: 强制调整到指定大小 (宽, 高) # 转换为RGB if image.mode ! RGB: image image.convert(RGB) img_array np.array(image) # 强制调整大小 if force_size is not None: new_width, new_height force_size img_array cv2.resize(img_array, (new_width, new_height), interpolationcv2.INTER_LANCZOS4) # 按比例调整 elif max_size is not None: height, width img_array.shape[:2] if max(height, width) max_size: scale max_size / max(height, width) new_width int(width * scale) new_height int(height * scale) img_array cv2.resize(img_array, (new_width, new_height), interpolationcv2.INTER_LANCZOS4) return img_array6. 常见问题与解决方案在开发和使用过程中你可能会遇到一些问题。这里是我总结的常见问题及解决方法。6.1 模型加载问题问题首次运行时下载模型失败或速度很慢解决# 方法1设置镜像源在代码开头添加 import os os.environ[MODELSCOPE_CACHE] ./model_cache # 自定义缓存目录 os.environ[MODELSCOPE_ENDPOINT] https://mirrors.aliyun.com/modelscope/ # 使用阿里云镜像 # 方法2手动下载模型如果自动下载失败 # 从 https://modelscope.cn/models/damo/cv_unet_image-colorization 手动下载 # 然后放到 ~/.cache/modelscope/hub/damo/cv_unet_image-colorization/6.2 内存不足问题问题处理大图片时出现内存错误解决# 在预处理函数中限制图片大小 def preprocess_image(image, max_size768): # 减小最大尺寸 # ... 原有代码 ... # 或者在Streamlit侧边栏添加选项 max_size st.slider(最大处理尺寸, 512, 2048, 1024, 128)6.3 GPU相关问题问题1有GPU但程序仍然使用CPU解决# 明确指定使用GPU device cuda if torch.cuda.is_available() else cpu st.info(f使用设备: {device}) # 如果检测到GPU但无法使用尝试 import torch print(fCUDA可用: {torch.cuda.is_available()}) print(fGPU数量: {torch.cuda.device_count()}) if torch.cuda.is_available(): print(fGPU名称: {torch.cuda.get_device_name(0)})问题2GPU内存不足解决# 减少批量大小如果支持 # 或者处理前缩小图片尺寸 # 清理GPU缓存 torch.cuda.empty_cache()6.4 其他常见错误错误AttributeError: module torch has no attribute load解决这通常是因为兼容性补丁没有正确应用。确保在导入其他torch相关模块前应用补丁。错误RuntimeError: Failed to load image解决检查图片格式确保使用PIL或OpenCV支持的标准格式。7. 总结通过这篇文章我们完成了一个完整的项目从解决PyTorch 2.6的兼容性问题到构建一个功能完善的本地黑白照片上色工具。7.1 关键收获让我们回顾一下最重要的几点兼容性修复是核心我们通过重写torch.load函数设置weights_onlyFalse解决了新旧PyTorch版本之间的兼容性问题。这个技巧不仅适用于这个项目也适用于其他需要加载旧模型的场景。本地化处理保护隐私所有照片处理都在本地完成不需要上传到任何服务器这对于处理私人或敏感照片非常重要。完整的端到端解决方案我们从环境搭建、兼容性修复、模型加载到界面开发覆盖了完整的技术栈。实用的工程技巧包括错误处理、进度提示、图片预处理、结果保存等这些都是实际项目中必不可少的。7.2 进一步探索的方向如果你对这个项目感兴趣还可以尝试以下扩展批量处理功能修改程序支持一次上传多张照片批量上色颜色调整在上色后允许用户手动调整颜色历史照片修复结合去噪、去划痕算法先修复再上色视频上色扩展功能处理黑白视频需要更复杂的帧处理自定义训练使用自己的数据集微调模型获得特定风格的上色效果7.3 最后的建议技术工具的价值在于解决实际问题。这个黑白照片上色工具不仅是一个技术演示更是一个能真正帮助人们重温历史、保存记忆的实用工具。当你看到一张几十年前的黑白家庭照片变成彩色看到祖辈的面容重新焕发光彩你会感受到技术的人文温度。这正是我们做技术开发的意义——用代码创造价值用技术连接情感。现在你可以运行这个工具试着为你家的老照片上色或者用它来修复那些珍贵的历史影像。每一次成功的上色都是对过去的一次致敬对未来的一次连接。获取更多AI镜像想探索更多AI镜像和应用场景访问 CSDN星图镜像广场提供丰富的预置镜像覆盖大模型推理、图像生成、视频生成、模型微调等多个领域支持一键部署。