基于图像理解模型的自动图片筛选工具

基于图像理解模型的自动图片筛选工具

功能概述

本工具通过智谱AI接口实现自动化图片筛选，主要功能包括：

• 递归遍历指定目录下的图片文件（支持png/jpg/webp等格式）
• 调用AI模型进行图片美感评分（0.0-9.0分）
• 按评分排序并自动筛选自定义百分比的高质量图片
• 可视化进度条和日志输出

技术原理

1. 核心处理流程

# 主要处理步骤：
1. 遍历目录收集图片 -> 2. 调用AI评分 -> 3. 过滤有效分数 -> 4. 排序筛选 -> 5. 移动优质图片

2. 关键技术点

• 智能评分：通过ZhipuAITool调用视觉大模型进行图片质量评估
• 安全转换：safe_float_convert处理模型返回的非标准化评分
• 进度可视化：使用tqdm实现带彩色进度条的处理界面
• 异常处理：捕获处理过程中的各种异常并高亮显示

使用方法

# 基本调用示例
select_top_images(
    image_path="/path/to/images",
    output_dir="/path/to/selected",
    top_percentage=0.2   # 选取前20%的图片
    prompt="请根据图片的美感程度进行评分，评分范围为0.0-9.0分，评分越高表示图片美感越好。"
)

代码解析

核心函数 get_image_score

def get_image_score(image_path: str, res: List[Dict] = None) -> List[Dict]:
    # ...初始化代码...
    
    # 递归收集所有图片文件
    for root, _, files in os.walk(image_path):
        # 过滤图片文件...
    
    # 带进度条的AI评分
    for file_path in tqdm(image_files):
        try:
            # 调用AI接口获取评分
            score_str = image_tool.chat_with_image(...)
            # 转换并存储评分
            res.append({"path": file_path, "score": score})
        except Exception as e:
            # 异常处理...
    
    return res

筛选逻辑 select_top_images

def select_top_images(...):
    # ...初始化代码...
    
    # 按分数降序排序
    sorted_items = sorted(valid_items, key=lambda x: x['score'], reverse=True)
    
    # 计算选取数量（至少保留1张）
    num_select = max(1, int(len(sorted_items) * top_percentage))
    
    # 移动优质图片
    for item in sorted_items[:num_select]:
        os.rename(...)

工具优势

1. 自动化处理：一键完成从评分到筛选的全流程
2. 智能评估：基于AI的图片质量判断，比传统规则更灵活
3. 灵活配置：支持调整分数阈值和选取比例
4. 可视化反馈：彩色进度条和实时日志显示处理状态

应用场景

• 炼丹素材预处理
• 图片质量分级管理
• 自动化素材库维护
• 训练数据清洗

注意事项

1. 需要替换有效的智谱API Key
2. 确保输入/输出目录有读写权限
3. 评分标准可根据具体需求调整prompt
4. 建议先在小规模数据集上测试

提示：可通过调整threshold和top_percentage参数平衡筛选严格度，建议从宽松设置开始逐步收紧标准。

具体代码:

import os
from typing import List, Dict
from zhipu_imgToPng import ZhipuAITool
from tqdm import tqdm
from termcolor import colored


def safe_float_convert(score_str: str) -> float:
    try:
        return float(score_str.replace('分', '').strip())
    except (ValueError, TypeError):
        print(colored(f"Warning: Could not convert score '{score_str}' to float", 'yellow'))
        return 0.0

def get_image_score(image_path: str,prompt:str, res: List[Dict] = None) -> List[Dict]:
    if res is None:
        res = []
    
    image_tool = ZhipuAITool()
    
    
    image_files = []
    for root, _, files in os.walk(image_path):
        for file in files:
            file_path = os.path.join(root, file)
            if file_path.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp', '.webp')):
                image_files.append(file_path)
    
    for file_path in tqdm(image_files, desc=colored("Processing Images", "cyan"), unit="image"):
        try:
            color = 'white'
            tqdm.write(colored(f"Processing: {os.path.basename(file_path)}", color))
            
            score_str = image_tool.chat_with_image(file_path, prompt)
            
            score = safe_float_convert(score_str)
            
            score_color = 'green' if score > 5 else 'blue' if score > 3 else 'white'
            tqdm.write(colored(f"Score: {score}", score_color))
            
            res.append({"path": file_path, "score": score})
        except Exception as e:
            tqdm.write(colored(f"Error processing {file_path}: {e}", 'red', attrs=['bold']))
    
    return res

def select_top_images(image_path: str, output_dir: str,prompt:str , top_percentage: float = 0.2):
    os.makedirs(output_dir, exist_ok=True)
    
    res = get_image_score(image_path,prompt)
    
    valid_items = [item for item in res if item['score'] is not None]
    
    if not valid_items:
        print(colored("No valid image scores found.", "red"))
        return
    
    sorted_items = sorted(valid_items, key=lambda x: x['score'], reverse=True)
    
    num_select = max(1, int(len(sorted_items) * top_percentage))
    
    for item in sorted_items[:num_select]:
        file_name = os.path.basename(item['path'])
        output_path = os.path.join(output_dir, file_name)
        os.rename(item['path'], output_path)
        print(colored(f"Moved {file_name} with score {item['score']}", "green"))

总结:

更多是提供一个思路，可以更改不同要求让模型去评分/筛选，比如可以筛选出符合特定风格的图片，或者筛选出符合特定主题的图片。