llm_image.py 6.98 KB
import base64
import shutil
import subprocess
import sys
from pathlib import Path
from typing import Optional
from urllib.parse import urlparse
import cv2
from aabd.stream_chain import video2frames
import json
import logging
import hashlib

logger = logging.getLogger(__name__)

_IS_ABSOLUTE_UTC_THRESHOLD = 1_000_000_000


def download_to_mp4(url: str, output_path: str, duration: Optional[int] = None) -> None:
    cmd = ["ffmpeg", "-y", "-fflags", "+discardcorrupt", "-i", url]
    if duration is not None:
        cmd.extend(["-t", str(duration)])
    cmd.extend(["-c", "copy", "-bsf:a", "aac_adtstoasc", "-movflags", "+faststart", output_path])
    result = subprocess.run(cmd, capture_output=True, text=True, encoding="utf-8", errors="replace")
    if result.returncode != 0:
        raise RuntimeError(f"ffmpeg 下载失败: {result.stderr}")


def resize_frame(frame, max_px_area):
    if max_px_area is None:
        return frame
    h, w = frame.shape[:2]
    area = h * w
    if area > max_px_area:
        scale = (max_px_area / area) ** 0.5
        new_w = int(w * scale)
        new_h = int(h * scale)
        frame = cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_AREA)
    return frame


class Video2Frame:
    def __init__(self, cache_dir, save_frames_enable=False):
        self.cache_dir = Path(cache_dir)
        self.save_frames_enable = save_frames_enable

    def get_root_path(self, url, cache=None):
        if cache is None:
            parsed_url = urlparse(url)
            raw = f"{parsed_url.scheme}/{parsed_url.netloc}{parsed_url.path}"
            cache = hashlib.md5(raw.encode('utf-8')).hexdigest()
        return self.cache_dir / cache

    def iter_mp4(self, file, start, end, fps):

        if (start is not None and start > _IS_ABSOLUTE_UTC_THRESHOLD) or (
                end is not None and end > _IS_ABSOLUTE_UTC_THRESHOLD):
            video_start_time = 0
            with video2frames.AVAnyKeyStreamDecoder(file, sei_enable=True, tqdm_enable=False) as (sd, _):
                for frame in sd:
                    if start is not None:
                        video_start_time = start - frame.get('src_frame_time')
                    if end is not None:
                        video_end_time = end - frame.get('src_frame_time') or 0
                    break
        else:
            video_start_time = start if start is not None else 0
            video_end_time = end if end is not None else sys.maxsize

        with video2frames.AVAnyKeyStreamDecoder(file, start=video_start_time, end=video_end_time, control_type='time',
                                                sei_enable=True, max_fps=fps, frame_type='numpy_bgr') as (sd, _):
            yield from sd

    def to_frames(self, url, cache=None, fps=None, start=None, end=None, roi=None, max_px_area=None) -> list:
        video_root_path = self.get_root_path(url, cache)

        download_video_path = video_root_path / "video.mp4"

        if not download_video_path.exists():
            logger.info(f"download_video: {url} -> {download_video_path}")
            video_root_path.mkdir(parents=True, exist_ok=True)
            download_to_mp4(url, str(download_video_path), duration=None)
        logger.info(f"video_path: {download_video_path}")
        cache_dir = video_root_path / 'caches' / f"{start}-{end}-{'_'.join(roi) if roi else 'None'}-{max_px_area}"
        cache_frames = cache_dir / "frames"
        cache_name = cache_dir / "data.json"

        if cache_name.exists():
            logger.info(f"use_cache: {cache_name}")
            return json.load(cache_name.open('r'))
        else:
            cache_dir.mkdir(parents=True, exist_ok=True)
            data = {
                'fps': fps,
                'frames': []
            }
            if self.save_frames_enable:
                cache_frames.mkdir(parents=True, exist_ok=True)
            for f in self.iter_mp4(str(download_video_path), start, end, fps):
                frame = f['frame']
                time = f['src_frame_time']
                fps = f['fps']
                data['fps'] = fps
                if roi is not None:
                    x1, y1, x2, y2 = roi
                    frame = frame[y1:y2, x1:x2]
                if max_px_area is not None:
                    frame = resize_frame(frame, max_px_area)

                _, buffer = cv2.imencode('.jpg', frame)
                b64_str = base64.b64encode(buffer).decode('utf-8')

                data['frames'].append({'time': time, 'image': b64_str})

                if self.save_frames_enable:
                    frame_path = cache_frames / f"{time:015d}.jpg"
                    cv2.imwrite(frame_path, frame)

            cache_name.parent.mkdir(parents=True, exist_ok=True)
            json.dump(data, cache_name.open('w'), indent=4)
            logger.info(f"save_cache: {cache_name}")
            return data

    def to_llm_contents(self, url, cache=None, fps=None, start=None, end=None, roi=None, max_px_area=None,
                        prompt_start=None,
                        prompt_end=None):
        data = self.to_frames(url, cache, fps, start, end, roi, max_px_area)
        fps = data['fps']
        frames = data['frames']

        contents = []
        if prompt_start is not None:
            contents.append({"type": "text", "text": prompt_start})
        video_prompt = (
            f"以下是从视频中按时间顺序提取的 {len(frames)} 帧画面,fps={fps},请将它们视为一个连续的视频进行分析。"
        )
        contents.append({"type": "text", "text": video_prompt})

        for frame in frames:
            contents.append({
                "type": "image_url",
                "image_url": {
                    "url": f"data:image/jpeg;base64,{frame['image']}"
                }
            })
        if prompt_end is not None:
            contents.append({"type": "text", "text": prompt_end})
        return contents


if __name__ == '__main__':
    from aabd.base.patched_logging import init_logging
    import os
    os.environ["APP_LOG_TYPE"] = "console"
    init_logging()
    url = "http://video.mam.miguvideo.com/mnt6/fastclip3/wsc/2026/04/29/df760b8d15394cbc9ed0d0a4a9c4b786_1080PS/29133605/vodtmp/3b3e99cf4ca84c3782503d8817242de2.m3u8"
    # url = rf"R:\wdx\202604\20260420_download_football_video\finished\69dd5845dd0412067b8d5587-auto-1776074760653\live\videos\00-14-09-052.mp4"

    vf = Video2Frame(rf'D:\Code\migu\aigc-embedding-service\src\football_replay_match\core', save_frames_enable=True)
    a = vf.to_llm_contents(url, fps=2, start=1777444502543, end=1777444531063, max_px_area=1920 * 1080 // 6)
    from langchain_openai import ChatOpenAI
    from langchain_core.messages import SystemMessage, HumanMessage

    model = ChatOpenAI(base_url="http://192.168.1.59:11434/v1", model="Qwen3.6-35B-A3B-UD-Q8_K_XL.gguf",
                       temperature=0.7, api_key='no_key',
                       extra_body={"chat_template_kwargs": {"enable_thinking": False}})

    result = model.invoke([SystemMessage(content="描述一下视频的内容。"), HumanMessage(content=a)])
    print(result)