Source code for enc.io.io

from enc.utils.codingstructure import FrameData
from enc.io.types import FRAME_DATA_TYPE, POSSIBLE_BITDEPTH
from enc.io.format.ppm import read_ppm, write_ppm
from enc.io.format.yuv import read_yuv, write_yuv
from enc.io.format.png import read_png, write_png

import os


[docs] def load_frame_data_from_file(file_path: str, idx_display_order: int) -> FrameData: """Load the idx_display_order-th frame from a .yuv file or .png file. For the latter, idx_display_order must be equal to 0 as there is only one frame in a png. Args: file_path (str): Absolute path of the file from which the frame is loaded. idx_display_order (int): Index of the frame in display order Returns: FrameData: The loaded frame, wrapped as a FrameData object. """ POSSIBLE_EXT = [".yuv", ".png", ".ppm"] assert file_path[-4:] in POSSIBLE_EXT, ( "The function load_frame_data_from_file() expects a file ending with " f"{POSSIBLE_EXT}. Found {file_path}" ) if file_path.endswith(".yuv"): # ! We only consider yuv420 and 444 planar bitdepth: POSSIBLE_BITDEPTH = 8 if "_8b" in file_path else 10 frame_data_type: FRAME_DATA_TYPE = "yuv420" if "420" in file_path else "yuv444" data = read_yuv(file_path, idx_display_order, frame_data_type, bitdepth) elif file_path.endswith(".png"): frame_data_type: FRAME_DATA_TYPE = "rgb" data, bitdepth = read_png(file_path) elif file_path.endswith(".ppm"): frame_data_type: FRAME_DATA_TYPE = "rgb" data, bitdepth = read_ppm(file_path) return FrameData(bitdepth, frame_data_type, data)
def save_frame_data_to_file(frame_data: FrameData, file_path: str) -> None: """Save the data of a FrameData into a PNG, PPM or YUV file. file_path extension must match the FrameData type e.g. PNG or PPM for RGB and YUV for YUV420 or YUV444 Args: frame_data (FrameData): The data to save file_path (str): Absolute path of the file from which the frame is stored. """ POSSIBLE_EXT = [".yuv", ".png", ".ppm"] cur_extension = os.path.splitext(file_path)[1] assert cur_extension in POSSIBLE_EXT, ( "The function save_frame_data_to_file() expects a file ending with " f"{POSSIBLE_EXT}. Found {file_path}" ) if cur_extension == ".png": assert frame_data.frame_data_type == "rgb", ( "The function save_frame_data_to_file() can only save a RGB data " f"into a PNG file. Found frame_data_type = {frame_data.frame_data_type}." ) assert frame_data.bitdepth == 8, ( "The function save_frame_data_to_file() can only write 8-bit data " f"into a PNG file. Found bitdepth = {frame_data.bitdepth}." ) write_png(frame_data.data, file_path) elif cur_extension == ".ppm": assert frame_data.frame_data_type == "rgb", ( "The function save_frame_data_to_file() can only save a RGB data " f"into a PPM file. Found frame_data_type = {frame_data.frame_data_type}." ) write_ppm(frame_data.data, frame_data.bitdepth, file_path, norm=True) elif cur_extension == ".yuv": assert frame_data.frame_data_type in ["yuv420", "yuv444"], ( "The function save_frame_data_to_file() can only save a YUV data " f"into a YUV file. Found frame_data_type = {frame_data.frame_data_type}." ) write_yuv( frame_data.data, frame_data.bitdepth, frame_data.frame_data_type, file_path, norm=True, )