app/utils/qwenvl_analyzer.py

import json
from typing import List, Union, Dict
import os
from pathlib import Path
from loguru import logger
from tqdm import tqdm
import asyncio
from tenacity import retry, stop_after_attempt, RetryError, wait_exponential
from openai import OpenAI
import PIL.Image
import base64
import io
import traceback


class QwenAnalyzer:
    """千问视觉分析器类"""

    def __init__(self, model_name: str = "qwen-vl-max-latest", api_key: str = None, base_url: str = None):
        """
        初始化千问视觉分析器
        
        Args:
            model_name: 模型名称，默认使用 qwen-vl-max-latest
            api_key: 阿里云API密钥
            base_url: API基础URL，如果为None则使用默认值
        """
        if not api_key:
            raise ValueError("必须提供API密钥")

        self.model_name = model_name
        self.api_key = api_key
        self.base_url = base_url

        # 配置API客户端
        self._configure_client()

    def _configure_client(self):
        """
        配置API客户端
        使用最简化的参数配置，避免不必要的参数
        """
        try:
            self.client = OpenAI(
                api_key=self.api_key,
                base_url=self.base_url
            )
        except Exception as e:
            logger.error(f"初始化OpenAI客户端失败: {str(e)}")
            raise

    def _image_to_base64(self, image: PIL.Image.Image) -> str:
        """
        将PIL图片对象转换为base64字符串
        """
        buffered = io.BytesIO()
        image.save(buffered, format="JPEG")
        return base64.b64encode(buffered.getvalue()).decode("utf-8")

    @retry(
        stop=stop_after_attempt(3),
        wait=wait_exponential(multiplier=1, min=4, max=10)
    )
    async def _generate_content_with_retry(self, prompt: str, batch: List[PIL.Image.Image]):
        """使用重试机制的内部方法来调用千问API"""
        try:
            # 构建消息内容
            content = []

            # 添加图片
            for img in batch:
                base64_image = self._image_to_base64(img)
                content.append({
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/jpeg;base64,{base64_image}"
                    }
                })

            # 添加文本提示
            content.append({
                "type": "text",
                "text": prompt % (len(content), len(content), len(content))
            })

            # 调用API
            response = await asyncio.to_thread(
                self.client.chat.completions.create,
                model=self.model_name,
                messages=[{
                    "role": "user",
                    "content": content
                }]
            )

            return response.choices[0].message.content

        except Exception as e:
            logger.error(f"API调用错误: {str(e)}")
            raise RetryError("API调用失败")

    async def analyze_images(self,
                             images: Union[List[str], List[PIL.Image.Image]],
                             prompt: str,
                             batch_size: int) -> List[Dict]:
        """
        批量分析多张图片
        Args:
            images: 图片路径列表或PIL图片对象列表
            prompt: 分析提示词
            batch_size: 批处理大小
        Returns:
            分析结果列表
        """
        try:
            # 保存原始图片路径（如果是路径列表的话）
            original_paths = images if isinstance(images[0], str) else None

            # 加载图片
            if isinstance(images[0], str):
                images = self.load_images(images)

            # 验证图片列表
            if not images:
                raise ValueError("图片列表为空")

            # 验证每个图片对象
            valid_images = []
            valid_paths = []
            for i, img in enumerate(images):
                if not isinstance(img, PIL.Image.Image):
                    logger.error(f"无效的图片对象，索引 {i}: {type(img)}")
                    continue
                valid_images.append(img)
                if original_paths:
                    valid_paths.append(original_paths[i])

            if not valid_images:
                raise ValueError("没有有效的图片对象")

            images = valid_images
            results = []
            # 视频帧总数除以批量处理大小，如果有小数则+1
            batches_needed = len(images) // batch_size
            if len(images) % batch_size > 0:
                batches_needed += 1
                
            logger.debug(f"视频帧总数:{len(images)}, 每批处理 {batch_size} 帧, 需要访问 VLM {batches_needed} 次")

            with tqdm(total=batches_needed, desc="分析进度") as pbar:
                for i in range(0, len(images), batch_size):
                    batch = images[i:i + batch_size]
                    batch_paths = valid_paths[i:i + batch_size] if valid_paths else None
                    retry_count = 0

                    while retry_count < 3:
                        try:
                            # 在每个批次处理前添加小延迟
                            # if i > 0:
                            #     await asyncio.sleep(0.5)

                            # 确保每个批次的图片都是有效的
                            valid_batch = [img for img in batch if isinstance(img, PIL.Image.Image)]
                            if not valid_batch:
                                raise ValueError(f"批次 {i // batch_size} 中没有有效的图片")

                            response = await self._generate_content_with_retry(prompt, valid_batch)
                            result_dict = {
                                'batch_index': i // batch_size,
                                'images_processed': len(valid_batch),
                                'response': response,
                                'model_used': self.model_name
                            }

                            # 添加图片路径信息（如果有的话）
                            if batch_paths:
                                result_dict['image_paths'] = batch_paths

                            results.append(result_dict)
                            break

                        except Exception as e:
                            retry_count += 1
                            error_msg = f"批次 {i // batch_size} 处理出错: {str(e)}"
                            logger.error(error_msg)

                            if retry_count >= 3:
                                results.append({
                                    'batch_index': i // batch_size,
                                    'images_processed': len(batch),
                                    'error': error_msg,
                                    'model_used': self.model_name,
                                    'image_paths': batch_paths if batch_paths else []
                                })
                            else:
                                logger.info(f"批次 {i // batch_size} 处理失败，等待60秒后重试当前批次...")
                                await asyncio.sleep(60)

                    pbar.update(1)

            return results

        except Exception as e:
            error_msg = f"图片分析过程中发生错误: {str(e)}\n{traceback.format_exc()}"
            logger.error(error_msg)
            raise Exception(error_msg)

    def save_results_to_txt(self, results: List[Dict], output_dir: str):
        """将分析结果保存到txt文件"""
        # 确保输出目录存在
        os.makedirs(output_dir, exist_ok=True)

        for i, result in enumerate(results):
            response_text = result['response']

            # 如果有图片路径信息，用它来生成文件名
            if result.get('image_paths'):
                image_paths = result['image_paths']
                img_name_start = Path(image_paths[0]).stem.split('_')[-1]
                img_name_end = Path(image_paths[-1]).stem.split('_')[-1]
                file_name = f"frame_{img_name_start}_{img_name_end}.txt"
            else:
                # 如果没有路径信息，使用批次索引
                file_name = f"batch_{result['batch_index']}.txt"

            txt_path = os.path.join(output_dir, file_name)

            # 保存结果到txt文件
            with open(txt_path, 'w', encoding='utf-8') as f:
                f.write(response_text.strip())
            logger.info(f"已保存分析结果到: {txt_path}")

    def load_images(self, image_paths: List[str]) -> List[PIL.Image.Image]:
        """
        加载多张图片
        Args:
            image_paths: 图片路径列表
        Returns:
            加载后的PIL Image对象列表
        """
        images = []
        failed_images = []

        for img_path in image_paths:
            try:
                if not os.path.exists(img_path):
                    logger.error(f"图片文件不存在: {img_path}")
                    failed_images.append(img_path)
                    continue

                img = PIL.Image.open(img_path)
                # 确保图片被完全加载
                img.load()
                # 转换为RGB模式
                if img.mode != 'RGB':
                    img = img.convert('RGB')
                images.append(img)

            except Exception as e:
                logger.error(f"无法加载图片 {img_path}: {str(e)}")
                failed_images.append(img_path)

        if failed_images:
            logger.warning(f"以下图片加载失败:\n{json.dumps(failed_images, indent=2, ensure_ascii=False)}")

        if not images:
            raise ValueError("没有成功加载任何图片")

        return images