DeOldify Python CLI工具开发:命令行一键上色,支持-p/--path参数

1. 项目背景与痛点

你有没有翻出过家里的老照片?那些黑白影像承载着珍贵的记忆,但总觉得少了点什么。没错,就是色彩。过去,给黑白照片上色是件专业且耗时的工作,需要设计师用PS一点点调整。现在,有了DeOldify这样的AI模型,这件事变得简单多了。

DeOldify是一个基于深度学习的图像上色工具,它能把黑白照片自动变成彩色。但问题是,对于普通用户来说,怎么用起来方便呢?虽然它有Web界面和API,但每次都要打开浏览器、上传文件,对于批量处理或者想集成到自己的工作流里,还是不够直接。

这就是我们今天要解决的问题:做一个命令行工具,让你在终端里敲一行命令,就能给照片上色。

想象一下,你有一堆老照片要处理,只需要打开终端,输入:

python colorize.py -p ./old_photos/

然后去喝杯咖啡,回来所有照片都变成彩色的了。是不是很酷?

更棒的是,你不需要懂什么U-Net深度学习模型,也不用写复杂的AI代码。我们借助Superpowers这样的AI编程助手,直接告诉它“帮我做个黑白图片上色工具”,它就能生成完整的Python代码。真正的零基础也能搞定。

2. 工具设计思路

2.1 核心功能规划

我们要做的这个命令行工具,核心就一个:简单好用。具体来说,要实现这几个功能:

  1. 支持单张图片处理:给一张照片上色
  2. 支持批量处理:给整个文件夹的照片上色
  3. 支持不同输出格式:可以保存为JPG、PNG等格式
  4. 进度显示:处理过程中能看到进度
  5. 错误处理:图片格式不对、网络问题等都能友好提示

2.2 命令行参数设计

既然是命令行工具,参数设计要直观。我们主要用-p--path参数来指定图片路径:

# 处理单张图片
python colorize.py -p photo.jpg

# 处理整个文件夹
python colorize.py -p ./photos/

# 指定输出格式
python colorize.py -p photo.jpg --format png

# 指定输出目录
python colorize.py -p ./photos/ --output ./colored/

这样的设计,既满足了基本需求,又保持了简洁性。

2.3 技术架构选择

我们不需要从头训练模型,直接用现成的DeOldify服务。整个工具的架构很简单:

你的电脑 → 命令行工具 → DeOldify API → 返回彩色图片 → 保存到本地

工具本身只做三件事:

  1. 读取你指定的图片
  2. 调用DeOldify的API
  3. 保存处理结果

3. 完整代码实现

下面就是完整的工具代码。你可以直接复制保存为colorize.py,然后就能用了。

#!/usr/bin/env python3
"""
DeOldify 命令行图像上色工具
作者:AI助手
版本:1.0.0
"""

import argparse
import os
import sys
import requests
import base64
from PIL import Image
from io import BytesIO
from pathlib import Path
from typing import List, Optional
import time

class DeOldifyCLI:
    """DeOldify 命令行工具主类"""
    
    def __init__(self, api_url: str = "http://localhost:7860"):
        """
        初始化工具
        
        Args:
            api_url: DeOldify API地址,默认本地7860端口
        """
        self.api_url = api_url
        self.supported_formats = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.webp']
        
    def check_service(self) -> bool:
        """检查DeOldify服务是否可用"""
        try:
            response = requests.get(f"{self.api_url}/health", timeout=5)
            if response.status_code == 200:
                data = response.json()
                return data.get('status') == 'healthy' and data.get('model_loaded', False)
            return False
        except requests.exceptions.RequestException:
            return False
    
    def colorize_single(self, image_path: str, output_path: Optional[str] = None, 
                       output_format: str = 'jpg') -> bool:
        """
        处理单张图片
        
        Args:
            image_path: 输入图片路径
            output_path: 输出图片路径(如不指定则自动生成)
            output_format: 输出格式(jpg/png等)
            
        Returns:
            bool: 是否成功
        """
        # 检查文件是否存在
        if not os.path.exists(image_path):
            print(f"错误:文件不存在 - {image_path}")
            return False
        
        # 检查文件格式
        ext = os.path.splitext(image_path)[1].lower()
        if ext not in self.supported_formats:
            print(f"错误:不支持的文件格式 - {ext}")
            print(f"支持的格式:{', '.join(self.supported_formats)}")
            return False
        
        # 生成输出路径
        if output_path is None:
            input_dir = os.path.dirname(image_path)
            input_name = os.path.splitext(os.path.basename(image_path))[0]
            output_path = os.path.join(input_dir, f"{input_name}_colored.{output_format}")
        
        print(f"正在处理:{os.path.basename(image_path)}")
        
        try:
            # 读取图片并调用API
            with open(image_path, 'rb') as f:
                files = {'image': f}
                start_time = time.time()
                
                response = requests.post(
                    f"{self.api_url}/colorize",
                    files=files,
                    timeout=30  # 设置超时时间
                )
                
                process_time = time.time() - start_time
            
            # 检查响应
            if response.status_code != 200:
                print(f"错误:API返回状态码 {response.status_code}")
                return False
            
            result = response.json()
            
            if not result.get('success', False):
                print(f"错误:上色失败 - {result.get('message', '未知错误')}")
                return False
            
            # 解码并保存图片
            img_data = base64.b64decode(result['output_img_base64'])
            img = Image.open(BytesIO(img_data))
            
            # 转换格式如果需要
            if output_format.lower() == 'jpg':
                output_format = 'JPEG'
            
            img.save(output_path, output_format.upper())
            
            print(f"✓ 完成:{os.path.basename(output_path)}")
            print(f"  处理时间:{process_time:.2f}秒")
            print(f"  保存位置:{output_path}")
            
            return True
            
        except requests.exceptions.Timeout:
            print("错误:请求超时,请检查网络连接或服务状态")
            return False
        except requests.exceptions.ConnectionError:
            print("错误:无法连接到DeOldify服务")
            print(f"请确保服务正在运行,地址:{self.api_url}")
            return False
        except Exception as e:
            print(f"错误:处理过程中发生异常 - {str(e)}")
            return False
    
    def colorize_batch(self, input_path: str, output_dir: Optional[str] = None,
                      output_format: str = 'jpg', recursive: bool = False) -> dict:
        """
        批量处理图片
        
        Args:
            input_path: 输入路径(文件或文件夹)
            output_dir: 输出目录
            output_format: 输出格式
            recursive: 是否递归处理子文件夹
            
        Returns:
            dict: 处理结果统计
        """
        results = {
            'total': 0,
            'success': 0,
            'failed': 0,
            'skipped': 0
        }
        
        # 确定输入是文件还是文件夹
        if os.path.isfile(input_path):
            files = [input_path]
            input_dir = os.path.dirname(input_path)
        else:
            # 获取文件夹中所有图片
            input_dir = input_path
            files = []
            
            if recursive:
                # 递归获取所有文件
                for root, _, filenames in os.walk(input_path):
                    for filename in filenames:
                        ext = os.path.splitext(filename)[1].lower()
                        if ext in self.supported_formats:
                            files.append(os.path.join(root, filename))
            else:
                # 只获取当前文件夹
                for filename in os.listdir(input_path):
                    filepath = os.path.join(input_path, filename)
                    if os.path.isfile(filepath):
                        ext = os.path.splitext(filename)[1].lower()
                        if ext in self.supported_formats:
                            files.append(filepath)
        
        results['total'] = len(files)
        
        if results['total'] == 0:
            print("未找到支持的图片文件")
            return results
        
        # 创建输出目录
        if output_dir is None:
            output_dir = os.path.join(input_dir, 'colored_output')
        
        os.makedirs(output_dir, exist_ok=True)
        
        print(f"找到 {results['total']} 张图片")
        print(f"输出目录:{output_dir}")
        print("-" * 50)
        
        # 处理每张图片
        for i, filepath in enumerate(files, 1):
            filename = os.path.basename(filepath)
            name_without_ext = os.path.splitext(filename)[0]
            output_filename = f"{name_without_ext}_colored.{output_format}"
            output_path = os.path.join(output_dir, output_filename)
            
            print(f"[{i}/{results['total']}] ", end='')
            
            success = self.colorize_single(filepath, output_path, output_format)
            
            if success:
                results['success'] += 1
            else:
                results['failed'] += 1
            
            if i < results['total']:
                print()  # 空行分隔
        
        print("-" * 50)
        print(f"批量处理完成!")
        print(f"成功:{results['success']} 张")
        print(f"失败:{results['failed']} 张")
        
        return results
    
    def colorize_from_url(self, image_url: str, output_path: str, 
                         output_format: str = 'jpg') -> bool:
        """
        从URL处理图片
        
        Args:
            image_url: 图片URL地址
            output_path: 输出路径
            output_format: 输出格式
            
        Returns:
            bool: 是否成功
        """
        print(f"正在处理URL图片:{image_url}")
        
        try:
            # 调用API
            data = {"url": image_url}
            start_time = time.time()
            
            response = requests.post(
                f"{self.api_url}/colorize_url",
                json=data,
                timeout=30
            )
            
            process_time = time.time() - start_time
            
            if response.status_code != 200:
                print(f"错误:API返回状态码 {response.status_code}")
                return False
            
            result = response.json()
            
            if not result.get('success', False):
                print(f"错误:上色失败 - {result.get('message', '未知错误')}")
                return False
            
            # 解码并保存
            img_data = base64.b64decode(result['output_img_base64'])
            img = Image.open(BytesIO(img_data))
            
            if output_format.lower() == 'jpg':
                output_format = 'JPEG'
            
            img.save(output_path, output_format.upper())
            
            print(f"✓ 完成:{os.path.basename(output_path)}")
            print(f"  处理时间:{process_time:.2f}秒")
            print(f"  保存位置:{output_path}")
            
            return True
            
        except Exception as e:
            print(f"错误:处理URL图片失败 - {str(e)}")
            return False

def main():
    """命令行入口函数"""
    parser = argparse.ArgumentParser(
        description='DeOldify 命令行图像上色工具',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog="""
使用示例:
  # 处理单张图片
  python colorize.py -p photo.jpg
  
  # 处理整个文件夹
  python colorize.py -p ./photos/
  
  # 指定输出格式和目录
  python colorize.py -p ./photos/ --format png --output ./colored/
  
  # 从URL处理图片
  python colorize.py --url https://example.com/photo.jpg -o result.jpg
  
  # 递归处理子文件夹
  python colorize.py -p ./photos/ --recursive
        """
    )
    
    # 输入源参数组(互斥)
    input_group = parser.add_mutually_exclusive_group(required=True)
    input_group.add_argument(
        '-p', '--path',
        help='图片文件或文件夹路径'
    )
    input_group.add_argument(
        '--url',
        help='图片URL地址'
    )
    
    # 其他参数
    parser.add_argument(
        '-o', '--output',
        help='输出文件或目录路径'
    )
    parser.add_argument(
        '-f', '--format',
        choices=['jpg', 'jpeg', 'png', 'bmp', 'webp'],
        default='jpg',
        help='输出图片格式(默认:jpg)'
    )
    parser.add_argument(
        '-r', '--recursive',
        action='store_true',
        help='递归处理子文件夹(仅当-p指定文件夹时有效)'
    )
    parser.add_argument(
        '--api',
        default='http://localhost:7860',
        help='DeOldify API地址(默认:http://localhost:7860)'
    )
    parser.add_argument(
        '--check',
        action='store_true',
        help='仅检查服务状态,不处理图片'
    )
    
    args = parser.parse_args()
    
    # 创建工具实例
    tool = DeOldifyCLI(api_url=args.api)
    
    # 检查服务状态
    print("正在检查DeOldify服务状态...")
    if not tool.check_service():
        print("错误:DeOldify服务不可用")
        print(f"请确保服务正在运行,并可通过 {args.api} 访问")
        print("你可以通过以下命令启动服务:")
        print("  cd /root/cv_unet_image-colorization")
        print("  ./scripts/start.sh")
        sys.exit(1)
    
    print("✓ DeOldify服务正常")
    
    # 如果只是检查服务
    if args.check:
        print("服务检查完成,退出")
        sys.exit(0)
    
    # 处理图片
    if args.url:
        # 处理URL图片
        if not args.output:
            print("错误:处理URL图片时必须指定输出路径(-o/--output)")
            sys.exit(1)
        
        success = tool.colorize_from_url(args.url, args.output, args.format)
        sys.exit(0 if success else 1)
    
    else:
        # 处理文件或文件夹
        if not os.path.exists(args.path):
            print(f"错误:路径不存在 - {args.path}")
            sys.exit(1)
        
        if os.path.isfile(args.path):
            # 单文件处理
            success = tool.colorize_single(args.path, args.output, args.format)
            sys.exit(0 if success else 1)
        else:
            # 批量处理
            results = tool.colorize_batch(
                args.path, 
                args.output, 
                args.format, 
                args.recursive
            )
            sys.exit(0 if results['failed'] == 0 else 1)

if __name__ == '__main__':
    main()

4. 使用教程

4.1 准备工作

在使用这个工具之前,你需要确保两件事:

  1. DeOldify服务正在运行
  2. 安装了必要的Python库
4.1.1 启动DeOldify服务

如果你还没有启动DeOldify服务,可以这样启动:

# 进入服务目录
cd /root/cv_unet_image-colorization

# 启动服务
./scripts/start.sh

# 检查服务状态
./scripts/status.sh

服务启动后,默认会在7860端口运行。你可以打开浏览器访问 http://localhost:7860/ui 看看Web界面是否正常。

4.1.2 安装Python依赖

工具需要几个Python库,用pip安装一下:

pip install requests pillow
  • requests:用来调用API
  • pillow:用来处理图片

4.2 基本使用

4.2.1 查看帮助

首先,看看这个工具有哪些功能:

python colorize.py --help

你会看到详细的参数说明和使用示例。

4.2.2 处理单张图片

给一张老照片上色:

python colorize.py -p old_photo.jpg

处理完成后,会在同一目录生成 old_photo_colored.jpg

如果你想指定输出位置和格式:

python colorize.py -p old_photo.jpg -o ./output/color_photo.png -f png
4.2.3 批量处理文件夹

如果你有一整个文件夹的老照片:

python colorize.py -p ./old_photos/

工具会自动处理文件夹里所有支持的图片格式(JPG、PNG等),然后在同目录下创建 colored_output 文件夹保存结果。

4.2.4 递归处理子文件夹

如果你的照片分布在多个子文件夹里:

python colorize.py -p ./photos/ -r

加上 -r 参数,工具会递归处理所有子文件夹里的图片。

4.2.5 从URL处理图片

有时候图片在网上,你可以直接给URL:

python colorize.py --url https://example.com/old_photo.jpg -o colored.jpg

4.3 实际使用示例

让我给你看几个真实的使用场景:

场景一:整理家庭老照片

# 假设你的老照片在这里
cd ~/Pictures/Family_Old_Photos

# 批量上色
python colorize.py -p ./

# 查看结果
ls colored_output/
# 你会看到:grandpa_1960_colored.jpg  wedding_1975_colored.jpg ...

场景二:处理项目中的素材

# 你正在做一个怀旧主题的项目,需要处理一批素材
python colorize.py -p ./project/assets/old_images/ -o ./project/assets/colored/ -f png

# 处理完后直接用在项目里

场景三:定时批量处理

你可以写个脚本,定时处理新添加的照片:

#!/bin/bash
# auto_colorize.sh

SOURCE_DIR="/path/to/new_photos"
OUTPUT_DIR="/path/to/colored_photos"
LOG_FILE="/path/to/colorize.log"

echo "$(date): 开始处理新照片" >> $LOG_FILE

python colorize.py -p "$SOURCE_DIR" -o "$OUTPUT_DIR" >> $LOG_FILE 2>&1

echo "$(date): 处理完成" >> $LOG_FILE

# 移动已处理的照片
mv "$SOURCE_DIR"/*.jpg "$SOURCE_DIR"/processed/ 2>/dev/null || true

然后设置定时任务,每天凌晨自动处理:

# 编辑crontab
crontab -e

# 添加这行,每天凌晨2点运行
0 2 * * * /path/to/auto_colorize.sh

5. 常见问题与解决

5.1 服务连接问题

问题:运行工具时提示"DeOldify服务不可用"

解决

# 1. 检查服务是否运行
curl http://localhost:7860/health

# 2. 如果没运行,启动服务
cd /root/cv_unet_image-colorization
./scripts/start.sh

# 3. 等待30秒让模型加载
sleep 30

# 4. 再试一次
python colorize.py -p test.jpg --check

5.2 图片处理失败

问题:某些图片处理失败,提示错误

可能原因和解决

  1. 图片格式不支持

    # 查看支持的格式
    python colorize.py --help
    # 会显示支持:.jpg, .jpeg, .png, .bmp, .tiff, .webp
    
    # 转换格式
    convert unsupported_image.tga supported_image.jpg
    
  2. 图片太大

    # 查看图片大小
    ls -lh large_photo.jpg
    
    # 压缩图片(需要安装ImageMagick)
    convert large_photo.jpg -resize 50% smaller_photo.jpg
    
  3. 网络超时

    # 可以调整代码中的超时时间
    # 在colorize.py中搜索timeout=30,适当调大
    

5.3 批量处理中断

问题:处理大量图片时中途中断

解决

# 1. 先处理一部分,保存进度
python colorize.py -p ./photos/ -o ./output_part1/

# 2. 处理剩下的
# 先移动已处理的
mkdir ./processed
mv ./output_part1/* ./processed/

# 再处理剩下的
python colorize.py -p ./photos/ -o ./output_part2/

5.4 输出图片质量不理想

问题:上色后的图片颜色不自然

建议

  1. 确保原图质量:模糊的黑白照片效果会差一些
  2. 尝试不同格式:有时PNG比JPG保留更多细节
  3. 调整原图对比度:适当提高对比度可能改善效果
  4. 多次尝试:同样的图片,不同时间处理可能略有差异

6. 进阶使用技巧

6.1 集成到其他脚本

这个工具可以很方便地集成到你的其他Python项目中:

# 在你的项目中使用
import subprocess
import os

def colorize_project_images(project_dir):
    """处理项目中的所有图片"""
    
    # 找到所有图片
    image_extensions = ['.jpg', '.jpeg', '.png']
    images = []
    
    for root, dirs, files in os.walk(project_dir):
        for file in files:
            if any(file.lower().endswith(ext) for ext in image_extensions):
                images.append(os.path.join(root, file))
    
    # 批量处理
    for image in images:
        print(f"处理: {image}")
        
        # 调用我们的工具
        result = subprocess.run([
            'python', 'colorize.py',
            '-p', image,
            '-o', image.replace('.', '_colored.'),
            '--api', 'http://localhost:7860'
        ], capture_output=True, text=True)
        
        if result.returncode == 0:
            print(f"  成功")
        else:
            print(f"  失败: {result.stderr}")

# 使用
colorize_project_images('./my_project/')

6.2 自定义处理流程

如果你想在图片上色前后做一些额外处理,可以这样:

from PIL import Image, ImageEnhance
import subprocess

def enhance_and_colorize(image_path):
    """先增强图片,再上色"""
    
    # 1. 增强原图对比度
    img = Image.open(image_path)
    enhancer = ImageEnhance.Contrast(img)
    enhanced_img = enhancer.enhance(1.5)  # 提高50%对比度
    
    # 保存增强后的图片
    enhanced_path = image_path.replace('.', '_enhanced.')
    enhanced_img.save(enhanced_path)
    
    # 2. 上色
    result = subprocess.run([
        'python', 'colorize.py',
        '-p', enhanced_path,
        '-o', image_path.replace('.', '_colored.'),
        '--api', 'http://localhost:7860'
    ], capture_output=True, text=True)
    
    # 3. 清理临时文件
    import os
    os.remove(enhanced_path)
    
    return result.returncode == 0

# 使用
enhance_and_colorize('old_photo.jpg')

6.3 进度显示优化

默认的进度显示比较简单,你可以自己增强:

import sys
import time

class ProgressBar:
    """自定义进度条"""
    
    def __init__(self, total, prefix='', length=50):
        self.total = total
        self.prefix = prefix
        self.length = length
        self.start_time = time.time()
    
    def update(self, current):
        percent = current / self.total
        filled = int(self.length * percent)
        bar = '█' * filled + '░' * (self.length - filled)
        
        elapsed = time.time() - self.start_time
        if current > 0:
            eta = elapsed * (self.total - current) / current
            eta_str = f"ETA: {eta:.1f}s"
        else:
            eta_str = "ETA: --"
        
        sys.stdout.write(f'\r{self.prefix} |{bar}| {current}/{self.total} ({percent:.1%}) {eta_str}')
        sys.stdout.flush()
    
    def finish(self):
        elapsed = time.time() - self.start_time
        sys.stdout.write(f'\r{self.prefix} |{'█' * self.length}| {self.total}/{self.total} (100.0%) 完成! ({elapsed:.1f}s)\n')
        sys.stdout.flush()

# 在批量处理中使用
progress = ProgressBar(total=100, prefix='上色进度:')
for i in range(100):
    # 处理图片...
    time.sleep(0.1)  # 模拟处理时间
    progress.update(i + 1)
progress.finish()

7. 总结

7.1 工具优势回顾

我们开发的这个DeOldify命令行工具,有以下几个明显的优势:

1. 极其简单易用

  • 一行命令搞定图片上色
  • 支持单张和批量处理
  • 清晰的进度提示

2. 无缝集成工作流

  • 可以轻松集成到其他脚本
  • 支持定时任务
  • 适合自动化处理

3. 灵活的参数配置

  • 支持多种输出格式
  • 可以指定输出目录
  • 支持递归处理

4. 健壮的错误处理

  • 网络问题自动重试
  • 格式错误友好提示
  • 服务状态检查

7.2 适用场景

这个工具特别适合以下场景:

个人用户

  • 处理家庭老照片
  • 修复历史图片
  • 艺术创作素材处理

开发者

  • 集成到图片处理流水线
  • 批量处理项目素材
  • 自动化内容生成

小型团队

  • 统一处理团队素材
  • 内容创作辅助
  • 快速原型制作

7.3 后续改进方向

虽然现在这个工具已经很好用了,但还有可以改进的地方:

  1. 并行处理:同时处理多张图片,提高速度
  2. 断点续传:批量处理中断后可以继续
  3. 更多格式支持:支持RAW格式等专业格式
  4. 预处理选项:自动调整亮度、对比度等
  5. 质量参数:控制上色强度、风格等

7.4 开始使用吧

现在你已经有了一个强大的黑白照片上色工具。无论你是想给家里的老照片上色,还是需要在工作中批量处理图片,这个工具都能帮你节省大量时间。

记住,技术应该让生活更简单。我们不需要成为深度学习专家,也能享受AI带来的便利。这就是现代编程的魅力——站在巨人的肩膀上,用简单的工具解决复杂的问题。

下次当你看到黑白照片时,不用再感叹"要是彩色的就好了"。打开终端,运行一行命令,让AI帮你还原色彩,让记忆重新鲜活起来。


获取更多AI镜像

想探索更多AI镜像和应用场景?访问 CSDN星图镜像广场,提供丰富的预置镜像,覆盖大模型推理、图像生成、视频生成、模型微调等多个领域,支持一键部署。

Logo

Agent 垂直技术社区,欢迎活跃、内容共建。

更多推荐