视频-语言多模态

import torch
import torch.nn as nn
from einops import rearrange

class TimeSformerBlock(nn.Module):
    """TimeSformer 时空分解注意力块"""

    def __init__(self, dim: int, num_heads: int, num_frames: int,
                 num_patches: int, mlp_ratio: float = 4.0, dropout: float = 0.1):
        super().__init__()
        self.num_frames = num_frames
        self.num_patches = num_patches

        # 时间注意力：沿时间维度计算
        self.temporal_norm = nn.LayerNorm(dim)
        self.temporal_attn = nn.MultiheadAttention(dim, num_heads, dropout=dropout, batch_first=True)

        # 空间注意力：沿空间维度计算
        self.spatial_norm = nn.LayerNorm(dim)
        self.spatial_attn = nn.MultiheadAttention(dim, num_heads, dropout=dropout, batch_first=True)

        self.mlp_norm = nn.LayerNorm(dim)
        self.mlp = nn.Sequential(
            nn.Linear(dim, int(dim * mlp_ratio)),
            nn.GELU(),
            nn.Dropout(dropout),
            nn.Linear(int(dim * mlp_ratio), dim),
        )

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # x: [B, T * P, D] -> [B, T, P, D]
        B, _, D = x.shape
        x = x.reshape(B, self.num_frames, self.num_patches, D)

        # 时间注意力：固定空间位置，跨帧交互
        x_t = rearrange(x, "b t p d -> (b p) t d")
        x_t = x_t + self.temporal_attn(self.temporal_norm(x_t), x_t, x_t)[0]

        # 空间注意力：固定时间帧，空间位置交互
        x_t = rearrange(x_t, "(b p) t d -> (b t) p d", b=B, p=self.num_patches)
        x_t = x_t + self.spatial_attn(self.spatial_norm(x_t), x_t, x_t)[0]

        # MLP
        x_t = rearrange(x_t, "(b t) p d -> b t p d", b=B, p=self.num_patches)
        x_t = x_t.reshape(B, self.num_frames * self.num_patches, D)
        x_t = x_t + self.mlp(self.mlp_norm(x_t))
        return x_t

# 3D CNN vs TimeSformer 计算量对比
import torch

def calc_3d_cnn_flops(input_size, kernel_size, channels_in, channels_out):
    """计算 3D 卷积的 FLOPs
    公式: T*H*W * Co * Ci * Kt * Kh * Kw
    """
    T, H, W = input_size
    Kt, Kh, Kw = kernel_size
    flops = T * H * W * channels_out * channels_in * Kt * Kh * Kw
    return flops

def calc_timesformer_flops(num_frames, num_patches, dim, num_heads):
    """计算 TimeSformer 单块的 FLOPs
    时间注意力: P * T^2 * D  (固定空间位置，跨 T 帧)
    空间注意力: T * P^2 * D  (固定时间帧，跨 P 空间块)
    相比全注意力 T^2 * P^2 * D，节省了大量计算
    """
    temporal_flops = num_patches * (num_frames ** 2) * dim
    spatial_flops = num_frames * (num_patches ** 2) * dim
    total = temporal_flops + spatial_flops
    full_attn_flops = (num_frames  2) * (num_patches  2) * dim
    savings = (1 - total / full_attn_flops) * 100
    return total, full_attn_flops, savings

# 对比：16帧、224x224 输入
flops_3d = calc_3d_cnn_flops((16, 224, 224), (3, 3, 3), 64, 128)
tf_total, tf_full, savings = calc_timesformer_flops(16, 196, 768, 12)
print(f"3D Conv FLOPs:  {flops_3d/1e9:.1f}G")
print(f"TimeSformer FLOPs: {tf_total/1e9:.1f}G (分解)")
print(f"Full Attention:    {tf_full/1e9:.1f}G")
print(f"节省: {savings:.1f}%")

import torch
import torch.nn as nn
import torch.nn.functional as F

class VideoCLIPTrainer(nn.Module):
    """VideoCLIP 对比学习训练框架"""

    def __init__(self, video_encoder, text_encoder,
                 embed_dim: int = 512, temperature: float = 0.07):
        super().__init__()
        self.video_encoder = video_encoder
        self.text_encoder = text_encoder
        self.video_proj = nn.Linear(video_encoder.output_dim, embed_dim)
        self.text_proj = nn.Linear(text_encoder.output_dim, embed_dim)
        self.logit_scale = nn.Parameter(torch.ones([]) * np.log(1 / temperature))

    def forward(self, videos: torch.Tensor, texts: torch.Tensor,
                text_mask: torch.Tensor) -> dict:
        # 编码
        video_feat = self.video_encoder(videos)  # [B, T, D_v]
        text_feat = self.text_encoder(texts, text_mask)  # [B, D_t]

        # 池化视频特征：对帧做平均池化
        video_feat = video_feat.mean(dim=1)  # [B, D_v]

        # 投影到统一嵌入空间
        v_emb = F.normalize(self.video_proj(video_feat), dim=-1)
        t_emb = F.normalize(self.text_proj(text_feat), dim=-1)

        # 相似度矩阵
        logit_scale = self.logit_scale.exp()
        logits = logit_scale * v_emb @ t_emb.T

        # InfoNCE 双向损失
        labels = torch.arange(logits.size(0), device=logits.device)
        loss_v2t = F.cross_entropy(logits, labels)
        loss_t2v = F.cross_entropy(logits.T, labels)
        loss = (loss_v2t + loss_t2v) / 2
        return {"loss": loss, "logits": logits}

# CLIP4Clip 的 Mean Pooling vs CLS Pooling 对比
import torch
import torch.nn.functional as F

def cls_pooling(frame_features: torch.Tensor) -> torch.Tensor:
    """使用 CLS token（第一帧位置）作为视频表示"""
    return frame_features[:, 0, :]  # [B, D]

def mean_pooling(frame_features: torch.Tensor,
                 attention_mask: torch.Tensor = None) -> torch.Tensor:
    """对所有帧做平均池化（CLIP4Clip 方法）
    attention_mask: [B, T] 标记有效帧
    """
    if attention_mask is not None:
        # 只平均有效帧
        mask = attention_mask.unsqueeze(-1).float()  # [B, T, 1]
        summed = (frame_features * mask).sum(dim=1)  # [B, D]
        counts = mask.sum(dim=1).clamp(min=1e-9)  # [B, 1]
        return summed / counts
    return frame_features.mean(dim=1)  # [B, D]

def max_pooling(frame_features: torch.Tensor) -> torch.Tensor:
    """对所有帧做最大池化（捕获关键帧信息）"""
    return frame_features.max(dim=1)[0]  # [B, D]

# 实验比较三种池化策略
frame_features = torch.randn(32, 64, 512)  # 32个样本，每段64帧
mask = torch.ones(32, 64)
# 模拟：后20帧为填充帧
mask[:, 44:] = 0

cls_emb = cls_pooling(frame_features)
mean_emb = mean_pooling(frame_features, mask)
max_emb = max_pooling(frame_features)

print(f"CLS pooling:  {cls_emb.shape}, 取第 0 帧")
print(f"Mean pooling: {mean_emb.shape}, 平均有效 44 帧")
print(f"Max pooling:  {max_emb.shape}, 取各维度最大值")

import torch
import torch.nn as nn
from transformers import CLIPVisionModel, LlamaForCausalLM

class VideoLLaVA(nn.Module):
    """Video-LLaVA: 将视频 token 注入大语言模型"""

    def __init__(self, vision_model_name: str, llm_name: str,
                 num_frames: int = 8):
        super().__init__()
        self.num_frames = num_frames
        self.vision_model = CLIPVisionModel.from_pretrained(vision_model_name)
        self.llm = LlamaForCausalLM.from_pretrained(llm_name)

        # 冻结大语言模型和视觉编码器
        for param in self.vision_model.parameters():
            param.requires_grad = False
        for param in self.llm.parameters():
            param.requires_grad = False

        # 视觉-语言对齐投影层
        hidden_dim = self.llm.config.hidden_size
        vision_dim = self.vision_model.config.hidden_size
        self.vision_proj = nn.Sequential(
            nn.Linear(vision_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, hidden_dim),
        )

    def encode_video(self, frames: torch.Tensor) -> torch.Tensor:
        # frames: [B, T, C, H, W]
        B, T, C, H, W = frames.shape
        frames = frames.view(B * T, C, H, W)
        # 提取 patch 特征: [B*T, N, D_v]
        vision_out = self.vision_model(frames, output_hidden_states=True)
        features = vision_out.hidden_states[-1][:, 1:, :]  # 去掉 CLS token
        features = features.view(B, T, -1, vision_out.hidden_states[-1].size(-1))
        # 平均池化时间维度
        features = features.mean(dim=1)  # [B, N, D_v]
        return self.vision_proj(features)  # [B, N, D_llm]

    def forward(self, frames, input_ids, attention_mask, labels):
        video_tokens = self.encode_video(frames)  # [B, N_v, D]
        # 拼接视频 token 和文本 token 的嵌入
        text_embeds = self.llm.get_input_embeddings()(input_ids)
        combined = torch.cat([video_tokens, text_embeds], dim=1)
        # 更新 attention mask
        video_mask = torch.ones(frames.size(0), video_tokens.size(1),
                                device=attention_mask.device, dtype=attention_mask.dtype)
        full_mask = torch.cat([video_mask, attention_mask], dim=1)
        return self.llm(inputs_embeds=combined, attention_mask=full_mask, labels=labels)

# VideoQA 评估指标计算
from collections import Counter
import re

def exact_match_accuracy(predictions: list[str],
                         ground_truths: list[list[str]]) -> float:
    """精确匹配准确率
    每个样本可能有多个可接受的标准答案
    """
    def normalize(text: str) -> str:
        text = text.lower().strip()
        text = re.sub(r'[^a-z0-9s]', '', text)
        return text.strip()

    correct = 0
    for pred, truths in zip(predictions, ground_truths):
        pred_norm = normalize(pred)
        for truth in truths:
            if pred_norm == normalize(truth):
                correct += 1
                break
    return correct / len(predictions)

def open_ended_accuracy(pred: str, answers: list[str]) -> float:
    """开放式问题准确率（类似 VQA v2 的评分）
    答案出现 3 次或以上计 1 分，否则计 (出现次数 * 0.3)
    """
    pred_norm = pred.lower().strip()
    answer_counts = Counter(a.lower().strip() for a in answers)
    if answer_counts[pred_norm] >= 3:
        return 1.0
    return min(answer_counts[pred_norm] * 0.3, 1.0)

# 示例评估
predictions = ["playing basketball", "cooking", "dancing"]
ground_truths = [
    ["playing basketball", "basketball"],
    ["cooking", "preparing food"],
    ["dancing", "doing a dance"]
]
acc = exact_match_accuracy(predictions, ground_truths)
print(f"Exact Match Accuracy: {acc:.2%}")

import torch
import torch.nn as nn
from transformers import GPT2LMHeadModel

class VideoCaptioner(nn.Module):
    """基于预训练语言模型的视频描述生成器"""

    def __init__(self, vision_encoder, llm_model: str = "gpt2"):
        super().__init__()
        self.vision_encoder = vision_encoder
        self.llm = GPT2LMHeadModel.from_pretrained(llm_model)
        llm_dim = self.llm.config.hidden_size

        # 多层投影将视觉特征映射到语言空间
        vision_dim = vision_encoder.feature_dim
        self.vision_proj = nn.Sequential(
            nn.Linear(vision_dim, llm_dim),
            nn.LayerNorm(llm_dim),
            nn.GELU(),
            nn.Linear(llm_dim, llm_dim),
        )
        # 特殊 token 表示视频输入的起始
        self.video_token_id = self.llm.config.vocab_size  # 新增 token

    def generate_caption(self, video_frames: torch.Tensor,
                        max_length: int = 50, temperature: float = 0.7) -> str:
        # 编码视频: [B, T, C, H, W] -> [B, N, D]
        video_features = self.vision_encoder(video_frames)
        video_tokens = self.vision_proj(video_features)  # [B, N, D_llm]

        # 生成起始 token: [BOS]
        batch_size = video_tokens.size(0)
        decoder_input = torch.full((batch_size, 1),
                                   self.llm.config.bos_token_id,
                                   device=video_tokens.device)

        # 将视频 token 作为前缀输入
        inputs_embeds = torch.cat([video_tokens, self.llm.wte(decoder_input)], dim=1)

        # 自回归生成
        output = self.llm.generate(
            inputs_embeds=inputs_embeds,
            max_length=video_tokens.size(1) + max_length,
            temperature=temperature,
            do_sample=True,
            pad_token_id=self.llm.config.eos_token_id,
        )
        # 解码为文本（跳过视频 token 部分）
        text_tokens = output[:, video_tokens.size(1):]
        return self.llm.tokenizer.decode(text_tokens[0], skip_special_tokens=True)

# 评估指标：CIDEr 实现
import numpy as np
from collections import Counter
import math

class CIDErScorer:
    """CIDEr (Consensus-based Image Description Evaluation)
    使用 TF-IDF 加权 n-gram 匹配度评估描述质量
    """

    def __init__(self, n: int = 4, sigma: float = 6.0):
        self.n = n
        self.sigma = sigma
        self.document_frequency = Counter()  # DF: 每个 n-gram 出现在多少文档中
        self.num_documents = 0

    def fit(self, references: list[list[str]]):
        """从参考描述中学习 IDF 权重
        references: 每个样本的参考描述列表
        """
        self.num_documents = len(references)
        for refs in references:
            # 同一样本的多个参考视为一个文档
            seen = set()
            for ref in refs:
                for ngram in self._extract_ngrams(ref):
                    if ngram not in seen:
                        self.document_frequency[ngram] += 1
                        seen.add(ngram)

    def _extract_ngrams(self, sentence: str) -> list[str]:
        words = sentence.lower().split()
        ngrams = []
        for order in range(1, self.n + 1):
            for i in range(len(words) - order + 1):
                ngrams.append(" ".join(words[i:i + order]))
        return ngrams

    def score(self, hypothesis: str, references: list[str]) -> float:
        hyp_ngrams = Counter(self._extract_ngrams(hypothesis))
        if not hyp_ngrams:
            return 0.0

        score = 0.0
        for n in range(1, self.n + 1):
            # 计算当前 order 的 CIDEr 分数
            cand_ngrams = {k: v for k, v in hyp_ngrams.items() if k.count(" ") == n - 1}
            ref_counts = []
            for ref in references:
                ref_ngrams = Counter(self._extract_ngrams(ref))
                ref_counts.append({k: v for k, v in ref_ngrams.items() if k.count(" ") == n - 1})

            # TF-IDF 加权
            ngram_score = 0.0
            for ngram, count in cand_ngrams.items():
                tf = count
                idf = math.log(self.num_documents / (1 + self.document_frequency.get(ngram, 0)))
                # 与参考的共识
                max_ref = max(rc.get(ngram, 0) for rc in ref_counts)
                ngram_score += tf * idf * min(count, max_ref) * idf
            score += ngram_score

        return score / self.n

# 使用示例
scorer = CIDErScorer(n=4)
refs = [["a person is cooking in the kitchen"], ["someone prepares food"]]
scorer.fit(refs)
score = scorer.score("a person is cooking in a kitchen", refs[0])
print(f"CIDEr: {score:.4f}")

import torch
import torch.nn as nn
from torch.nn import functional as F

class MomentDETR(nn.Module):
    """基于 DETR 的时序定位模型"""

    def __init__(self, d_model: int = 256, num_moments: int = 10,
                 num_encoder_layers: int = 4, num_decoder_layers: int = 4,
                 nhead: int = 8):
        super().__init__()
        self.num_moments = num_moments

        # 视频-文本交叉注意力编码器
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=d_model, nhead=nhead, dim_feedforward=d_model * 4,
            batch_first=True)
        self.encoder = nn.TransformerEncoder(encoder_layer, num_encoder_layers)

        # 解码器：使用可学习查询预测时间段
        decoder_layer = nn.TransformerDecoderLayer(
            d_model=d_model, nhead=nhead, dim_feedforward=d_model * 4,
            batch_first=True)
        self.decoder = nn.TransformerDecoder(decoder_layer, num_decoder_layers)

        # 可学习的时刻查询
        self.moment_queries = nn.Parameter(torch.randn(num_moments, d_model))

        # 预测头：起止时间和 IoU 分数
        self.start_head = nn.Linear(d_model, 1)
        self.end_head = nn.Linear(d_model, 1)
        self.iou_head = nn.Linear(d_model, 1)

        # 位置编码
        self.pos_embed = nn.Parameter(torch.randn(1, 1, d_model))

    def forward(self, video_features: torch.Tensor,
                text_features: torch.Tensor) -> dict:
        B, T, D_v = video_features.shape

        # 融合视频和文本特征
        video_fused = video_features + text_features.mean(dim=1, keepdim=True)
        video_fused = video_fused + self.pos_embed[:, :T, :]

        # 编码
        memory = self.encoder(video_fused)

        # 解码：可学习查询与记忆交互
        queries = self.moment_queries.unsqueeze(0).expand(B, -1, -1)
        decoder_out = self.decoder(queries, memory)

        # 预测
        start_logits = self.start_head(decoder_out).squeeze(-1)  # [B, Q]
        end_logits = self.end_head(decoder_out).squeeze(-1)  # [B, Q]
        iou_scores = self.iou_head(decoder_out).squeeze(-1)  # [B, Q]

        return {
            "start_logits": start_logits,
            "end_logits": end_logits,
            "iou_scores": iou_scores,
            "queries": decoder_out
        }

# 时序 IoU 计算和评估指标
import numpy as np

def temporal_iou(span_a: tuple, span_b: tuple) -> float:
    """计算两个时间段的交并比"""
    start_a, end_a = span_a
    start_b, end_b = span_b
    intersection_start = max(start_a, start_b)
    intersection_end = min(end_a, end_b)
    if intersection_start >= intersection_end:
        return 0.0
    intersection = intersection_end - intersection_start
    union = (end_a - start_a) + (end_b - start_b) - intersection
    return intersection / union if union > 0 else 0.0

def recall_at_iou(predictions: list[tuple],
                  ground_truths: list[tuple],
                  iou_threshold: float = 0.5) -> float:
    """计算 Recall@IoU
    predictions: 预测的起止时间列表
    ground_truths: 真实的起止时间列表
    """
    assert len(predictions) == len(ground_truths)
    correct = 0
    for pred, gt in zip(predictions, ground_truths):
        if temporal_iou(pred, gt) >= iou_threshold:
            correct += 1
    return correct / len(predictions)

def mean_iou(predictions: list[tuple],
             ground_truths: list[tuple]) -> float:
    """计算平均 IoU"""
    ious = [temporal_iou(p, g) for p, g in zip(predictions, ground_truths)]
    return np.mean(ious)

# 示例
predictions = [(2.5, 8.3), (15.0, 22.1), (30.0, 35.5)]
ground_truths = [(3.0, 8.0), (14.5, 21.0), (28.0, 36.0)]
for thresh in [0.3, 0.5, 0.7]:
    r = recall_at_iou(predictions, ground_truths, thresh)
    print(f"Recall@{thresh}: {r:.2%}")
print(f"Mean IoU: {mean_iou(predictions, ground_truths):.4f}")

import torch
import torch.nn as nn

class QFormerVideoCompressor(nn.Module):
    """使用 Q-Former 压缩视频特征"""

    def __init__(self, vision_dim: int, llm_dim: int,
                 num_queries: int = 128, num_layers: int = 3,
                 num_heads: int = 8):
        super().__init__()
        self.num_queries = num_queries

        # 可学习的查询向量
        self.queries = nn.Parameter(torch.randn(num_queries, llm_dim))

        # 视觉特征投影
        self.vision_proj = nn.Linear(vision_dim, llm_dim)

        # Q-Former 编码器
        encoder_layer = nn.TransformerEncoderLayer(
            d_model=llm_dim, nhead=num_heads,
            dim_feedforward=llm_dim * 4,
            batch_first=True)
        self.qformer = nn.TransformerEncoder(encoder_layer, num_layers)

        # LayerNorm
        self.norm = nn.LayerNorm(llm_dim)

    def forward(self, video_features: torch.Tensor,
                text_features: torch.Tensor = None) -> torch.Tensor:
        """
        video_features: [B, T, N, D_vision]  (T帧，每帧N个patch)
        text_features:  [B, L, D_text] (可选的文本引导)
        返回: [B, num_queries, D_llm]
        """
        B, T, N, D = video_features.shape

        # 展平时空维度: [B, T*N, D]
        video_flat = video_features.reshape(B, T * N, D)
        video_tokens = self.vision_proj(video_flat)  # [B, T*N, D_llm]

        # 扩展查询
        queries = self.queries.unsqueeze(0).expand(B, -1, -1)

        # 拼接查询和视频token，通过 Transformer 进行交叉注意力
        # 查询在前，视频token在后
        concat = torch.cat([queries, video_tokens], dim=1)

        # 通过 Q-Former 处理
        output = self.qformer(concat)

        # 提取查询部分的结果
        compressed = output[:, :self.num_queries, :]
        return self.norm(compressed)

# 视频 token 压缩策略对比
import torch
import time

def benchmark_compression_methods(num_frames: int = 64,
                                   patches_per_frame: int = 256,
                                   dim: int = 1024):
    """对比不同压缩策略的 token 数量和计算时间"""
    video_features = torch.randn(1, num_frames, patches_per_frame, dim)
    total_tokens = num_frames * patches_per_frame

    strategies = {
        "无压缩": {
            "output_tokens": total_tokens,
            "compression_ratio": 1.0,
        },
        "均匀采样(8帧)": {
            "output_tokens": 8 * patches_per_frame,
            "compression_ratio": total_tokens / (8 * patches_per_frame),
        },
        "Q-Former(128查询)": {
            "output_tokens": 128,
            "compression_ratio": total_tokens / 128,
        },
        "Perceiver(64查询)": {
            "output_tokens": 64,
            "compression_ratio": total_tokens / 64,
        },
        "时空池化(4x4)": {
            "output_tokens": (num_frames // 4) * (patches_per_frame // 4),
            "compression_ratio": total_tokens / ((num_frames // 4) * (patches_per_frame // 4)),
        },
    }

    print(f"输入: {num_frames} 帧 x {patches_per_frame} patch = {total_tokens:,} token")
    print("-" * 60)
    for name, info in strategies.items():
        print(f"{name:25s} | {info['output_tokens']:6,} token | "
              f"压缩比: {info['compression_ratio']:.1f}x")

benchmark_compression_methods(64, 256, 1024)

import torch
import torchvision.transforms as transforms
from PIL import Image
import numpy as np

class VideoCLIPInference:
    """VideoCLIP 零样本视频检索推理"""

    def __init__(self, model, preprocess):
        self.model = model
        self.model.eval()
        self.preprocess = preprocess
        # 视频预处理：与图像相同的变换
        self.video_transform = transforms.Compose([
            transforms.Resize((224, 224)),
            transforms.ToTensor(),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225]),
        ])

    def encode_video(self, frames: list[Image.Image],
                     num_frames: int = 8) -> torch.Tensor:
        """将视频帧编码为统一的嵌入向量"""
        # 均匀采样 num_frames 帧
        if len(frames) < num_frames:
            indices = list(range(len(frames)))
        else:
            indices = np.linspace(0, len(frames) - 1, num_frames, dtype=int)

        selected_frames = [frames[i] for i in indices]
        processed = torch.stack([self.preprocess(f) for f in selected_frames])
        processed = processed.unsqueeze(0)  # [1, T, C, H, W]

        with torch.no_grad():
            video_features = self.model.encode_video(processed)
            video_features = video_features / video_features.norm(dim=-1, keepdim=True)
        return video_features

    def encode_text(self, texts: list[str]) -> torch.Tensor:
        """将文本编码为嵌入向量"""
        with torch.no_grad():
            text_features = self.model.encode_text(texts)
            text_features = text_features / text_features.norm(dim=-1, keepdim=True)
        return text_features

    def search(self, query: str, video_database: dict[str, torch.Tensor],
               top_k: int = 5) -> list[tuple[str, float]]:
        """在视频库中检索最相关的视频"""
        query_embed = self.encode_text([query])  # [1, D]
        # 计算与所有视频嵌入的相似度
        video_ids = list(video_database.keys())
        video_embeds = torch.stack([video_database[vid] for vid in video_ids])
        similarities = (query_embed @ video_embeds.T).squeeze(0)
        # Top-K
        top_indices = similarities.argsort(descending=True)[:top_k]
        results = []
        for idx in top_indices:
            results.append((video_ids[idx], similarities[idx].item()))
        return results

# 使用示例
# videos_db = {"video_001": embed1, "video_002": embed2, ...}
# results = searcher.search("a dog playing in the park", videos_db, top_k=5)
# for vid, score in results:
#     print(f"{vid}: {score:.4f}")

# 完整的视频检索 Pipeline
import os
import cv2
import torch
from typing import Optional

def extract_frames_from_video(video_path: str,
                              num_frames: int = 8) -> list:
    """从视频文件中均匀提取 num_frames 帧"""
    cap = cv2.VideoCapture(video_path)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))

    if total_frames == 0:
        cap.release()
        return []

    indices = np.linspace(0, total_frames - 1, num_frames, dtype=int)
    frames = []
    for idx in indices:
        cap.set(cv2.CAP_PROP_POS_FRAMES, idx)
        ret, frame = cap.read()
        if ret:
            frames.append(Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)))
    cap.release()
    return frames

def build_video_database(video_dir: str,
                         inference: VideoCLIPInference,
                         num_frames: int = 8) -> dict[str, torch.Tensor]:
    """批量编码视频目录中的所有视频"""
    database = {}
    video_files = [f for f in os.listdir(video_dir)
                   if f.endswith((".mp4", ".avi", ".mov"))]

    for i, vf in enumerate(video_files):
        path = os.path.join(video_dir, vf)
        frames = extract_frames_from_video(path, num_frames)
        if frames:
            embed = inference.encode_video(frames)
            database[vf] = embed.squeeze(0)
        if (i + 1) % 10 == 0:
            print(f"已编码 {i + 1}/{len(video_files)} 个视频")

    return database

def multi_query_search(queries: list[str],
                       video_db: dict,
                       inference: VideoCLIPInference,
                       top_k: int = 3) -> dict[str, list[tuple]]:
    """多查询批量检索"""
    all_results = {}
    # 批量编码所有查询
    all_query_embeds = inference.encode_text(queries)

    video_ids = list(video_db.keys())
    video_embeds = torch.stack([video_db[vid] for vid in video_ids])

    for i, query in enumerate(queries):
        query_embed = all_query_embeds[i:i+1]
        sims = (query_embed @ video_embeds.T).squeeze(0)
        top_idx = sims.argsort(descending=True)[:top_k]
        all_results[query] = [
            (video_ids[j], sims[j].item()) for j in top_idx
        ]

    return all_results