Spaces:
Sleeping
Sleeping
File size: 3,141 Bytes
5ac1897 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 |
import imageio
import numpy as np
from tqdm import tqdm
from typing import Union, List
from pathlib import Path
from glob import glob
from .edit import flex_resize_img, flex_resize_video
def load_img_meta(
img_path : Union[str, Path],
):
''' Read the image meta from the given path without opening image. '''
assert Path(img_path).exists(), f'Image not found: {img_path}'
H, W = imageio.v3.improps(img_path).shape[:2]
meta = {'w': W, 'h': H}
return meta
def load_img(
img_path : Union[str, Path],
mode : str = 'RGB',
):
''' Read the image from the given path. '''
assert Path(img_path).exists(), f'Image not found: {img_path}'
img = imageio.v3.imread(img_path, plugin='pillow', mode=mode)
meta = {
'w': img.shape[1],
'h': img.shape[0],
}
return img, meta
def save_img(
img : np.ndarray,
output_path : Union[str, Path],
resize_ratio : Union[float, None] = None,
**kwargs,
):
''' Save the image. '''
assert img.ndim == 3, f'Invalid image shape: {img.shape}'
if resize_ratio is not None:
img = flex_resize_img(img, ratio=resize_ratio)
imageio.v3.imwrite(output_path, img, **kwargs)
def load_video(
video_path : Union[str, Path],
):
''' Read the video from the given path. '''
if isinstance(video_path, str):
video_path = Path(video_path)
assert video_path.exists(), f'Video not found: {video_path}'
if video_path.is_dir():
print(f'Found {video_path} is a directory. It will be regarded as a image folder.')
imgs_path = sorted(glob(str(video_path / '*')))
frames = []
for img_path in tqdm(imgs_path):
frames.append(imageio.imread(img_path))
fps = 30 # default fps
else:
print(f'Found {video_path} is a file. It will be regarded as a video file.')
reader = imageio.get_reader(video_path, format='FFMPEG')
frames = []
for frame in tqdm(reader, total=reader.count_frames()):
frames.append(frame)
fps = reader.get_meta_data()['fps']
frames = np.stack(frames, axis=0) # (L, H, W, 3)
meta = {
'fps': fps,
'w' : frames.shape[2],
'h' : frames.shape[1],
'L' : frames.shape[0],
}
return frames, meta
def save_video(
frames : Union[np.ndarray, List[np.ndarray]],
output_path : Union[str, Path],
fps : float = 30,
resize_ratio : Union[float, None] = None,
quality : Union[int, None] = None,
):
''' Save the frames as a video. '''
if isinstance(frames, List):
frames = np.stack(frames, axis=0)
assert frames.ndim == 4, f'Invalid frames shape: {frames.shape}'
if resize_ratio is not None:
frames = flex_resize_video(frames, ratio=resize_ratio)
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
writer = imageio.get_writer(output_path, fps=fps, quality=quality)
output_seq_name = str(output_path).split('/')[-1]
for frame in tqdm(frames, desc=f'Saving {output_seq_name}'):
writer.append_data(frame)
writer.close() |