| 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 |
| |
|
|
