视频帧描述与摘要生成示例

本示例展示如何使用Scalar UDF和UDAF对视频帧进行内容描述与摘要生成。通过av库从真实视频文件中提取帧，存入数据库表后，使用BLIP模型生成每帧的自然语言描述，并使用T5模型将多帧描述汇总为视频摘要。

场景描述

在视频内容分析场景中，需要从视频中提取关键帧，对每帧生成自然语言描述，并将多帧描述综合为连贯的视频摘要。通过Class形式的Scalar UDF，在__init__中加载BLIP模型，在__call__中对每帧执行图像描述生成。通过UDAF（用户定义聚合函数），将多帧描述聚合后使用T5模型生成摘要。本示例使用av库从真实视频文件中提取帧作为测试数据。

约束与限制

UDF运行时需通过pip安装transformers包，请确保运行时环境网络可用。
Scalar UDF的__init__参数仅支持传入常量。
注册UDF前建议先调用delete_function清理同名函数，避免重复执行时报错。

前提条件

已安装Aura DataFrame SDK。具体操作，请参见Aura DataFrame SDK获取与安装。
已创建AuraJobV2端点，并获取端点名称。具体操作，请参见创建Aura引擎端点。
已创建LakeFormation实例，并在该实例中创建Catalog及数据库。具体操作，请参见创建LakeFormation实例并规划元数据。
已创建并行文件系统。具体操作，请参见创建并行文件系统。
代码中涉及os.getenv()中的参数已配置到环境变量。
运行时环境要求Python版本为3.11。

准备工作

本示例使用BLIP图像描述模型和T5文本摘要模型，需要准备模型文件、视频文件和Python依赖。

安装Python依赖：

pip install torch==2.6.0
pip install transformers==5.3.0 
pip install pillow 
pip install av

安装huggingface命令行工具：
```
pip install huggingface-hub
```
下载测试视频文件到示例脚本同级目录，重命名为test_video.mp4：
下载一段包含真实场景内容的MP4格式短视频，建议10秒以内。
下载BLIP图像描述模型文件到示例脚本同级目录：
```
hf download Salesforce/blip-image-captioning-base --local-dir ./blip-image-captioning-base
```
如果hf download命令无法使用，可通过浏览器手动下载：
1. 在Hugging Face网站搜索 Salesforce/blip-image-captioning-base，进入模型主页。
2. 单击“Files and versions”标签页。
3. 下载以下文件，保存到本地 blip-image-captioning-base 目录下：
  - config.json
  - preprocessor_config.json
  - pytorch_model.bin
  - special_tokens_map.json
  - tokenizer_config.json
  - tokenizer.json
  - vocab.txt
4. 将blip-image-captioning-base目录放置在与示例脚本同级目录下。
下载T5文本摘要模型文件到示例脚本同级目录：
```
hf download google-t5/t5-small --local-dir ./t5-small
```
如果hf download命令无法使用，可通过浏览器手动下载：
1. 在Hugging Face网站搜索google-t5/t5-small，进入模型主页。
2. 单击“Files and versions”标签页。
3. 下载以下文件，保存到本地t5-small目录下：
  - config.json
  - pytorch_model.bin
  - spiece.model
  - tokenizer.json
  - tokenizer_config.json
  - generation_config.json
4. 将t5-small目录放置在与示例脚本同级目录下。

示例

本示例将视频帧描述与摘要生成的完整流程整合为一段可独立运行的代码，包含连接管理、Scalar UDF定义（基于BLIP模型生成帧描述）、UDAF定义（基于T5模型生成视频摘要）、函数注册、视频帧提取、帧级描述生成和视频摘要聚合的全过程。

导入依赖并定义create_connect上下文管理器，封装Aura服务连接的创建与释放。
定义Scalar UDF类VideoContentAnalyzer，在__init__中加载BLIP模型，在__call__中对每帧JPEG数据生成描述并分析内容特征，返回JSON格式结果。
定义VideoSummaryGenerator类，加载T5模型，将多帧描述聚合后生成视频摘要。
定义UDAF类VideoSummarySystem，通过accumulate收集每帧描述，finish中调用T5模型生成视频摘要。
连接AuraJobV2端点并设置UDF工作空间，删除已存在的同名UDF和UDAF后注册到目标数据库。
用av库从视频文件中提取最多10帧（JPEG格式），创建video_frames表写入二进制帧数据。
加载video_frames数据集，通过map_batches调用Scalar UDF分析每帧生成descriptions列。
通过AggregateFnBuilder调用UDAF聚合多帧描述生成视频摘要，执行查询并打印结果。

import ibis
import typing
import os
from contextlib import contextmanager
from aura_frame.multimodal import ai_lake
from aura_frame.multimodal.function import AggregateFnBuilder

target_database = "test"  #在LakeFormation中创建数据库，例如：test。具体操作，请参见创建LakeFormation实例并规划元数据
@contextmanager
def create_connect():
#连接AuraJobV2端点，具体参数说明，请参见建立与AuraJobV2端点的连接
    con = ai_lake.connect(
        aura_endpoint=os.getenv("aura_endpoint"),         # AuraJobV2端点所在的终端节点（Endpoint）
        aura_endpoint_id=os.getenv("aura_endpoint_id"),    # AuraJobV2端点ID
        aura_endpoint_name=os.getenv("aura_endpoint_name"),# AuraJobV2端点名称
        aura_workspace_id=os.getenv("aura_workspace_id"),  # 工作空间ID
        lf_catalog_name=os.getenv("lf_catalog_name"),      # LakeFormation Catalog名称
        access_key=os.getenv("access_key"),                # 访问密钥ID（AK）
        secret_key=os.getenv("secret_key"),                # 访问密钥（SK）
        default_database=target_database,                       # LakeFormation数据库名称
        use_single_cn_mode=True,   # 是否开启单CN模式
        logging_level=logging.INFO # 设置日志级别
    )
    print(f"Establish ai_lake.connect: {con}")
    try:
        yield con
    finally:
        # 关闭连接
        con.close()

class VideoContentAnalyzer:
    def __init__(self, model_name="./blip-image-captioning-base"):
        import torch
        from transformers import BlipProcessor, BlipForConditionalGeneration
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.processor = BlipProcessor.from_pretrained(model_name)
        self.model = BlipForConditionalGeneration.from_pretrained(
            model_name,
            torch_dtype=torch.float16 if self.device == "cuda" else torch.float32,
        ).to(self.device)
        self.model.eval()
    def __call__(self, frame: bytes, frame_index: int) -> str:
        import json
        analyse = self.analyze_frame(frame, frame_index)
        return json.dumps({"description": analyse})
    def analyze_frame(self, frame: bytes, frame_index: int):
        try:
            frame_description = self._generate_frame_description(frame)
            content_features = self._analyze_content_features(frame_description)
            analyse = {
                "description": frame_description,
                "content_type": content_features["content_type"],
                "key_objects": content_features["key_objects"],
                "scene_context": content_features["scene_context"],
            }
        except Exception:
            analyse = {}
        return analyse
    def _generate_frame_description(self, frame: bytes):
        import io
        from PIL import Image
        import torch
        pil_img = Image.open(io.BytesIO(frame)).convert("RGB")
        inputs = self.processor(images=pil_img, return_tensors="pt").to(self.device)
        with torch.no_grad():
            outputs = self.model.generate(
                **inputs,
                max_length=50,
                num_beams=5,
                early_stopping=True,
            )
        description = self.processor.decode(outputs[0], skip_special_tokens=True)
        return description.strip()
    def _analyze_content_features(self, description):
        description_lower = description.lower()
        content_categories = {
            "educational": ["tutorial", "lesson", "explain", "teach", "learn", "education"],
            "entertainment": ["funny", "comedy", "entertainment", "show", "performance"],
            "nature": ["outdoor", "nature", "landscape", "mountain", "forest", "animal"],
            "sports": ["sport", "game", "player", "team", "match", "competition"],
            "food": ["food", "cooking", "recipe", "meal", "restaurant", "kitchen"],
            "technology": ["computer", "tech", "device", "electronic", "software"],
            "people": ["person", "people", "man", "woman", "child", "group"],
        }
        detected_categories = []
        for category, keywords in content_categories.items():
            if any(keyword in description_lower for keyword in keywords):
                detected_categories.append(category)
        common_objects = [
            "person", "people", "man", "woman", "child", "car", "building",
            "tree", "house", "street", "water", "sky", "food", "animal",
        ]
        key_objects = [obj for obj in common_objects if obj in description_lower]
        if any(word in description_lower for word in ["indoors", "inside", "room"]):
            scene_context = "indoor"
        elif any(word in description_lower for word in ["outdoors", "outside", "park", "street"]):
            scene_context = "outdoor"
        else:
            scene_context = "unknown"
        return {
            "content_type": detected_categories[:2] if detected_categories else ["general"],
            "key_objects": key_objects[:5],
            "scene_context": scene_context,
        }

class VideoSummaryGenerator:
    def __init__(self, model):
        from transformers import pipeline
        import torch
        self.summarizer = pipeline(
            "text-generation",
            model=model,
            device=0 if torch.cuda.is_available() else -1,
        )
    def generate_summary(self, frame_analyses):
        all_descriptions = [analysis.get("description", "") for analysis in frame_analyses]
        if not all_descriptions:
            return "Unable to analyze the video content."
        content_analysis = self._analyze_video_content(frame_analyses)
        description_text = self._create_description_text(all_descriptions, content_analysis)
        summary = self._generate_detailed_summary(description_text, content_analysis)
        return summary
    def _analyze_video_content(self, frame_analyses):
        from collections import Counter
        content_types = []
        all_objects = []
        for analysis in frame_analyses:
            content_types.extend(analysis.get("content_type", []))
            all_objects.extend(analysis.get("key_objects", []))
        common_content_types = [item for item, count in Counter(content_types).most_common(2)]
        common_objects = [item for item, count in Counter(all_objects).most_common(5)]
        return {
            "primary_categories": common_content_types,
            "dominant_objects": common_objects,
            "total_frames_analyzed": len(frame_analyses),
        }
    def _create_description_text(self, descriptions, content_analysis):
        from collections import Counter
        desc_counter = Counter(descriptions)
        representative_descs = [desc for desc, count in desc_counter.most_common(3)]
        text = "Video Content Analysis:\n"
        text += f"Primary Categories: {', '.join(content_analysis['primary_categories'])}\n"
        text += f"Key Objects: {', '.join(content_analysis['dominant_objects'][:3])}\n"
        text += "Representative Scenes:\n"
        for i, desc in enumerate(representative_descs, 1):
            text += f"{i}. {desc}\n"
        return text
    def _generate_concise_summary(self, content_analysis):
        import random
        primary_category = content_analysis["primary_categories"][0] if content_analysis["primary_categories"] else "General content"
        main_objects = ", ".join(content_analysis["dominant_objects"][:2])
        templates = [
            f"This is a {primary_category} video that includes elements like {main_objects}.",
            f"The video presents a {primary_category} scenario, with its focus on {main_objects}.",
            f"The content centers on {primary_category}, highlighting elements such as {main_objects}.",
        ]
        return random.choice(templates)
    def _generate_detailed_summary(self, description_text, content_analysis):
        try:
            summary = self.summarizer(
                "summarize: " + description_text,
                max_length=150,
                min_length=50,
                do_sample=False,
            )[0]["generated_text"]
            return summary
        except Exception:
            return self._generate_concise_summary(content_analysis)

class VideoSummarySystem:
    def __init__(self, model="./t5-small"):
        self.summary_generator = VideoSummaryGenerator(model)
        self.frame_descriptions = []
    def accumulate(self, description: str) -> None:
        import json
        data = json.loads(description)
        self.frame_descriptions.append(data.get("description", data))
    def finish(self) -> str:
        import json
        summary = self.summary_generator.generate_summary(self.frame_descriptions)
        result = {
            "video_summary": summary,
            "frames_analyzed": len(self.frame_descriptions),
            "content_categories": list(
                set(cat for analysis in self.frame_descriptions for cat in analysis.get("content_type", []))
            ),
            "key_objects": list(
                set(obj for analysis in self.frame_descriptions for obj in analysis.get("key_objects", []))
            )[:8],
            "processing_details": {
                "total_frames": len(self.frame_descriptions),
                "successful_analyses": len(self.frame_descriptions),
            },
        }
        return json.dumps(result)
    @property
    def aggregate_state(self) -> typing.Dict[str, typing.Any]:
        return {"frame_descriptions": self.frame_descriptions}
    def merge(self, other_state: typing.Dict[str, typing.Any]) -> None:
        self.frame_descriptions += other_state.get("frame_descriptions")

with create_connect() as con:
    con.set_function_staging_workspace(
        obs_directory_base=os.getenv("obs_directory_base"),
        obs_bucket_name=os.getenv("obs_bucket_name"),
        obs_server=os.getenv("obs_server"),
        access_key=os.getenv("access_key"),
        secret_key=os.getenv("secret_key")
    )
    print("finish set_function_staging_workspace")
    con.delete_function("videocontentanalyzer", database=target_database, if_it_exists=True)
    con.create_scalar_function(
        VideoContentAnalyzer,
        database=target_database,
        packages=["transformers==5.3.0"],
        imports=["./blip-image-captioning-base"],
    )
    con.delete_function("videosummarysystem", database=target_database, if_it_exists=True)
    con.create_agg_function(
        VideoSummarySystem,
        database=target_database,
        packages=["transformers==5.3.0"],
        imports=["./t5-small"],
    )
    # 用 av 从真实视频文件提取帧
    import av
    import io
    from PIL import Image
    video_path = "test_video.mp4"
    container = av.open(video_path)
    stream = container.streams.video[0]
    frames = []
    frame_count = 0
    max_frames = 10
    for frame in container.decode(stream):
        if frame_count >= max_frames:
            break
        img = frame.to_image()
        buf = io.BytesIO()
        img.save(buf, format="JPEG")
        frames.append((frame_count, buf.getvalue()))
        frame_count += 1
    container.close()
    print(f"从视频中提取了 {len(frames)} 帧")
    # 建表并插入帧数据
    table_name = "video_frames"
    table_schema = ibis.schema({"frame_index": "int32", "frame_data": "binary"})
    con.sql(f'DROP TABLE IF EXISTS "{target_database}"."{table_name}";')
    con.create_table(
        table_name,
        schema=table_schema,
        database=target_database,
        table_format="orc",
        if_not_exists=True,
    )
    for frame_index, frame_bytes in frames:
        hex_str = frame_bytes.hex()
        con.sql(f'''
        INSERT INTO "{target_database}"."{table_name}" (frame_index, frame_data) VALUES
            ({frame_index}, '\\x{hex_str}'::bytea)
        ''')
    print("视频帧数据写入完成")
    # 加载视频帧数据集
    ds = con.load_dataset("video_frames", database=target_database)
    ds.show(limit=1)
    # 使用 Scalar UDF 分析每个帧
    udf = con.get_function("videocontentanalyzer", database=target_database)
    ds = ds.map_batches(
        fn=udf,
        on=[ds.frame_data, ds.frame_index],
        as_col="descriptions",
        num_cpus=2,
        memory=4096,
        arch="x86",
        model_name="./blip-image-captioning-base",
    )
    # ds.select_columns([ds.frame_index, ds.descriptions]).show(1)
    # 使用UDAF生成视频摘要
    udaf = con.get_function("videosummarysystem", database=target_database)
    agg_builder = AggregateFnBuilder(
        fn=udaf,
        on=[ds.descriptions],
        as_col="describe",
        num_cpus=1,
        memory=2048,
        arch="x86",
    )
    summary = ds.aggregate(agg_builder).execute()
    for row in summary.itertuples():
        print(row)