当前位置：首页 > news >正文

零门槛AI上色：cv_unet_image-colorization+Streamlit可视化工具教程

news 2026/8/2 19:25:51

零门槛AI上色：cv_unet_image-colorization+Streamlit可视化工具教程

你是否有一堆泛黄的老照片，想要让它们重现当年的色彩？或者想给黑白照片添加一些艺术气息？今天我要介绍的这个工具，能让这些需求变得轻而易举。不需要任何AI背景，跟着这篇教程，你就能在自己的电脑上搭建一个专业的照片上色工具。

1. 工具简介与准备工作

1.1 什么是cv_unet_image-colorization

cv_unet_image-colorization是一个基于深度学习的图像上色模型，它采用了ResNet编码器和UNet生成对抗网络(GAN)架构。这个模型能够理解图像内容，并为其填充符合现实的合理颜色。最棒的是，它完全在本地运行，你的照片数据不会上传到任何服务器，确保了隐私安全。

1.2 你需要准备什么

在开始之前，请确保你的电脑满足以下要求：

操作系统：Windows 10/11、macOS或Linux
显卡：支持CUDA的NVIDIA显卡（可选但推荐）
Python 3.8或更高版本
至少4GB可用存储空间（用于存放模型）

2. 环境安装与配置

2.1 安装Python依赖

首先，我们需要安装运行这个工具所需的所有Python库。打开你的命令行工具（Windows上是CMD或PowerShell，Mac/Linux上是终端），输入以下命令：

pip install modelscope streamlit opencv-python pillow numpy torch

这条命令会安装以下组件：

modelscope：阿里开源的模型社区框架
streamlit：用于构建可视化界面
opencv-python和Pillow：图像处理库
numpy：科学计算基础库
torch：PyTorch深度学习框架

如果下载速度慢，可以加上国内镜像源：

pip install -i https://pypi.tuna.tsinghua.edu.cn/simple modelscope streamlit opencv-python pillow numpy torch

2.2 下载模型文件

接下来，我们需要获取模型文件：

访问ModelScope官网，搜索"cv_unet_image-colorization"
找到对应的模型页面，下载模型文件
在本地创建一个文件夹存放模型，例如D:\ai_models\image_colorization
将下载的模型文件夹（通常名为iic/cv_unet_image-colorization）放入上述目录

记住这个路径，我们稍后会用到。

3. 创建照片上色应用

3.1 编写Python脚本

创建一个新文件，命名为photo_colorizer.py，然后添加以下代码：

import streamlit as st from modelscope.pipelines import pipeline from modelscope.utils.constant import Tasks from PIL import Image import cv2 import numpy as np import os # 设置页面标题和布局 st.set_page_config(page_title="AI老照片上色器", layout="wide") st.title("AI智能老照片上色修复工具")

3.2 加载AI模型

添加模型加载代码：

@st.cache_resource def load_colorization_model(): """加载图像上色模型，使用缓存避免重复加载""" # 重要：将下面的路径替换为你自己存放模型的真实路径 model_dir = 'D:/ai_models/image_colorization/iic/cv_unet_image-colorization' # 示例路径，请修改 try: colorizer = pipeline(Tasks.image_colorization, model=model_dir) st.success("AI上色模型加载成功！") return colorizer except Exception as e: st.error(f"模型加载失败，请检查模型路径是否正确: {e}") return None colorizer = load_colorization_model()

注意：务必将model_dir的路径改为你实际存放模型的路径。

3.3 构建用户界面

添加界面代码：

# 侧边栏控制面板 with st.sidebar: st.header("上传图片") uploaded_file = st.file_uploader("选择一张黑白或老照片", type=['jpg', 'jpeg', 'png']) st.header("⚙ 操作") col1, col2 = st.columns(2) with col1: process_btn = st.button("开始上色", type="primary", use_container_width=True) with col2: clear_btn = st.button("清除所有", use_container_width=True) st.markdown("---") st.caption("提示：支持JPG、PNG格式。图片越清晰，上色效果通常越好。") # 主展示区 col_left, col_right = st.columns(2) # 初始化session_state来保存图片状态 if 'original_image' not in st.session_state: st.session_state.original_image = None if 'colored_image' not in st.session_state: st.session_state.colored_image = None # 处理"清除"按钮 if clear_btn: st.session_state.original_image = None st.session_state.colored_image = None st.rerun()

3.4 实现上色功能

最后添加核心功能代码：

# 处理上传的图片 if uploaded_file is not None: image = Image.open(uploaded_file).convert('RGB') st.session_state.original_image = image with col_left: st.subheader("原始图片") if st.session_state.original_image: st.image(st.session_state.original_image, use_column_width=True) else: st.info("请从左侧上传一张黑白或老照片。") # 处理"开始上色"按钮 with col_right: st.subheader("AI上色结果") if process_btn and st.session_state.original_image and colorizer: with st.spinner('AI正在为照片注入色彩，请稍候...'): try: img_np = np.array(st.session_state.original_image) result = colorizer(img_np) colored_np = result['output_img'] colored_pil = Image.fromarray(colored_np.astype('uint8')) st.session_state.colored_image = colored_pil st.success("上色完成！") except Exception as e: st.error(f"上色过程中出现错误: {e}") if st.session_state.colored_image: st.image(st.session_state.colored_image, use_column_width=True) # 提供下载按钮 buf = st.session_state.colored_image st.download_button( label="下载彩色图片", data=buf, file_name="colored_photo.png", mime="image/png", use_container_width=True ) else: st.info("点击「开始上色」按钮，查看AI修复后的彩色效果。")