image_data_processing.py

import re
import os
import time
from nltk.tokenize import word_tokenize
from nltk.corpus import stopwords
from nltk.stem import WordNetLemmatizer
from rich.progress import Progress, TextColumn, BarColumn, TimeElapsedColumn
from data_processing_common import sanitize_filename  # Import sanitize_filename


def get_text_from_generator(generator):
    """从生成器响应中提取文本内容。

    Args:
        generator: 生成器对象，产生包含文本内容的响应

    Returns:
        str: 从生成器中提取的完整文本内容
    """
    response_text = ""
    try:
        while True:
            # 获取下一个响应
            response = next(generator)
            # 获取响应中的choices字段
            choices = response.get('choices', [])
            for choice in choices:
                # 获取delta字段（增量内容）
                delta = choice.get('delta', {})
                # 如果delta中包含content，则追加到响应文本中
                if 'content' in delta:
                    response_text += delta['content']
    except StopIteration:
        # 生成器结束时捕获StopIteration异常
        pass
    return response_text


def process_single_image(image_path, image_inference, text_inference, silent=False, log_file=None):
    """处理单个图像文件以生成元数据。

    Args:
        image_path: 图像文件路径
        image_inference: 图像推理模型
        text_inference: 文本推理模型
        silent: 是否静默模式（不在控制台输出）
        log_file: 日志文件路径

    Returns:
        dict: 包含文件路径、文件夹名、文件名和描述的字典
    """
    # 记录开始时间
    start_time = time.time()

    # 为当前文件创建进度条实例
    with Progress(
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TimeElapsedColumn()
    ) as progress:
        # 添加处理任务到进度条
        task_id = progress.add_task(f"正在处理 {os.path.basename(image_path)}", total=1.0)
        # 生成图像元数据（文件夹名、文件名、描述）
        foldername, filename, description = generate_image_metadata(
            image_path, progress, task_id, image_inference, text_inference
        )

    # 记录结束时间
    end_time = time.time()
    # 计算处理耗时
    time_taken = end_time - start_time

    # 构建输出消息
    message = (
        f"文件: {image_path}\n"
        f"处理耗时: {time_taken:.2f} 秒\n"
        f"图片描述: {description}\n"
        f"文件夹名称: {foldername}\n"
        f"生成的文件名: {filename}\n"
    )

    # 根据静默模式决定输出方式
    if silent:
        # 静默模式：如果指定了日志文件，则写入日志
        if log_file:
            with open(log_file, 'a', encoding='utf-8') as f:
                f.write(message + '\n')
    else:
        # 非静默模式：直接打印到控制台
        print(message)

    # 返回处理结果字典
    return {
        'file_path': image_path,  # 原始文件路径
        'foldername': foldername,  # 生成的文件夹名称
        'filename': filename,  # 生成的文件名
        'description': description  # 图片描述
    }


def process_image_files(image_paths, image_inference, text_inference, silent=False, log_file=None):
    """Process image files sequentially."""
    data_list = []
    for image_path in image_paths:
        data = process_single_image(image_path, image_inference, text_inference, silent=silent, log_file=log_file)
        data_list.append(data)
    return data_list


"""
图像数据处理模块 - 使用 LM Studio API
处理图像文件的分析、元数据生成和组织
"""

from pathlib import Path
from typing import List, Dict, Any, Tuple, Optional
from rich.progress import Progress, SpinnerColumn, TextColumn, BarColumn, TaskProgressColumn
import time
import re
import os
from nltk.corpus import stopwords
from nltk.tokenize import word_tokenize
from nltk.stem import WordNetLemmatizer

# 初始化 NLTK 资源
try:
    stopwords.words('english')
except LookupError:
    import nltk

    nltk.download('stopwords', quiet=True)
    nltk.download('punkt', quiet=True)
    nltk.download('punkt_tab', quiet=True)
    nltk.download('wordnet', quiet=True)


def process_image_files(
        image_files: List[Path],
        image_inference,
        text_inference,
        silent: bool = False,
        log_file=None
) -> List[Dict[str, Any]]:
    """
    批量处理图像文件

    Args:
        image_files: 图像文件路径列表
        image_inference: VLM 模型实例
        text_inference: 文本模型实例
        silent: 是否静默模式
        log_file: 日志文件句柄

    Returns:
        图像数据列表
    """
    data_images = []

    with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TaskProgressColumn(),
    ) as progress:

        for image_path in image_files:
            # ✅ 确保转换为 Path 对象
            if isinstance(image_path, str):
                image_path = Path(image_path)

            task_id = progress.add_task(f"Processing {image_path.name}", total=100)

            try:
                data = process_single_image(
                    image_path,
                    image_inference,
                    text_inference,
                    silent=silent,
                    log_file=log_file
                )
                data_images.append(data)

                progress.update(task_id, completed=100)

            except Exception as e:
                error_msg = f"Failed to process {image_path.name}: {e}"
                print(f"[ERROR] {error_msg}")
                if log_file:
                    log_file.write(f"{error_msg}\n")

                progress.update(task_id, completed=100, description=f"❌ {image_path.name}")

    return data_images

def process_single_image(
        image_path: Path,
        image_inference,
        text_inference,
        silent: bool = False,
        log_file=None
) -> Dict[str, Any]:
    """
    处理单个图像文件

    Args:
        image_path: 图像文件路径
        image_inference: VLM 模型实例
        text_inference: 文本模型实例
        silent: 是否静默模式
        log_file: 日志文件句柄

    Returns:
        图像数据字典
    """
    # ✅ 确保转换为 Path 对象
    if isinstance(image_path, str):
        image_path = Path(image_path)

    with Progress(
            SpinnerColumn(),
            TextColumn("[progress.description]{task.description}"),
            BarColumn(),
            TaskProgressColumn(),
    ) as progress:
        task_id = progress.add_task(f"Analyzing {image_path.name}", total=100)

        # 生成元数据
        foldername, filename, description = generate_image_metadata(
            image_path,
            progress,
            task_id,
            image_inference,
            text_inference
        )

        # 构建数据字典
        data = {
            "path": str(image_path),
            "foldername": foldername,
            "filename": filename,
            "description": description
        }

        progress.update(task_id, completed=100)

        return data


def generate_image_metadata(image_path, progress, task_id, image_inference, text_inference):
    """
    为图像文件生成描述、文件夹名称和文件名

    参数:
        image_path: 图像文件路径
        progress: 进度条对象
        task_id: 任务ID
        image_inference: 图像推理模型
        text_inference: 文本推理模型

    返回:
        (文件夹名, 文件名, 描述) 的元组
    """

    # 处理图像的总步骤数
    total_steps = 3

    # ============ 步骤1: 使用图像推理模型生成描述 ============
    description_prompt = "Please provide a detailed description of this image, focusing on the main subject and any important details."

    try:
        # 使用 create_completion 替代 _chat
        result = image_inference.create_completion(
            prompt=description_prompt,
            image_path=str(image_path),
            max_tokens=200,
            temperature=0.3
        )
        description = result['choices'][0]['text'].strip()

        if not description:
            description = "图像分析不可用"
            print(f"[警告] 图像描述为空，使用默认值")

    except Exception as e:
        print(f"[ERROR] 图像描述生成失败: {e}")
        description = "图像分析失败"

    # 更新进度条
    progress.update(task_id, advance=100 / total_steps)

    # ============ 步骤2: 使用文本推理模型生成文件名 ============
    filename_prompt = f"""Based on the description below, generate a specific and descriptive filename for the image.
Limit the filename to a maximum of 3 words. Use nouns and avoid starting with verbs like 'depicts', 'shows', 'presents', etc.
Do not include any data type words like 'image', 'jpg', 'png', etc. Use only letters and connect words with underscores.

Description: {description}

Example:
Description: A photo of a sunset over the mountains.
Filename: sunset_over_mountains

Now generate the filename.

Output only the filename, without any additional text.

Filename:"""

    try:
        filename_response = text_inference.create_completion(
            prompt=filename_prompt,
            max_tokens=30,
            temperature=0.5
        )
        filename = filename_response['choices'][0]['text'].strip()
        # 移除可能存在的 'Filename:' 前缀
        filename = re.sub(r'^Filename:\s*', '', filename, flags=re.IGNORECASE).strip()
        print(f"[INFO] 文件名生成成功: {filename}")

    except Exception as e:
        print(f"[ERROR] 文件名生成失败: {e}")
        filename = ""

    # 更新进度条
    progress.update(task_id, advance=100 / total_steps)

    # ============ 步骤3: 从描述中生成文件夹名称 ============
    foldername_prompt = f"""Based on the description below, generate a general category or theme that best represents the main subject of this image.
This will be used as the folder name. Limit the category to a maximum of 2 words. Use nouns and avoid verbs.
Do not include specific details, words from the filename, or any generic terms like 'untitled' or 'unknown'.

Description: {description}

Examples:
1. Description: A photo of a sunset over the mountains.
   Category: landscapes

2. Description: An image of a smartphone displaying a storage app with various icons and information.
   Category: technology

3. Description: A close-up of a blooming red rose with dew drops.
   Category: nature

Now generate the category.

Output only the category, without any additional text.

Category:"""

    try:
        foldername_response = text_inference.create_completion(
            prompt=foldername_prompt,
            max_tokens=20,
            temperature=0.5
        )
        foldername = foldername_response['choices'][0]['text'].strip()
        # 移除可能存在的 'Category:' 前缀
        foldername = re.sub(r'^Category:\s*', '', foldername, flags=re.IGNORECASE).strip()
        print(f"[INFO] 文件夹名称生成成功: {foldername}")

    except Exception as e:
        print(f"[ERROR] 文件夹名称生成失败: {e}")
        foldername = ""

    # 更新进度条
    progress.update(task_id, advance=100 / total_steps)

    # ============ 清理和处理AI输出 ============
    # 定义不需要的词汇和停用词
    unwanted_words = set([
        'the', 'and', 'based', 'generated', 'this', 'is', 'filename', 'file', 'image', 'picture', 'photo',
        'folder', 'category', 'output', 'only', 'below', 'text', 'jpg', 'png', 'jpeg', 'gif', 'bmp', 'svg',
        'logo', 'in', 'on', 'of', 'with', 'by', 'for', 'to', 'from', 'a', 'an', 'as', 'at', 'red', 'blue',
        'green', 'color', 'colors', 'colored', 'text', 'graphic', 'graphics', 'main', 'subject', 'important',
        'details', 'description', 'depicts', 'show', 'shows', 'display', 'illustrates', 'presents', 'features',
        'provides', 'covers', 'includes', 'demonstrates', 'describes'
    ])
    stop_words = set(stopwords.words('english'))
    all_unwanted_words = unwanted_words.union(stop_words)
    lemmatizer = WordNetLemmatizer()

    # 清理和处理AI输出的函数
    def clean_ai_output(text, max_words):
        """
        清理AI生成的文本输出

        参数:
            text: 需要清理的文本
            max_words: 最大词数限制

        返回:
            清理后的文本，使用下划线连接
        """
        # 移除文件扩展名和特殊字符
        text = re.sub(r'\.\w{1,4}$', '', text)  # 移除如 .jpg, .png 等文件扩展名
        text = re.sub(r'[^\w\s]', ' ', text)  # 移除特殊字符
        text = re.sub(r'\d+', '', text)  # 移除数字
        text = text.strip()

        # 分割连接的单词 (例如: 'GoogleChrome' -> 'Google Chrome')
        text = re.sub(r'([a-z])([A-Z])', r'\1 \2', text)

        # 分词并词形还原
        words = word_tokenize(text)
        words = [word.lower() for word in words if word.isalpha()]
        words = [lemmatizer.lemmatize(word) for word in words]

        # 移除不需要的词汇和重复项
        filtered_words = []
        seen = set()
        for word in words:
            if word not in all_unwanted_words and word not in seen:
                filtered_words.append(word)
                seen.add(word)

        # 限制最大词数
        filtered_words = filtered_words[:max_words]
        return '_'.join(filtered_words)

    # 处理文件名
    print(f"[INFO] 开始清理文件名...")
    filename = clean_ai_output(filename, max_words=3)
    if not filename or filename.lower() in ('untitled', ''):
        # 如果文件名为空或无效，从描述中提取关键词
        print(f"[警告] 文件名无效，尝试从描述中提取关键词")
        filename = clean_ai_output(description, max_words=3)
    if not filename:
        # 如果仍然为空，使用原始文件名
        filename = 'image_' + os.path.splitext(os.path.basename(str(image_path)))[0]
        print(f"[警告] 无法生成有效文件名，使用默认值: {filename}")

    # 清理文件名，确保符合文件系统规范
    sanitized_filename = sanitize_filename(filename, max_words=3)
    print(f"[INFO] 最终文件名: {sanitized_filename}")

    # 处理文件夹名
    print(f"[INFO] 开始清理文件夹名称...")
    foldername = clean_ai_output(foldername, max_words=2)
    if not foldername or foldername.lower() in ('untitled', ''):
        # 如果文件夹名为空或无效，尝试从描述中提取关键词
        print(f"[警告] 文件夹名称无效，尝试从描述中提取关键词")
        foldername = clean_ai_output(description, max_words=2)
        if not foldername:
            # 如果仍然为空，使用默认值
            foldername = 'images'
            print(f"[警告] 无法生成有效文件夹名称，使用默认值: {foldername}")

    # 清理文件夹名，确保符合文件系统规范
    sanitized_foldername = sanitize_filename(foldername, max_words=2)
    print(f"[INFO] 最终文件夹名称: {sanitized_foldername}")

    print(f"[成功] 图像元数据生成完成 - 文件夹: {sanitized_foldername}, 文件名: {sanitized_filename}")

    return sanitized_foldername, sanitized_filename, description


def sanitize_filename(filename, max_words=None):
    """
    清理文件名，移除非法字符

    Args:
        filename: 原始文件名
        max_words: 最大单词数

    Returns:
        清理后的文件名
    """
    # Remove any characters that are not alphanumeric or underscores
    filename = re.sub(r'[^\w\s-]', '', filename)
    # Replace spaces with underscores
    filename = re.sub(r'\s+', '_', filename)
    # Remove consecutive underscores
    filename = re.sub(r'_+', '_', filename)
    # Convert to lowercase
    filename = filename.lower()
    # Limit to max_words if specified
    if max_words:
        words = filename.split('_')
        filename = '_'.join(words[:max_words])
    # Trim underscores from the beginning and end
    filename = filename.strip('_')
    return filename


def get_text_from_generator(generator):
    """
    从生成器中提取文本（兼容旧代码）
    现在直接返回字符串即可

    Args:
        generator: 文本生成器或字符串

    Returns:
        提取的文本
    """
    if isinstance(generator, str):
        return generator

    # 如果是生成器，收集所有文本
    text_parts = []
    for chunk in generator:
        if isinstance(chunk, dict) and 'choices' in chunk:
            text_parts.append(chunk['choices'][0].get('text', ''))
        elif isinstance(chunk, str):
            text_parts.append(chunk)

    return ''.join(text_parts)