"""
Author: Ismael Seidel (ismael.seidel@ufsc.br)
Affiliation: Embedded Computing Lab (ECL), Federal University of Santa Catarina (UFSC)
Description:
This module contains the implementation of YUV->PGX conversion.
"""
from pathlib import Path
from typing import Union
import yuvio
from lfc_toolkit.src.data_handlers.encoding_orders import (
ScanFunctionType,
get_serpentine_scan_list,
)
from lfc_toolkit.src.data_handlers.formatters import get_view_name_with_leading_zeros
from lfc_toolkit.src.data_handlers.lightfield import (
RAW_BT709_FR_PGX_LightField_Data,
RAWLightFieldData,
)
from lfc_toolkit.src.file.pgx_handler import PGXHandler
[docs]
def generate_pgx_from_yuv(
raw_light_field_data: RAWLightFieldData,
pgx_path: Union[str, Path],
scan_order: ScanFunctionType = get_serpentine_scan_list,
remove_after_using: bool = False,
):
"""
Get RAW PGX light field object from RAW YUV light field data.
:param raw_light_field_data: RAW YUV light field to be converted
:type raw_light_field_data: RAWLightFieldData
:param pgx_path: Path to save the PGX files
:type pgx_path: Union[str, Path]
:param scan_order: Function defining the scan order
:type scan_order: ScanFunctionType
:param remove_after_using: Whether to remove intermediate files
:type remove_after_using: bool
:return: Converted PGX light field data
:rtype: RAW_BT709_FR_PGX_LightField_Data
"""
assert raw_light_field_data.colour_space == "BT.709 Full Range"
assert raw_light_field_data.pix_fmt == "yuv444p10le"
lightfield = raw_light_field_data.copy()
output_raw_lf = RAW_BT709_FR_PGX_LightField_Data(
lightfield=lightfield,
pgx_path=pgx_path,
bpp_for_naming=raw_light_field_data.bpp_for_naming,
remove_after_using=remove_after_using,
)
hand = PGXHandler()
scan_list = scan_order(
raw_light_field_data.n_views_width,
raw_light_field_data.n_views_height,
0,
0,
1,
1,
) # always in this format 0,0,1,1 as it is expected by the JPLM
input_yuv_filename = raw_light_field_data.raw_path
reader = yuvio.get_reader(
input_yuv_filename,
raw_light_field_data.view_width,
raw_light_field_data.view_height,
raw_light_field_data.pix_fmt,
)
for (x, y), yuv_frame in zip(scan_list, reader):
view_filename = f"{get_view_name_with_leading_zeros(x, y)}.pgx"
for i, channel_data in enumerate([yuv_frame.y, yuv_frame.u, yuv_frame.v]):
path = output_raw_lf.raw_path / str(i) / view_filename
path.parent.mkdir(parents=True, exist_ok=True)
hand.write(path, channel_data)
return output_raw_lf
[docs]
def main():
import argparse
# 1. Create the ArgumentParser object with a description
parser = argparse.ArgumentParser(description="Convert a single YUV LF to PGX.")
# 2. Add a positional argument (required input, no dashes)
parser.add_argument(
"--name", "-n", required=True, type=str, help="The name of the light field."
)
parser.add_argument("--n-views-height", "-t", required=True, type=int, help="T")
parser.add_argument("--n-views-width", "-s", required=True, type=int, help="S")
parser.add_argument("--view-height", "-v", required=True, type=int, help="V")
parser.add_argument("--view-width", "-u", required=True, type=int, help="U")
parser.add_argument(
"--input_path", "-i", required=True, type=str, help="Input path"
)
parser.add_argument(
"--output_path", "-o", required=True, type=str, help="Output path"
)
# 3. Parse the arguments from the command line
args = parser.parse_args()
from lfc_toolkit.src.data_handlers.lightfield import LightField
lf = RAWLightFieldData(
lightfield=LightField(
name=args.name,
view_width=args.view_width,
view_height=args.view_height,
n_views_width=args.n_views_width,
n_views_height=args.n_views_height,
),
raw_path=args.input_path,
type="YUV",
pix_fmt="yuv444p10le",
)
generate_pgx_from_yuv(
raw_light_field_data=lf, pgx_path=args.output_path, remove_after_using=False
)
if __name__ == "__main__":
main()