<reponame>shinying/METER<filename>meter/datasets/util.py
import glob
import os.path as op
import random

# import av
# import decord
import numpy as np
import torch
# decord.bridge.set_bridge('torch')


# Borrow from https://github.com/m-bain/frozen-in-time/blob/main/base/base_dataset.py
def sample_frames(num_frames, vlen, sample='rand', crop=1., fix_start=-1):
    acc_samples = min(num_frames, vlen)
    rem = (1. - crop) / 2
    start, stop = int(vlen * rem), int(vlen * (1 - rem))
    intv = np.linspace(start=start, stop=stop, num=acc_samples+1).astype(int)
    if sample == 'rand':
        frame_idxs = [random.randrange(intv[i], intv[i+1]) for i in range(len(intv)-1)]
    elif fix_start >= 0:
        fix_start = int(fix_start)
        frame_idxs = [intv[i]+fix_start for i in range(len(intv)-1)]
    elif sample == 'uniform':
        frame_idxs = [(intv[i]+intv[i+1]-1) // 2 for i in range(len(intv)-1)]
    else:
        raise NotImplementedError
    return frame_idxs


def read_frames(video_path, num_frames, sample='rand', crop=1., fix_start=-1):
    frames = glob.glob(op.join(video_path, '*.png'))
    if not len(frames):
        raise FileNotFoundError("no such videos")

    frames.sort(key=lambda n: int(op.basename(n)[:-4]))
    while len(frames) < num_frames / crop: # duplicate frames
        frames = [f for frame in frames for f in (frame, frame)]
    frame_idxs = sample_frames(num_frames, len(frames), sample, crop, fix_start)
    return [frames[i] for i in frame_idxs]

"""
def read_frames_decord(video_path, num_frames, sample='rand', crop=1., fix_start=-1):
    video_reader = decord.VideoReader(video_path, num_threads=8)
    vlen = len(video_reader)
    frame_idxs = sample_frames(num_frames, vlen, sample, fix_start, crop)
    frames = video_reader.get_batch(frame_idxs)
    frames = frames.float() / 255
    frames = frames.permute(0, 3, 1, 2)
    return frames


def read_frames_av(video_path, num_frames, sample='rand', fix_start=-1, crop=1):
    reader = av.open(video_path)
    try:
        frames = []
        frames = [torch.from_numpy(f.to_rgb().to_ndarray()) for f in reader.decode(video=0)]
    except (RuntimeError, ZeroDivisionError) as exception:
        print('{}: WEBM reader cannot open {}. Empty '
              'list returned.'.format(type(exception).__name__, video_path))
    vlen = len(frames)
    frame_idxs = sample_frames(num_frames, vlen, sample, fix_start, crop)
    frames = torch.stack([frames[idx] for idx in frame_idxs]).float() / 255
    frames = frames.permute(0, 3, 1, 2)
    return frames
"""

import glob
import os.path as op
import random

# import av
# import decord
import numpy as np
import torch
# decord.bridge.set_bridge('torch')


# Borrow from https://github.com/m-bain/frozen-in-time/blob/main/base/base_dataset.py
def sample_frames(num_frames, vlen, sample='rand', crop=1., fix_start=-1):
    acc_samples = min(num_frames, vlen)
    rem = (1. - crop) / 2
    start, stop = int(vlen * rem), int(vlen * (1 - rem))
    intv = np.linspace(start=start, stop=stop, num=acc_samples+1).astype(int)
    if sample == 'rand':
        frame_idxs = [random.randrange(intv[i], intv[i+1]) for i in range(len(intv)-1)]
    elif fix_start >= 0:
        fix_start = int(fix_start)
        frame_idxs = [intv[i]+fix_start for i in range(len(intv)-1)]
    elif sample == 'uniform':
        frame_idxs = [(intv[i]+intv[i+1]-1) // 2 for i in range(len(intv)-1)]
    else:
        raise NotImplementedError
    return frame_idxs


def read_frames(video_path, num_frames, sample='rand', crop=1., fix_start=-1):
    frames = glob.glob(op.join(video_path, '*.png'))
    if not len(frames):
        raise FileNotFoundError("no such videos")

    frames.sort(key=lambda n: int(op.basename(n)[:-4]))
    while len(frames) < num_frames / crop: # duplicate frames
        frames = [f for frame in frames for f in (frame, frame)]
    frame_idxs = sample_frames(num_frames, len(frames), sample, crop, fix_start)
    return [frames[i] for i in frame_idxs]

"""
def read_frames_decord(video_path, num_frames, sample='rand', crop=1., fix_start=-1):
    video_reader = decord.VideoReader(video_path, num_threads=8)
    vlen = len(video_reader)
    frame_idxs = sample_frames(num_frames, vlen, sample, fix_start, crop)
    frames = video_reader.get_batch(frame_idxs)
    frames = frames.float() / 255
    frames = frames.permute(0, 3, 1, 2)
    return frames


def read_frames_av(video_path, num_frames, sample='rand', fix_start=-1, crop=1):
    reader = av.open(video_path)
    try:
        frames = []
        frames = [torch.from_numpy(f.to_rgb().to_ndarray()) for f in reader.decode(video=0)]
    except (RuntimeError, ZeroDivisionError) as exception:
        print('{}: WEBM reader cannot open {}. Empty '
              'list returned.'.format(type(exception).__name__, video_path))
    vlen = len(frames)
    frame_idxs = sample_frames(num_frames, vlen, sample, fix_start, crop)
    frames = torch.stack([frames[idx] for idx in frame_idxs]).float() / 255
    frames = frames.permute(0, 3, 1, 2)
    return frames
"""
# import av # import decord # decord.bridge.set_bridge('torch') # Borrow from https://github.com/m-bain/frozen-in-time/blob/main/base/base_dataset.py # duplicate frames def read_frames_decord(video_path, num_frames, sample='rand', crop=1., fix_start=-1):     video_reader = decord.VideoReader(video_path, num_threads=8)     vlen = len(video_reader)     frame_idxs = sample_frames(num_frames, vlen, sample, fix_start, crop)     frames = video_reader.get_batch(frame_idxs)     frames = frames.float() / 255     frames = frames.permute(0, 3, 1, 2)     return frames   def read_frames_av(video_path, num_frames, sample='rand', fix_start=-1, crop=1):     reader = av.open(video_path)     try:         frames = []         frames = [torch.from_numpy(f.to_rgb().to_ndarray()) for f in reader.decode(video=0)]     except (RuntimeError, ZeroDivisionError) as exception:         print('{}: WEBM reader cannot open {}. Empty '               'list returned.'.format(type(exception).__name__, video_path))     vlen = len(frames)     frame_idxs = sample_frames(num_frames, vlen, sample, fix_start, crop)     frames = torch.stack([frames[idx] for idx in frame_idxs]).float() / 255     frames = frames.permute(0, 3, 1, 2)     return frames