Phenaki生成视频写入MP4报错：期望3维数组却得到4维

我尝试将Phenaki的make_video输出写入MP4文件，使用lucidrains/phenaki-pytorch实现，代码如下：

phenaki = Phenaki(
    cvivit = cvivit,
    maskgit = maskgit
)


entire_video, scenes = make_video(phenaki, texts = [
    'blah blah',
], num_frames = (17, 14, 14), prime_lengths = (5, 5))

entire_video.shape # (1, 3, 17 + 14 + 14 = 45, 256, 256)
torchvision.io.write_video(filename= "test.mp4", video_array= entire_video, fps=24)

执行时出现如下错误：

File "/.../GitHub/phenaki-pytorch/run.py", line 49, in <module>
    torchvision.io.write_video(filename= "test.mp4", video_array= entire_video, fps=24)
  File "/opt/homebrew/lib/python3.10/site-packages/torchvision/io/video.py", line 132, in write_video
    frame = av.VideoFrame.from_ndarray(img, format="rgb24")
  File "av/video/frame.pyx", line 408, in av.video.frame.VideoFrame.from_ndarray
  File "av/utils.pyx", line 72, in av.utils.check_ndarray
ValueError: Expected numpy array with ndim `3` but got `4`

请问我哪里操作有误？为什么av.VideoFrame.from_ndarray期望3维numpy数组？

错误原因 #

• 张量维度不匹配：entire_video的形状是(1, 3, 45, 256, 256)，对应(批次大小, 通道数, 帧数, 高度, 宽度)，但torchvision.io.write_video要求输入维度为(帧数, 高度, 宽度, 通道数)或(批次大小, 帧数, 高度, 宽度, 通道数)，你的张量不仅多了不必要的批次维度，通道数的位置也不符合要求。
• 单帧处理要求：av.VideoFrame.from_ndarray是用来处理单帧图像的，单帧图像的numpy数组维度必须是(高度, 宽度, 通道数)（3维），而你的代码传入的是带批次维度的4维数据，因此触发报错。

解决方法 #

调整张量维度顺序并移除多余的批次维度，代码修改如下：

import torch

# 移除批次维度，同时调整维度顺序为 (帧数, 高度, 宽度, 通道数)
video_tensor = entire_video.squeeze(0).permute(1, 2, 3, 0)
# 模型输出通常为0-1的浮点数，需转换为视频要求的0-255整数格式
video_tensor = (video_tensor * 255).byte()

torchvision.io.write_video(filename="test.mp4", video_array=video_tensor, fps=24)

关键步骤说明 #

• squeeze(0)：去掉第一个批次维度（当前只有1个视频，批次维度无意义）。
• permute(1, 2, 3, 0)：将原维度(3, 45, 256, 256)转换为(45, 256, 256, 3)，完全符合write_video的输入格式要求。
• (video_tensor * 255).byte()：把模型输出的浮点值映射到视频标准的0-255整数范围，并转换为uint8数据类型。

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