test_fun.py

import warnings

warnings.filterwarnings("ignore", category=FutureWarning)
import json
import os
import subprocess
import time

import numpy as np
import torch
import torch.nn.functional as F
from decord import VideoReader
from transformers import AutoModel, AutoVideoProcessor

import src.datasets.utils.video.transforms as video_transforms
import src.datasets.utils.video.volume_transforms as volume_transforms
from src.models.attentive_pooler import AttentiveClassifier
from src.models.vision_transformer import vit_giant_xformers_rope, vit_base

IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)


def load_pretrained_vjepa_pt_weights(model, pretrained_weights):
    # Load weights of the VJEPA2 encoder
    # The PyTorch state_dict is already preprocessed to have the right key names
    pretrained_dict = torch.load(pretrained_weights, weights_only=True, map_location="cpu")["encoder"]
    pretrained_dict = {k.replace("module.", ""): v for k, v in pretrained_dict.items()}
    pretrained_dict = {k.replace("backbone.", ""): v for k, v in pretrained_dict.items()}
    msg = model.load_state_dict(pretrained_dict, strict=False)
    print("Pretrained weights found at {} and loaded with msg: {}".format(pretrained_weights, msg))


def build_pt_video_transform(img_size):
    short_side_size = int(256.0 / 224 * img_size)
    # Eval transform has no random cropping nor flip
    eval_transform = video_transforms.Compose(
        [
            video_transforms.Resize(short_side_size, interpolation="bilinear"),
            video_transforms.CenterCrop(size=(img_size, img_size)),
            volume_transforms.ClipToTensor(),
            video_transforms.Normalize(mean=IMAGENET_DEFAULT_MEAN, std=IMAGENET_DEFAULT_STD),
        ]
    )
    return eval_transform


def get_video(sample_video_path, num_frames=80):
    vr = VideoReader(sample_video_path)
    total_frames = len(vr)
    # Choose evenly spaced frames, limited by available frames
    if total_frames < num_frames:
        frame_idx = np.arange(0, total_frames, 2)
    else:
        frame_idx = np.linspace(0, total_frames - 1, num_frames, dtype=int)
    video = vr.get_batch(frame_idx).asnumpy()
    return video


def forward_vjepa_video(model_pt, pt_transform, sample_video_path):
    # Run a sample inference with VJEPA
    with torch.inference_mode():
        # Read and pre-process the image
        video = get_video(sample_video_path)  # T x H x W x C
        video = torch.from_numpy(video).permute(0, 3, 1, 2)  # T x C x H x W
        print(video.shape)
        x_pt = pt_transform(video)[0].cuda().unsqueeze(0)
        print(x_pt.shape)
        # Extract the patch-wise features from the last layer
        out_patch_features_pt = model_pt(x_pt)

    return out_patch_features_pt



def run_single_sample_inference():
    # sample_video_path = "/mnt/data2/tzx/workspace/auto/pipeline/drivelaw/data/escape_data/danger_8hz/FR/FR_ARCF004_20221029151346.mp4"
    sample_video_path = "/workspace/vjepa/videos/-WH-lxmGJVY_000005_000015.mp4"

    encoder, predictor = torch.hub.load('/workspace/vjepa', 'vjepa2_1_vit_giant_384', source='local')
    encoder.cuda().eval()

    hf_transform = torch.hub.load('/workspace/vjepa', 'vjepa2_preprocessor', source='local')
    print('Successfully loaded VJEPA2 model and preprocessor from local PyTorch Hub.')
    # Inference on video
    out_patch_features_pt = forward_vjepa_video(
        encoder, hf_transform, sample_video_path
    )

    print(
        f"""
        Inference results on video:
        PyTorch output shape:     {out_patch_features_pt.shape}
        """
    )


def load_and_transform(video_path, pt_transform, num_frames=80):
    # 读取并统一采样为 num_frames（或尽量接近）
    vr = VideoReader(video_path)
    total_frames = len(vr)
    if total_frames < num_frames:
        frame_idx = np.arange(0, total_frames, max(1, total_frames // num_frames))
    else:
        frame_idx = np.linspace(0, total_frames - 1, num_frames, dtype=int)
    video = vr.get_batch(frame_idx).asnumpy()              # T x H x W x C
    video = torch.from_numpy(video).permute(0, 3, 1, 2)   # T x C x H x W
    # pt_transform 返回 list（可能多视角），取第一个视图（或根据需要修改）
    tensor = pt_transform(video)[0]                       # C x T x H x W
    return tensor

def forward_vjepa_multiview(encoder, pt_transform, video_paths, num_frames=80):
    """
    video_paths: list of paths, e.g. [FR, LF, RF]
    返回 encoder 输出（每个视角一个条目）
    """
    with torch.inference_mode():
        views = []
        for p in video_paths:
            t = load_and_transform(p, pt_transform, num_frames=num_frames)
            views.append(t)
        # 堆叠为 batch: (V, C, T, H, W)
        views = torch.stack(views, dim=0).cuda()
        encoder = encoder.cuda()
        encoder.eval()
        out = encoder(views)  # encoder 接受 B x C x T x H x W
    return out

def run_multi_sample_inference():
    FR_video_path = "/mnt/data2/tzx/workspace/auto/pipeline/drivelaw/data/escape_data/danger_8hz/FR/FR_ARCF004_20221029151346.mp4"
    LF_video_path = "/mnt/data2/tzx/workspace/auto/pipeline/drivelaw/data/escape_data/danger_8hz/LF/LF_ARCF004_20221029151346.mp4"
    RF_video_path = "/mnt/data2/tzx/workspace/auto/pipeline/drivelaw/data/escape_data/danger_8hz/RF/RF_ARCF004_20221029151346.mp4"

    # 从本地 hub 加载（返回 encoder, predictor）
    encoder, predictor = torch.hub.load('/workspace/vjepa', 'vjepa2_1_vit_base_384', source='local')
    print('Successfully loaded encoder and predictor from local hub.')

    # 预处理：用 hub 提供的 preprocessor 获取 crop size，然后构建 PT transform
    hf_transform = torch.hub.load('/workspace/vjepa', 'vjepa2_preprocessor', source='local')

    # 三视角一起推理
    video_paths = [LF_video_path, FR_video_path, RF_video_path]
    out = forward_vjepa_multiview(encoder, hf_transform, video_paths, num_frames=80)

    print(f"Encoder output for {len(video_paths)} views: {out.shape}")


if __name__ == "__main__":
    # Run with: `python -m notebooks.vjepa2_demo`
    run_single_sample_inference()