NVIDIA/DALI

fn.readers.video--extract frames details

Open

#4,797 创建于 2023年4月20日

在 GitHub 查看
 (5 评论) (0 反应) (1 负责人)C++ (670 fork)auto 404
Videohelp wanted

仓库指标

Star
 (5,722 star)
PR 合并指标
 (PR 指标待抓取)

描述

import os.path
import numpy as np
import shutil
from PIL import Image
from nvidia.dali import pipeline_def
import nvidia.dali.fn as fn
import nvidia.dali.types as types

def get_opts(video_path):
    meta = get_video_meta_info(video_path)
    width = meta['width']
    height = meta['height']
    input_fps = meta['fps']
    duration = meta['duration']
    nb_frames = meta['nb_frames']

    if nb_frames == 0:
        print('error: the num of video frames is 0!')

    scale = 1
    video_uniform_len = min(int(duration), 60)
    sequence_length = 32
    video_read_interval = int((duration / video_uniform_len) / scale) * scale  # uniform extract frames
    if video_read_interval == 1:
        sequence_step = int(input_fps)
    else:
        sequence_step = nb_frames // 60
    opts = {
        'spatial_sequence_length': 1,
        'motion_sequence_length': sequence_length,
        'stride': 1,
        'step': sequence_step,
        'n_iter': video_uniform_len,
        'batch_size': 1,
        'num_threads': 4,
        'height': height,
        'width': width
    }
    return opts

def save_images(frames, seq_len, directory):
    for j in range(seq_len):
        im = Image.fromarray(frames[j])
        im.save(os.path.join(directory, str(j)) + '.png')

def extract_frames(filenames):
    data_dir = "/workspace/code/FAST-VQA-and-FasterVQA/vqa_dataset/output"

    opts = get_opts(video_filename)
    @pipeline_def
    def video_pipe(filenames):
        video = fn.readers.video(device="gpu", filenames=filenames, sequence_length=opts['motion_sequence_length'],
                                 stride=opts['stride'], step=opts['step'],
                                 skip_vfr_check=True)
        return video

    pipe = video_pipe(filenames=video_filename, batch_size=opts['batch_size'], num_threads=opts['num_threads'],
                      device_id=0)
    pipe.build()

    for i in range(opts['n_iter']):
        pipe_out = pipe.run()
        frames = np.array(pipe_out[0][0].as_cpu())
        print(frames.shape)
        label_dir = os.path.join(data_dir, str(i))
        os.makedirs(label_dir)
        save_images(frames, opts['motion_sequence_length'], label_dir)

video_filename = "test.mp4"
extract_frames(video_filename)

Uploading 000.mp4…

Hello, when I use this API to extract frames, if I input a video, the length of the video is less than 60s, suppose 15s, and FPS=24, I plan to set sequence_length=32, and I find that the first sequence is directly copied in the last sequence. I would like to ask if we can not copy the first sequence, but directly copy the 24 frames of the last sequence into 32 frames to form a sequence of (15, 32, 3, h, w) 2681681982272_ pic @jantonguirao

贡献者指南