Direct_Transformers/infer.py

import os
import uuid
import time
import psutil
import torch
import cv2
import shutil
from models.qwen import Qwen2VL
from models.gemma import Gemma
from models.minicpm import MiniCPM
from models.lfm import LFM2
from video_processor import extract_frames, FrameSamplingMethod
import argparse
import json
import logging
from tqdm import tqdm
TEMP_VIDEO_DIR = "temp_videos"
def process_video(model, video_path, prompt, sampling_method_str="CONTENT_AWARE", sampling_rate=5):
    """
    直接处理视频和文本提示，进行推理并返回结果。
    
    Args:
        video_path (str): 视频文件路径
        prompt (str): 文本提示
        sampling_method_str (str): 采样方法字符串
        sampling_rate (int): 采样率或阈值
    
    Returns:
        dict: 推理结果
    """
    request_start_time = time.time()
    request_id = str(uuid.uuid4())
    logging.info(f"[{request_id}] Processing video: '{video_path}', Prompt: '{prompt}'")

    # 验证视频文件
    if not os.path.exists(video_path):
        raise FileNotFoundError(f"Video file not found: {video_path}")
    
    if not video_path.lower().endswith(('.mp4', '.avi', '.mov', '.mkv')):
        logging.warning(f"[{request_id}] File '{video_path}' may not be a video file.")

    # 转换采样方法字符串为枚举
    sampling_method_map = {
        "CONTENT_AWARE": FrameSamplingMethod.CONTENT_AWARE,
        "UNIFORM": FrameSamplingMethod.UNIFORM,
    }
    sampling_method = sampling_method_map.get(sampling_method_str, FrameSamplingMethod.CONTENT_AWARE)
    
    # 创建临时目录
    temp_frame_dir = os.path.join(TEMP_VIDEO_DIR, request_id)
    os.makedirs(temp_frame_dir, exist_ok=True)

    try:
        logging.info(f"[{request_id}] Extracting frames using method: {sampling_method.value}, rate/threshold: {sampling_rate}")
        
        frames = extract_frames(video_path, sampling_method, sampling_rate)
        if not frames:
            raise ValueError(f"Could not extract any frames from the video: {video_path}")
        
        logging.info(f"[{request_id}] Extracted {len(frames)} frames successfully. Saving to temporary files...")

        # 将帧保存到临时文件并获取其路径
        frame_paths = []
        for i, frame in enumerate(frames):
            frame_path = os.path.join(temp_frame_dir, f"frame_{i:04d}.jpg")
            cv2.imwrite(frame_path, frame)
            abs_frame_path = os.path.abspath(frame_path)
            frame_paths.append(abs_frame_path)

        logging.info(f"[{request_id}] {len(frame_paths)} frames saved to {temp_frame_dir}")

        # 进行推理
        output = model.generate(frame_paths, prompt)

        logging.info(f"Tokens per second: {output['tokens_per_second']}, Peak GPU memory MB: {output['peak_gpu_memory_mb']}")

        inference_end_time = time.time()
        cpu_usage = psutil.cpu_percent(interval=None)
        cpu_core_utilization = psutil.cpu_percent(interval=None, percpu=True)
        logging.info(f"[{request_id}] Inference time: {inference_end_time - request_start_time:.2f} seconds, CPU usage: {cpu_usage}%, CPU core utilization: {cpu_core_utilization}")
        
        # 添加性能指标到输出
        output["inference_time"] = inference_end_time - request_start_time
        output["cpu_usage"] = cpu_usage
        output["cpu_core_utilization"] = cpu_core_utilization
        output["num_generated_tokens"] = output["num_generated_tokens"]
        output["request_id"] = request_id
        
        return output

    except Exception as e:
        logging.error(f"[{request_id}] An error occurred during processing: {str(e)}", exc_info=True)
        raise e
    finally:
        # 清理临时文件
        if os.path.exists(temp_frame_dir):
            shutil.rmtree(temp_frame_dir)
            logging.info(f"[{request_id}] Cleaned up temporary frame directory: {temp_frame_dir}")


def main():
    """主函数"""
    try:
        parser = argparse.ArgumentParser()
        parser.add_argument("--model_path", type=str, default="Qwen/Qwen2.5-VL-3B-Instruct-AWQ")
        parser.add_argument("--video_dir", type=str, default="videos", help="视频")
        parser.add_argument("--prompt", type=str, default="Summarize the key observable events in this 1-minute convenience store video clip. Focus strictly on the physical actions and interactions of the people. Describe only what you can see; do not interpret intentions, relationships, or work efficiency. Avoid all repetitive descriptions of the store's layout or shelves.", help="文本提示")
        parser.add_argument("--sampling_method", type=str, default="UNIFORM", 
                        choices=["CONTENT_AWARE", "UNIFORM", "RANDOM"], 
                        help="帧采样方法")
        parser.add_argument("--sampling_rate", type=int, default=30, help="采样率或阈值")
        args = parser.parse_args()


        # --- 日志和临时文件目录配置 ---
        LOG_DIR = f"logs/{args.model_path.split('/')[-1]}"
        OUTPUT_DIR = f"outputs/{args.model_path.split('/')[-1]}"
        os.makedirs(LOG_DIR, exist_ok=True)
        os.makedirs(OUTPUT_DIR, exist_ok=True)
        os.makedirs(TEMP_VIDEO_DIR, exist_ok=True)
        start_time = time.strftime('%Y%m%d_%H%M%S')
        log_filename = f"{LOG_DIR}/{start_time}.log"
        logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s', datefmt='%Y-%m-%d %H:%M:%S', filename=log_filename, filemode='a')

        # --- 加载模型和处理器 ---
        logging.info(f"Loading model: {args.model_path}")
        model_load_start = time.time()
        if "qwen" in args.model_path.lower():
            model = Qwen2VL(args.model_path)
        elif "gemma" in args.model_path.lower():
            model = Gemma(args.model_path)
        elif "minicpm" in args.model_path.lower():
            model = MiniCPM(args.model_path)
        elif "lfm" in args.model_path.lower():
            model = LFM2(args.model_path)
        model_load_end = time.time()
        GPU_MEMORY_USAGE = f"{torch.cuda.memory_allocated(0)/1024**2:.2f} MB" if torch.cuda.is_available() else "N/A"
        logging.info(f"Model loaded in {model_load_end - model_load_start:.2f} seconds")
        logging.info(f"GPU Memory Usage after model load: {GPU_MEMORY_USAGE}")
        # 处理视频
        total_output = {}
        for video_path in tqdm(os.listdir(args.video_dir)):
            result = process_video(
                model=model,
                video_path=os.path.join(args.video_dir, video_path),
                prompt=args.prompt,
                sampling_method_str=args.sampling_method,
                sampling_rate=args.sampling_rate
            )
            total_output[video_path] = result
            # 保存结果到文件
            output_filename = f"{OUTPUT_DIR}/{start_time}.json"
            with open(output_filename, 'w', encoding='utf-8') as f:
                json.dump(total_output, f, ensure_ascii=False, indent=2)
            
            print(f"处理完成！结果已保存到: {output_filename}")
            print(f"推理时间: {result['inference_time']:.2f} 秒")
            print(f"生成的内容: {result.get('generated_text', 'N/A')}")
        
    except Exception as e:
        logging.error(f"处理失败: {str(e)}", exc_info=True)
        print(f"处理失败: {str(e)}")


if __name__ == "__main__":
    main()