|
import av |
|
import os |
|
import pims |
|
import numpy as np |
|
from torch.utils.data import Dataset |
|
from torchvision.transforms.functional import to_pil_image |
|
from PIL import Image |
|
|
|
|
|
class VideoReader(Dataset): |
|
def __init__(self, path, transform=None): |
|
self.video = pims.PyAVVideoReader(path) |
|
self.rate = self.video.frame_rate |
|
self.transform = transform |
|
|
|
@property |
|
def frame_rate(self): |
|
return self.rate |
|
|
|
def __len__(self): |
|
return len(self.video) |
|
|
|
def __getitem__(self, idx): |
|
frame = self.video[idx] |
|
frame = Image.fromarray(np.asarray(frame)) |
|
if self.transform is not None: |
|
frame = self.transform(frame) |
|
return frame |
|
|
|
|
|
class VideoWriter: |
|
def __init__(self, path, frame_rate, bit_rate=1000000): |
|
self.container = av.open(path, mode='w') |
|
self.stream = self.container.add_stream('h264', rate=f'{frame_rate:.4f}') |
|
self.stream.pix_fmt = 'yuv420p' |
|
self.stream.bit_rate = bit_rate |
|
|
|
def write(self, frames): |
|
|
|
self.stream.width = frames.size(3) |
|
self.stream.height = frames.size(2) |
|
if frames.size(1) == 1: |
|
frames = frames.repeat(1, 3, 1, 1) |
|
frames = frames.mul(255).byte().cpu().permute(0, 2, 3, 1).numpy() |
|
for t in range(frames.shape[0]): |
|
frame = frames[t] |
|
frame = av.VideoFrame.from_ndarray(frame, format='rgb24') |
|
self.container.mux(self.stream.encode(frame)) |
|
|
|
def close(self): |
|
self.container.mux(self.stream.encode()) |
|
self.container.close() |
|
|
|
|
|
class ImageSequenceReader(Dataset): |
|
def __init__(self, path, transform=None): |
|
self.path = path |
|
self.files = sorted(os.listdir(path)) |
|
self.transform = transform |
|
|
|
def __len__(self): |
|
return len(self.files) |
|
|
|
def __getitem__(self, idx): |
|
with Image.open(os.path.join(self.path, self.files[idx])) as img: |
|
img.load() |
|
if self.transform is not None: |
|
return self.transform(img) |
|
return img |
|
|
|
|
|
class ImageSequenceWriter: |
|
def __init__(self, path, extension='jpg'): |
|
self.path = path |
|
self.extension = extension |
|
self.counter = 0 |
|
os.makedirs(path, exist_ok=True) |
|
|
|
def write(self, frames): |
|
|
|
for t in range(frames.shape[0]): |
|
to_pil_image(frames[t]).save(os.path.join( |
|
self.path, str(self.counter).zfill(4) + '.' + self.extension)) |
|
self.counter += 1 |
|
|
|
def close(self): |
|
pass |
|
|
|
|
