class_analyzer.py

import torch
from qwen_vl_utils import process_vision_info
from typing import List, Dict, Any
import time
import numpy as np
from PIL import Image
from decord import VideoReader, cpu
from transformers import AutoProcessor
from modelling_qwen2_5_vl import Qwen2_5_VLForConditionalGeneration


class MLLMAnalyzer:
    def __init__(self, model_path: str):
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
            model_path,
            torch_dtype=torch.bfloat16,
            attn_implementation="flash_attention_2",
            device_map="auto",
        )
        self.processor = AutoProcessor.from_pretrained(model_path)
        self.model.eval()

    def load_video(self, video_path: str, events: List[Dict[str, Any]] = None) -> tuple:
        vr = VideoReader(video_path, ctx=cpu(0))
        total_frames = len(vr)

        if events is not None:
            print("Using events to load video")
            indices = []
            avg_fps = vr.get_avg_fps()
            high_fps_ranges = []
            for event in events:
                if event["time"] is not None:
                    start = max(0, int((event["time"] - 1) * avg_fps))
                    end = min(total_frames - 1, int((event["time"] + 1) * avg_fps))
                    high_fps_ranges.append((start, end))
            def merge_ranges(ranges):
                if not ranges:
                    return []
                ranges.sort()
                merged = [ranges[0]]
                for current in ranges[1:]:
                    prev = merged[-1]
                    if current[0] <= prev[1]:
                        merged[-1] = (prev[0], max(prev[1], current[1]))
                    else:
                        merged.append(current)
                return merged
            high_fps_ranges = merge_ranges(high_fps_ranges)

            indices = set()
            for start, end in high_fps_ranges:
                step = max(1, int(avg_fps // 2))
                indices.update(range(start, end + 1, step))
            indices = sorted(list(indices))
            frames = vr.get_batch(indices).asnumpy()
            timestamps = np.array([vr.get_frame_timestamp(idx) for idx in indices])
        else:
            print("Not using events to load video")
            indices = np.linspace(
                0, total_frames - 1, num=int(total_frames / vr.get_avg_fps()), dtype=int
            )
            frames = vr.get_batch(indices).asnumpy()
            timestamps = np.array([vr.get_frame_timestamp(idx) for idx in indices])

        # pixel level drop
        final_frames = []
        final_timestamps = []
        final_frames.append(frames[0])
        final_timestamps.append(timestamps[0])
        for idx in range(len(frames) - 1):
            prev_frame = frames[idx]
            curr_frame = frames[idx + 1]
            diff = np.abs(prev_frame - curr_frame).mean(axis=2).mean(axis=1)
            pixels_threshold = 30  # pixel level drop threshold;
            if diff.mean() < pixels_threshold:
                continue
            final_frames.append(frames[idx + 1])
            final_timestamps.append(timestamps[idx + 1])
        return np.array(final_frames), np.array(final_timestamps)

    def create_image_grid(self, images, num_columns=8):
        pil_images = [Image.fromarray(image) for image in images]
        import math
        num_rows = math.ceil(len(images) / num_columns)
        img_width, img_height = pil_images[0].size
        grid_width = num_columns * img_width
        grid_height = num_rows * img_height
        grid_image = Image.new("RGB", (grid_width, grid_height))
        for idx, image in enumerate(pil_images):
            row_idx = idx // num_columns
            col_idx = idx % num_columns
            position = (col_idx * img_width, row_idx * img_height)
            grid_image.paste(image, position)
        return grid_image

    def process_frames(self, frames: np.ndarray) -> List[Image.Image]:
        return [Image.fromarray(frame) for frame in frames]

    def generate_response(self, frames: List[Image.Image], query: str):
        messages = [
            {
                "role": "system",
                "content": "You are a helpful assistant. You will be given a video (list of frames) about GUI operations and a query. You need to answer the query based on the video. Answer in English.",
            },
            {
                "role": "user",
                "content": [
                    {"type": "text", "text": query},
                    {"type": "video", "video": frames},
                ],
            },
        ]

        text = self.processor.apply_chat_template(
            messages, tokenize=False, add_generation_prompt=True
        )
        image_inputs, video_inputs = process_vision_info([messages])
        inputs = self.processor(
            text=[text],
            images=image_inputs,
            videos=video_inputs,
            padding=True,
            return_tensors="pt",
        )
        inputs = inputs.to(self.device)
        output_ids = self.model.generate(
            **inputs,
            max_new_tokens=2048,
            drop_method="feature",
            drop_threshold=0.5, # feature level drop threshold
            drop_absolute=True,
            dr_save_path="res/drop_ratio.json",
            dp_save_path="res/drop_positions.json",
        )
        generated_ids = [
            output_ids[len(input_ids) :]
            for input_ids, output_ids in zip(inputs.input_ids, output_ids)
        ]
        output_text = self.processor.batch_decode(
            generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
        )
        return output_text[0]

    def test_video_understanding(
        self, video_path: str, query: str, events: List[Dict[str, Any]] = None
    ) -> Dict[str, Any]:
        start_time = time.time()
        frames, timestamps = self.load_video(video_path, events)
        pil_frames = self.process_frames(frames)
        response = self.generate_response(pil_frames, query)
        end_time = time.time()
        return {
            "response": response,
            "processing_time": end_time - start_time,
            "num_frames": len(frames),
        }