Source code for src.file.pgx_handler

"""
Author: André Filipe da Silva Fernandes
Affiliation: Embedded Computing Lab (ECL), Federal University of Santa Catarina (UFSC)
Contributors:
    - Ismael Seidel (ismael.seidel@ufsc.br) 

Description:
    This module is used to read and write PGX files.
"""

import array
from dataclasses import dataclass
from io import BufferedReader, BufferedWriter
from pathlib import Path
from typing import Union

import numpy as np


[docs] @dataclass class PGXHeader: width: int height: int depth: int byteorder: int
[docs] class PGXHandler: """Handler for reading and writing PGX image files."""
[docs] def read(self, path: Union[str, Path]) -> np.ndarray: """Reads a PGX file and returns the image data as a numpy array. :param path: Path to the PGX file :type path: Union[str, Path] :return: Image data as 2D numpy array :rtype: np.ndarray """ with open(path, "rb") as file: header = self._read_header(file) image = self._read_data(file, header) return image
[docs] def write(self, path: Union[str, Path], data: np.ndarray) -> None: """Writes image data to a PGX file. :param path: Path to the output PGX file :type path: Union[str, Path] :param data: Image data as 2D numpy array :type data: np.ndarray :return: None :rtype: None """ with open(path, "wb") as file: self._write_header(file, data) self._write_data(file, data)
def _read_header(self, file: BufferedReader) -> PGXHeader: """Reads and parses the PGX header from the file. :param file: Opened binary file in read mode :type file: BufferedReader :return: Parsed PGX header :rtype: PGXHeader """ header = file.readline().split() if len(header) != 5: raise ValueError("Invalid PGX Header") header_id = header[0].decode("utf-8") endianess = header[1].decode("utf-8") signal_depth = header[2].decode("utf-8") width = int(header[3]) height = int(header[4]) signal = signal_depth[0] depth = int(signal_depth[1:]) if header_id != "PG": raise ValueError(f'Invalid header id "{header_id}"') if endianess == "ML": byteorder = "big" elif endianess == "LM": byteorder = "little" else: raise ValueError(f'Invalid endianess "{endianess}"') if signal not in ["+", "-"]: raise ValueError(f'Invalid signal "{signal}"') if depth <= 0: raise ValueError(f'Invalid depth "{depth}"') if width <= 0: raise ValueError(f'Invalid width "{width}"') if height <= 0: raise ValueError(f'Invalid height "{height}"') return PGXHeader(width, height, depth, byteorder) def _read_data(self, file: BufferedReader, header: PGXHeader) -> np.ndarray: """Reads raw image data from the file according to the header. :param file: Opened binary file in read mode :type file: BufferedReader :param header: PGX header with dimensions and format :type header: PGXHeader :return: Image data as 2D numpy array :rtype: np.ndarray """ dtype = "i2" if header.byteorder == "big": dtype = ">" + dtype image_array = np.frombuffer(file.read(), dtype) shape = (header.height, header.width) image_array = image_array.reshape(shape) return image_array def _write_header( self, file: BufferedWriter, data: np.ndarray, byteorder: str = "big" ) -> None: """Writes the PGX header to the file. :param file: Opened binary file in write mode :type file: BufferedWriter :param data: Image data for dimensions :type data: np.ndarray :param byteorder: Byte order ("big" or "little"), defaults to "big" :type byteorder: str :return: None :rtype: None """ signal = "+" depth = 10 height, width = data.shape file.write(b"PG ") if byteorder == "big": file.write(b"ML ") else: file.write(b"LM ") file.write(bytes(signal, "utf8")) file.write(bytes(f"{depth} {width} {height} \n", "utf8")) def _write_data( self, file: BufferedWriter, data: np.ndarray, byteorder: str = "big" ) -> None: """Writes raw image data to the file. :param file: Opened binary file in write mode :type file: BufferedWriter :param data: Image data to write :type data: np.ndarray :param byteorder: Byte order ("big" or "little"), defaults to "big" :type byteorder: str :return: None :rtype: None """ if byteorder == "big": file.write(data.byteswap().tobytes()) else: file.write(data.tobytes())