51 lines
1.6 KiB
Python
51 lines
1.6 KiB
Python
import os
|
|
from typing import Iterable
|
|
|
|
import numpy as np
|
|
import torch
|
|
from PIL import Image
|
|
|
|
|
|
def _tensor_to_uint8(image: torch.Tensor) -> np.ndarray:
|
|
if image.is_floating_point():
|
|
image = image.clamp(0, 1).mul(255)
|
|
image = image.to(dtype=torch.uint8).cpu()
|
|
array = image.numpy()
|
|
if array.ndim == 3 and array.shape[2] >= 3:
|
|
array = array[:, :, :3]
|
|
return array
|
|
|
|
|
|
def save_frames_to_dir(frames: torch.Tensor, output_dir: str, name_pattern: str) -> None:
|
|
os.makedirs(output_dir, exist_ok=True)
|
|
if not isinstance(frames, torch.Tensor) or frames.dim() != 4:
|
|
raise ValueError("frames must be a 4D torch tensor [N, H, W, C].")
|
|
|
|
for index, frame in enumerate(frames, start=1):
|
|
filename = name_pattern % index
|
|
path = os.path.join(output_dir, filename)
|
|
array = _tensor_to_uint8(frame)
|
|
image = Image.fromarray(array)
|
|
image.save(path)
|
|
|
|
|
|
def load_frames_from_dir(frames_dir: str, extension: str = ".png") -> torch.Tensor:
|
|
if not os.path.isdir(frames_dir):
|
|
raise ValueError(f"Frames directory does not exist: {frames_dir}")
|
|
|
|
files = sorted(
|
|
f for f in os.listdir(frames_dir) if f.lower().endswith(extension)
|
|
)
|
|
if not files:
|
|
raise ValueError("No frames were generated by FFmpeg.")
|
|
|
|
images = []
|
|
for file in files:
|
|
path = os.path.join(frames_dir, file)
|
|
with Image.open(path) as image:
|
|
image = image.convert("RGB")
|
|
array = np.asarray(image).astype(np.float32) / 255.0
|
|
images.append(torch.from_numpy(array))
|
|
|
|
return torch.stack(images, dim=0)
|