使用compressai库做视频压缩报错，求视频输入网络及结果获取方案

问题背景 #

使用CompressAI库的ssf2020预训练模型测试视频压缩效果时，运行代码出现维度不匹配的RuntimeError，需要修正视频输入方式以正确获取压缩后的视频结果。

用户原代码 #

# pip install compressai
# pip install av

import math
import cv2
import io
import torch
from torchvision import transforms
import numpy as np
from pytorch_msssim import ms_ssim
from PIL import Image
from ipywidgets import interact, widgets
import matplotlib.pyplot as plt
from compressai.zoo import ssf2020
device = 'cuda' if torch.cuda.is_available() else 'cpu'
# 预训练模型
net = ssf2020(quality=4, pretrained=True).eval()
print(net)


import os
import av

def get_tensor_from_video(video_path, is_multi_thread_decode=True):
    """
    video_path: 视频路径
    is_multi_thread_decode: 是否启用多线程解码
    return: pytorch张量
    """
    if not os.access(video_path, os.F_OK):
        print('测试文件不存在')
        return
    container = av.open(video_path)
    if is_multi_thread_decode:
        container.streams.video[0].thread_type = "AUTO"
    container.seek(0, any_frame=False, backward=True, stream=container.streams.video[0])
    frames = []
    for frame in container.decode(video=0):
        frames.append(frame)
    container.close()

    result_frames = None

    result_frames = [frame.to_rgb().to_ndarray() for frame in frames]

    result_frames = torch.Tensor(np.stack(result_frames))

    # 输出张量维度为[视频帧数, 宽度, 高度, 通道数]
    return result_frames
cap = get_tensor_from_video('/content/drive/MyDrive/ColabNotebooks3/test.avi') 
cap = cap.permute(0, 3, 1, 2)  # 调整维度顺序，将cap转换为N,C,H,W格式
print(len(cap))    
print(cap.size())  

# 我猜测模型输入应为列表，每个元素为tensor格式，因此编写了以下代码
lists = []

for i in range(len(cap)):
    lists.append(cap[i])

with torch.no_grad():
    x_hat, likelihoods = net.forward(lists)

报错信息 #

RuntimeError：张量a的尺寸(5)必须与张量b的尺寸(192)在非单例维度0上匹配

问题原因 #

1. 输入格式错误：ssf2020是图像压缩模型，其forward方法接收的是批量张量（维度为[N, C, H, W]），而非帧列表。直接传入列表会导致模型内部张量维度匹配失败。
2. 数据类型与设备不匹配：原视频帧转成的张量是uint8类型，而模型期望输入为float类型；且模型未移至对应设备（cuda/cpu），导致参数与输入张量设备不一致。

修正后的代码 #

# pip install compressai
# pip install av

import torch
import numpy as np
import cv2
import av
from compressai.zoo import ssf2020

device = 'cuda' if torch.cuda.is_available() else 'cpu'
# 将预训练模型移至对应设备
net = ssf2020(quality=4, pretrained=True).eval().to(device)

def get_tensor_from_video(video_path, is_multi_thread_decode=True):
    if not os.access(video_path, os.F_OK):
        print('测试文件不存在')
        return
    container = av.open(video_path)
    if is_multi_thread_decode:
        container.streams.video[0].thread_type = "AUTO"
    container.seek(0, any_frame=False, backward=True, stream=container.streams.video[0])
    frames = []
    for frame in container.decode(video=0):
        frames.append(frame)
    container.close()

    # 转换为float32张量，归一化到[0,1]，调整维度为N,C,H,W并移至设备
    result_frames = [frame.to_rgb().to_ndarray() for frame in frames]
    result_frames = torch.tensor(np.stack(result_frames)).permute(0,3,1,2).float() / 255.0
    result_frames = result_frames.to(device)
    return result_frames

# 加载视频张量
cap = get_tensor_from_video('/content/drive/MyDrive/ColabNotebooks3/test.avi') 
print(f"视频帧数: {cap.size(0)}")    
print(f"张量尺寸: {cap.size()}")  

# 批量推理，直接传入N,C,H,W格式的张量
with torch.no_grad():
    x_hat, likelihoods = net(cap)  # 直接调用模型即可，无需显式调用forward

# 将输出的归一化张量转换回uint8格式，用于保存视频
x_hat = (x_hat.clamp(0, 1) * 255).byte().permute(0,2,3,1).cpu().numpy()

# 将压缩后的帧合成视频
output_path = '/content/drive/MyDrive/ColabNotebooks3/compressed_test.avi'
height, width = x_hat.shape[1], x_hat.shape[2]
fourcc = cv2.VideoWriter_fourcc(*'XVID')
out = cv2.VideoWriter(output_path, fourcc, 30.0, (width, height))

for frame in x_hat:
    # 转换为cv2默认的BGR格式
    bgr_frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
    out.write(bgr_frame)

out.release()
print(f"压缩视频已保存至: {output_path}")

关键修正点 #

• 模型初始化时添加.to(device)，确保模型与输入张量在同一设备运行。
• 视频张量转换为float32并归一化到[0,1]，符合模型输入要求。
• 直接将批量张量传入模型，无需拆分为列表。
• 将输出张量转换回uint8格式，并合成视频保存。

内容的提问来源于stack exchange，提问作者JiaXiang He