cv_unet_image-colorization部署教程:Docker Compose编排多用户Web服务方案
cv_unet_image-colorization部署教程:Docker Compose编排多用户Web服务方案
1. 引言
你是否遇到过这样的场景?家里有一堆珍贵的黑白老照片,想给它们上色却苦于没有专业的PS技能;或者作为一个开发者,想为你的社区或团队提供一个便捷的AI图像处理工具,却不知道如何快速部署和管理。
今天我要分享的,就是如何用Docker Compose一键部署一个基于UNet架构的AI图像上色Web服务。这个方案最大的好处是,你不需要懂复杂的深度学习框架,也不需要手动配置Python环境,只需要几条简单的命令,就能拥有一个支持多用户访问、稳定可靠的图像上色平台。
这个方案基于阿里魔搭开源的cv_unet_image-colorization模型,它能够智能识别黑白照片中的内容,并填充自然和谐的色彩。通过Docker Compose编排,我们可以轻松管理整个服务栈,包括Web界面、模型服务、文件存储等组件。
2. 项目架构与核心组件
2.1 整体架构设计
在开始部署之前,我们先了解一下整个系统的架构。这个方案采用了微服务化的设计思路,将不同的功能模块拆分成独立的容器,通过Docker Compose统一管理。
整个系统包含三个核心组件:
- Web前端服务:基于Streamlit构建的用户界面,提供图片上传、上色处理、结果展示和下载功能。
- 模型推理服务:专门负责运行UNet图像上色模型,处理图像上色任务。
- 共享存储卷:用于在不同容器之间共享模型文件和处理结果。
这种架构的好处很明显:各个组件职责清晰,可以独立升级和维护;资源隔离,一个组件出问题不会影响整个系统;扩展性强,未来可以轻松添加新的功能模块。
2.2 核心模型介绍
这个方案的核心是cv_unet_image-colorization模型,它基于UNet架构。UNet是一种在图像分割任务中表现优异的神经网络结构,它的特点是有一个对称的编码器-解码器设计。
简单来说,编码器部分负责理解图片的“大局”——识别出图片中有天空、草地、人物等元素;解码器部分则负责“精雕细琢”——为不同的区域填充合适的颜色。这种结构让模型既能把握整体色调,又能处理好细节边缘。
模型在海量的彩色-黑白配对图片上训练过,学会了各种物体的常见颜色。比如它知道天空通常是蓝色的(但也有可能是傍晚的橙红色),草地是绿色的,肤色是暖色调的。这种“常识”让它上色的结果看起来自然和谐。
3. 环境准备与快速部署
3.1 准备工作
在开始部署之前,你需要确保本地环境满足以下要求:
- 操作系统:Linux(Ubuntu/CentOS)、macOS或Windows(需要WSL2)
- Docker:版本20.10.0或更高
- Docker Compose:版本2.0.0或更高
- 磁盘空间:至少5GB可用空间(主要用于存放模型文件)
- 内存:建议8GB或以上
- 显卡:可选,有NVIDIA GPU可以加速处理
如果你还没有安装Docker和Docker Compose,可以按照官方文档进行安装。这里以Ubuntu系统为例:
# 安装Docker
sudo apt-get update
sudo apt-get install docker.io
# 安装Docker Compose
sudo curl -L "https://github.com/docker/compose/releases/download/v2.20.0/docker-compose-$(uname -s)-$(uname -m)" -o /usr/local/bin/docker-compose
sudo chmod +x /usr/local/bin/docker-compose
# 验证安装
docker --version
docker-compose --version
3.2 项目文件结构
创建一个项目目录,并组织好所需的文件:
cv-unet-colorization/
├── docker-compose.yml # Docker Compose配置文件
├── app/
│ ├── app.py # Streamlit应用主文件
│ ├── requirements.txt # Python依赖包
│ └── models/ # 模型文件目录(从ModelScope下载)
├── nginx/
│ └── nginx.conf # Nginx反向代理配置
└── README.md # 项目说明文档
3.3 Docker Compose配置文件
这是整个部署方案的核心文件,它定义了所有服务的配置和关系:
version: '3.8'
services:
# Web前端服务
web:
build: ./app
container_name: colorization-web
ports:
- "8501:8501"
volumes:
- ./app/models:/app/models
- ./app/uploads:/app/uploads
- ./app/results:/app/results
environment:
- MODEL_PATH=/app/models/cv_unet_image-colorization
- UPLOAD_FOLDER=/app/uploads
- RESULT_FOLDER=/app/results
depends_on:
- model-service
restart: unless-stopped
networks:
- colorization-network
# 模型推理服务(可选,用于分离推理负载)
model-service:
image: pytorch/pytorch:latest
container_name: colorization-model
volumes:
- ./app/models:/models
- ./app/model_server.py:/app/model_server.py
command: python /app/model_server.py
environment:
- MODEL_PATH=/models/cv_unet_image-colorization
ports:
- "5000:5000"
restart: unless-stopped
networks:
- colorization-network
# 如果有GPU,可以取消下面的注释
# deploy:
# resources:
# reservations:
# devices:
# - driver: nvidia
# count: 1
# capabilities: [gpu]
# Nginx反向代理(用于多用户访问和负载均衡)
nginx:
image: nginx:alpine
container_name: colorization-nginx
ports:
- "80:80"
- "443:443"
volumes:
- ./nginx/nginx.conf:/etc/nginx/nginx.conf
- ./ssl:/etc/nginx/ssl
depends_on:
- web
restart: unless-stopped
networks:
- colorization-network
networks:
colorization-network:
driver: bridge
volumes:
model-data:
upload-data:
result-data:
3.4 Streamlit应用代码
创建app/app.py文件,这是Web界面的核心代码:
import streamlit as st
import cv2
import numpy as np
from PIL import Image
import os
from datetime import datetime
import requests
import json
# 页面配置
st.set_page_config(
page_title="AI图像上色工具",
page_icon="🎨",
layout="wide"
)
# 初始化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 'processing' not in st.session_state:
st.session_state.processing = False
# 侧边栏配置
with st.sidebar:
st.title("🎨 AI图像上色")
st.markdown("---")
# 文件上传
uploaded_file = st.file_uploader(
"上传黑白图片",
type=['jpg', 'jpeg', 'png', 'bmp'],
help="支持JPG、JPEG、PNG、BMP格式"
)
# 处理选项
st.markdown("### 处理选项")
color_intensity = st.slider(
"色彩强度",
min_value=0.5,
max_value=2.0,
value=1.0,
step=0.1,
help="调整上色效果的鲜艳程度"
)
# 操作按钮
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
)
if clear_btn:
st.session_state.original_image = None
st.session_state.colored_image = None
st.rerun()
st.markdown("---")
st.markdown("### 使用说明")
st.info("""
1. 上传黑白或灰度图片
2. 调整色彩强度(可选)
3. 点击"开始上色"按钮
4. 等待处理完成并下载结果
""")
# 主界面
st.title("AI智能图像上色工具")
st.markdown("基于UNet深度学习模型的本地化图像上色解决方案")
# 如果上传了文件,显示预览
if uploaded_file is not None:
# 读取图片
image = Image.open(uploaded_file)
st.session_state.original_image = image
# 显示原图
col1, col2 = st.columns(2)
with col1:
st.subheader("原始图片")
st.image(image, use_column_width=True)
st.caption(f"尺寸: {image.size[0]}×{image.size[1]} | 格式: {image.format}")
# 处理逻辑
if process_btn and st.session_state.original_image is not None:
st.session_state.processing = True
with st.spinner("AI正在为图片上色,请稍候..."):
try:
# 这里应该是实际的模型调用代码
# 为了示例,我们模拟一个处理过程
import time
time.sleep(2) # 模拟处理时间
# 模拟处理结果(实际应该调用模型)
original_np = np.array(st.session_state.original_image.convert('RGB'))
# 这里应该是实际的模型推理代码
# colored_np = model.colorize(original_np)
# 模拟上色效果(实际项目中删除这部分)
if len(original_np.shape) == 2: # 灰度图
colored_np = cv2.cvtColor(original_np, cv2.COLOR_GRAY2BGR)
else: # 已经是彩色但饱和度低
colored_np = original_np.copy()
# 调整色彩强度(模拟)
hsv = cv2.cvtColor(colored_np, cv2.COLOR_BGR2HSV)
hsv[:, :, 1] = np.clip(hsv[:, :, 1] * color_intensity, 0, 255)
colored_np = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR)
# 转换为PIL Image
colored_image = Image.fromarray(colored_np)
st.session_state.colored_image = colored_image
st.session_state.processing = False
st.success("上色完成!")
except Exception as e:
st.session_state.processing = False
st.error(f"处理失败: {str(e)}")
# 显示处理结果
if st.session_state.colored_image is not None:
col1, col2 = st.columns(2)
with col1:
st.subheader("原始图片")
st.image(st.session_state.original_image, use_column_width=True)
with col2:
st.subheader("AI上色结果")
st.image(st.session_state.colored_image, use_column_width=True)
# 下载按钮
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
colored_bytes = st.session_state.colored_image.tobytes()
st.download_button(
label="📥 下载彩色图片",
data=colored_bytes,
file_name=f"colored_{timestamp}.png",
mime="image/png",
use_container_width=True
)
# 如果没有上传图片,显示示例
elif not st.session_state.processing:
st.markdown("---")
st.subheader("示例效果")
example_col1, example_col2, example_col3 = st.columns(3)
with example_col1:
st.image("https://via.placeholder.com/300x200/cccccc/969696?text=黑白风景",
caption="黑白风景 → 彩色风景")
with example_col2:
st.image("https://via.placeholder.com/300x200/cccccc/969696?text=老照片",
caption="黑白人像 → 彩色人像")
with example_col3:
st.image("https://via.placeholder.com/300x200/cccccc/969696?text=建筑",
caption="黑白建筑 → 彩色建筑")
# 页脚信息
st.markdown("---")
st.markdown("""
**技术说明**:
- 基于UNet深度学习模型
- 本地处理,保护隐私
- 支持GPU加速
- 开源模型来自阿里魔搭ModelScope
""")
3.5 一键部署命令
准备好所有文件后,只需要一条命令就能启动整个服务:
# 进入项目目录
cd cv-unet-colorization
# 下载模型文件(需要提前从ModelScope下载)
# 将模型文件放到 app/models/cv_unet_image-colorization/ 目录下
# 启动所有服务
docker-compose up -d
# 查看服务状态
docker-compose ps
# 查看日志
docker-compose logs -f web
服务启动后,你可以通过以下方式访问:
- Web界面:http://localhost:8501
- 模型API:http://localhost:5000
- 如果配置了Nginx:http://your-server-ip
4. 多用户Web服务配置
4.1 Nginx反向代理配置
为了让服务支持多用户同时访问,我们需要配置Nginx作为反向代理。创建nginx/nginx.conf文件:
events {
worker_connections 1024;
}
http {
upstream streamlit_servers {
# 可以配置多个Streamlit实例实现负载均衡
server web:8501;
# server web2:8501; # 如果有多个实例
# server web3:8501;
}
server {
listen 80;
server_name your-domain.com; # 替换为你的域名或IP
# 静态文件缓存
location /static {
alias /app/static;
expires 1y;
add_header Cache-Control "public, immutable";
}
# Streamlit代理配置
location / {
proxy_pass http://streamlit_servers;
proxy_http_version 1.1;
proxy_set_header Upgrade $http_upgrade;
proxy_set_header Connection "upgrade";
proxy_set_header Host $host;
proxy_set_header X-Real-IP $remote_addr;
proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
proxy_set_header X-Forwarded-Proto $scheme;
# 超时设置
proxy_connect_timeout 300s;
proxy_send_timeout 300s;
proxy_read_timeout 300s;
# 禁用缓冲
proxy_buffering off;
proxy_request_buffering off;
}
# 健康检查
location /health {
proxy_pass http://streamlit_servers/health;
access_log off;
}
}
}
4.2 会话管理与状态保持
在多用户场景下,需要确保每个用户的会话状态独立。Streamlit本身支持多会话,但我们需要做一些额外配置:
# 在app.py中添加会话管理
import streamlit as st
from streamlit.runtime.scriptrunner import get_script_run_ctx
def get_session_id():
"""获取当前会话ID"""
ctx = get_script_run_ctx()
if ctx is not None:
return ctx.session_id
return None
# 使用会话ID作为缓存键
@st.cache_resource(show_spinner=False)
def load_model():
"""加载模型,每个会话独立缓存"""
session_id = get_session_id()
# 这里加载模型,可以使用会话ID来区分不同用户的模型实例
# 实际项目中可能需要根据会话管理模型实例
return model
# 文件上传处理,确保文件路径唯一
def save_uploaded_file(uploaded_file, session_id):
"""保存上传的文件,使用会话ID避免冲突"""
import uuid
filename = f"{session_id}_{uuid.uuid4().hex}_{uploaded_file.name}"
filepath = os.path.join(UPLOAD_FOLDER, filename)
with open(filepath, "wb") as f:
f.write(uploaded_file.getbuffer())
return filepath
4.3 资源限制与优化
为了防止单个用户占用过多资源,我们可以配置Docker资源限制:
# 在docker-compose.yml中添加资源限制
services:
web:
# ... 其他配置 ...
deploy:
resources:
limits:
cpus: '1' # 限制CPU使用
memory: 2G # 限制内存使用
reservations:
cpus: '0.5'
memory: 1G
model-service:
# ... 其他配置 ...
deploy:
resources:
limits:
cpus: '2'
memory: 4G
# 如果有GPU,可以限制GPU使用
# devices:
# - capabilities: [gpu]
# count: 1
5. 高级功能与优化建议
5.1 批量处理功能
对于需要处理大量图片的用户,可以添加批量处理功能:
# 批量处理功能示例
def batch_process_images(image_files, output_dir):
"""批量处理图片"""
results = []
progress_bar = st.progress(0)
status_text = st.empty()
for i, image_file in enumerate(image_files):
# 更新进度
progress = (i + 1) / len(image_files)
progress_bar.progress(progress)
status_text.text(f"处理中: {i+1}/{len(image_files)}")
try:
# 处理单张图片
result = process_single_image(image_file)
# 保存结果
output_path = os.path.join(output_dir, f"result_{i}.png")
result.save(output_path)
results.append({
'original': image_file.name,
'result': output_path,
'success': True
})
except Exception as e:
results.append({
'original': image_file.name,
'error': str(e),
'success': False
})
progress_bar.empty()
status_text.empty()
return results
# 在Streamlit界面中添加批量上传
batch_files = st.file_uploader(
"批量上传图片",
type=['jpg', 'jpeg', 'png', 'bmp'],
accept_multiple_files=True,
help="支持多选,最多20张图片"
)
if batch_files and len(batch_files) > 0:
if st.button("🚀 批量处理"):
results = batch_process_images(batch_files, "./results")
# 显示处理结果
st.dataframe(pd.DataFrame(results))
# 提供批量下载
if st.button("📦 打包下载所有结果"):
# 创建ZIP文件
import zipfile
from io import BytesIO
zip_buffer = BytesIO()
with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
for result in results:
if result['success']:
zip_file.write(result['result'])
st.download_button(
label="下载ZIP包",
data=zip_buffer.getvalue(),
file_name="batch_results.zip",
mime="application/zip"
)
5.2 模型性能优化
为了提高处理速度,可以考虑以下优化措施:
# 模型推理优化示例
import torch
from torch import nn
import onnxruntime as ort
class OptimizedColorizationModel:
def __init__(self, model_path, use_gpu=True):
self.use_gpu = use_gpu and torch.cuda.is_available()
# 加载PyTorch模型
self.model = self.load_pytorch_model(model_path)
# 转换为ONNX格式以获得更好的性能
self.onnx_session = self.convert_to_onnx(model_path)
# 启用半精度推理(如果支持)
if self.use_gpu:
self.model.half() # 转换为半精度
def load_pytorch_model(self, model_path):
"""加载PyTorch模型"""
model = torch.load(model_path, map_location='cpu')
model.eval()
if self.use_gpu:
model.cuda()
return model
def convert_to_onnx(self, model_path):
"""转换为ONNX格式"""
onnx_path = model_path.replace('.pth', '.onnx')
if not os.path.exists(onnx_path):
# 导出为ONNX
dummy_input = torch.randn(1, 1, 256, 256)
torch.onnx.export(
self.model,
dummy_input,
onnx_path,
opset_version=11,
input_names=['input'],
output_names=['output']
)
# 创建ONNX Runtime会话
providers = ['CUDAExecutionProvider'] if self.use_gpu else ['CPUExecutionProvider']
session = ort.InferenceSession(onnx_path, providers=providers)
return session
def inference(self, image):
"""推理函数,自动选择最优后端"""
if self.use_gpu and image.shape[0] > 1: # 批量处理使用ONNX
return self.onnx_inference(image)
else: # 单张图片使用PyTorch
return self.pytorch_inference(image)
def pytorch_inference(self, image):
"""PyTorch推理"""
with torch.no_grad():
if self.use_gpu:
image = image.cuda()
if hasattr(self.model, 'half'):
image = image.half()
output = self.model(image)
return output.cpu()
def onnx_inference(self, image):
"""ONNX推理"""
input_name = self.onnx_session.get_inputs()[0].name
output_name = self.onnx_session.get_outputs()[0].name
# 准备输入
if self.use_gpu:
image = image.cpu().numpy()
else:
image = image.numpy()
# 推理
result = self.onnx_session.run(
[output_name],
{input_name: image}
)
return torch.from_numpy(result[0])
5.3 监控与日志
为了确保服务稳定运行,需要添加监控和日志功能:
# 在docker-compose.yml中添加监控服务
services:
# ... 其他服务 ...
prometheus:
image: prom/prometheus:latest
container_name: colorization-prometheus
volumes:
- ./monitoring/prometheus.yml:/etc/prometheus/prometheus.yml
- prometheus-data:/prometheus
command:
- '--config.file=/etc/prometheus/prometheus.yml'
- '--storage.tsdb.path=/prometheus'
- '--web.console.libraries=/etc/prometheus/console_libraries'
- '--web.console.templates=/etc/prometheus/consoles'
- '--storage.tsdb.retention.time=200h'
- '--web.enable-lifecycle'
ports:
- "9090:9090"
networks:
- colorization-network
restart: unless-stopped
grafana:
image: grafana/grafana:latest
container_name: colorization-grafana
volumes:
- grafana-data:/var/lib/grafana
- ./monitoring/dashboards:/etc/grafana/provisioning/dashboards
- ./monitoring/datasources:/etc/grafana/provisioning/datasources
environment:
- GF_SECURITY_ADMIN_PASSWORD=admin
- GF_USERS_ALLOW_SIGN_UP=false
ports:
- "3000:3000"
networks:
- colorization-network
restart: unless-stopped
depends_on:
- prometheus
# 添加日志收集
loki:
image: grafana/loki:latest
container_name: colorization-loki
ports:
- "3100:3100"
command: -config.file=/etc/loki/local-config.yaml
networks:
- colorization-network
promtail:
image: grafana/promtail:latest
container_name: colorization-promtail
volumes:
- /var/log:/var/log
- /var/lib/docker/containers:/var/lib/docker/containers:ro
- ./monitoring/promtail-config.yml:/etc/promtail/config.yml
command: -config.file=/etc/promtail/config.yml
networks:
- colorization-network
depends_on:
- loki
6. 常见问题与解决方案
6.1 部署问题排查
问题1:Docker Compose启动失败
# 查看详细错误信息
docker-compose logs
# 检查端口冲突
netstat -tulpn | grep :8501
# 重新构建镜像
docker-compose build --no-cache
# 清理旧容器和镜像
docker-compose down -v
docker system prune -a
问题2:模型加载失败
# 检查模型文件路径
docker exec -it colorization-web ls -la /app/models/
# 检查模型文件权限
docker exec -it colorization-web chmod -R 755 /app/models
# 查看模型加载日志
docker-compose logs model-service
问题3:内存不足
# 调整docker-compose.yml中的资源限制
services:
web:
deploy:
resources:
limits:
memory: 1G # 降低内存限制
6.2 性能优化建议
-
启用GPU加速:
- 确保主机有NVIDIA GPU并安装了正确的驱动
- 安装NVIDIA Container Toolkit
- 在docker-compose.yml中取消GPU相关注释
-
使用模型缓存:
@st.cache_resource(ttl=3600) # 缓存1小时 def get_model(): return load_model() -
图片预处理优化:
def optimize_image_processing(image, target_size=512): """优化图片处理流程""" # 调整尺寸 if max(image.size) > target_size: ratio = target_size / max(image.size) new_size = tuple(int(dim * ratio) for dim in image.size) image = image.resize(new_size, Image.Resampling.LANCZOS) # 转换为合适的格式 if image.mode != 'RGB': image = image.convert('RGB') return image
6.3 安全加固措施
-
添加身份验证:
# 简单的密码保护 import streamlit_authenticator as stauth credentials = { 'usernames': { 'user1': { 'name': 'User One', 'password': stauth.Hasher(['password123']).generate()[0] } } } authenticator = stauth.Authenticate( credentials, 'app_cookie', 'app_key', cookie_expiry_days=30 ) name, authentication_status, username = authenticator.login('Login', 'main') if authentication_status: # 显示应用内容 pass -
限制文件上传:
# 限制文件大小和类型 MAX_FILE_SIZE = 10 * 1024 * 1024 # 10MB ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'bmp'} def validate_uploaded_file(file): if file.size > MAX_FILE_SIZE: raise ValueError(f"文件大小不能超过{MAX_FILE_SIZE//1024//1024}MB") ext = file.name.split('.')[-1].lower() if ext not in ALLOWED_EXTENSIONS: raise ValueError(f"不支持的文件类型,请上传{ALLOWED_EXTENSIONS}格式") return True
7. 总结
通过Docker Compose部署cv_unet_image-colorization图像上色服务,我们实现了一个完整的多用户Web应用解决方案。这个方案有以下几个关键优势:
部署简单:只需要一条命令就能启动所有服务,无需手动配置复杂的Python环境。
易于维护:每个服务都在独立的容器中运行,更新或修复某个组件不会影响其他服务。
扩展性强:可以根据用户量轻松扩展Web服务实例,通过Nginx实现负载均衡。
资源隔离:每个用户的处理请求都在独立的会话中运行,互不干扰。
隐私安全:所有图片处理都在本地完成,数据不会上传到云端,保护用户隐私。
这个方案不仅适用于个人使用,也可以作为企业内部的图像处理工具,或者作为在线服务提供给更多用户。通过适当的优化和监控,它可以稳定地处理大量的图像上色请求。
未来还可以考虑添加更多功能,比如支持视频上色、批量处理队列、用户管理系统等,让这个工具更加完善和强大。
获取更多AI镜像
想探索更多AI镜像和应用场景?访问 CSDN星图镜像广场,提供丰富的预置镜像,覆盖大模型推理、图像生成、视频生成、模型微调等多个领域,支持一键部署。
更多推荐




所有评论(0)