upload.py

"""
批量处理贴纸图片的工具模块

功能：
1. 批量处理指定目录下的图片文件
2. 使用MD5哈希进行图片去重
3. 生成图片嵌入向量并存储到数据库
4. 可选地将处理后的图片移动到输出目录
"""
import os
from typing import List
from pathlib import Path

from app.services import sticker_service


def batch_process_stickers(image_dir: str = 'images.tmp', output_dir: str = 'processed.tmp') -> None:
    """
    批量处理目录中的贴纸图片
    
    参数:
        image_dir: 包含贴纸图片的目录路径
        output_dir: 存储处理后图片的输出目录(默认为'processed')
    
    功能:
        1. 遍历目录中的所有图片文件
        2. 使用MD5哈希进行去重
        3. 为每张图片生成嵌入向量并存储到数据库
        4. 可选地将处理后的图片移动到输出目录
    """
    # 创建输出目录(如果不存在)
    os.makedirs(output_dir, exist_ok=True)
    # 用于存储已处理图片的哈希集合
    # 遍历目录中的所有文件
    # 打印图片数量和列表
    print('>>> 待处理图片数量', len(os.listdir(image_dir)))
    for img_path in Path(image_dir).glob('*.*'):
        # 只处理jpg/jpeg/png格式的图片
        print('>>>> 正在处理图片', img_path)
        if img_path.suffix.lower() not in ['.jpg', '.jpeg', '.png']:
            continue
        
        try:
            # 加载图片并准备元数据
            from PIL import Image
            img = Image.open(img_path)
            title = img_path.stem  # 使用文件名作为标题
            
            # 调用上传服务
            sticker_service.upload_sticker(img, title, "", "")
            
            # 可选操作: 将处理后的图片移动到输出目录
            output_path = os.path.join(output_dir, img_path.name)
            os.rename(str(img_path), output_path)
            
            print(f"已处理: {img_path}")
        except Exception as e:
            print(f"处理异常 {img_path}: {str(e)}")


if __name__ == '__main__':
    """
    命令行入口
    用法: python batch_upload.py <图片目录> [--output 输出目录]
    """
    import argparse
    
    # 设置命令行参数解析
    parser = argparse.ArgumentParser(description='批量处理贴纸图片(带去重功能)')
    parser.add_argument('image_dir', help='包含贴纸图片的目录')
    parser.add_argument('--output', '-o', default='processed.tmp', 
                       help='处理后图片的输出目录(默认为processed)')
    
    # 解析参数并执行批量处理
    args = parser.parse_args()
    batch_process_stickers(args.image_dir, args.output)