Source code for src.file.ppm_handler

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

Description:
    This module is used to read and write PPM files.
    See: https://netpbm.sourceforge.net/doc/ppm.html
"""

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

import numpy as np


[docs] @dataclass class PPMHeader: width: int height: int depth: int
[docs] class PPMHandler: """Handler for reading and writing PPM image files."""
[docs] def read(self, path: Union[str, Path]) -> np.ndarray: """Reads a PPM file and returns the image data as a numpy array. :param path: Path to the PPM file :type path: Union[str, Path] :return: Image data as 3D numpy array (height, width, channels) :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, depth: int = 10) -> None: """Writes image data to a PPM file. :param path: Path to the output PPM file :type path: Union[str, Path] :param data: Image data as 3D numpy array (height, width, channels) :type data: np.ndarray :param depth: Bit depth for output, defaults to 10 :type depth: int :return: None :rtype: None """ with open(path, "wb") as file: self._write_header(file=file, data=data, depth=depth) self._write_data(file=file, data=data)
def _read_header(self, file: BufferedReader) -> PPMHeader: """Reads and parses the PPM header from the file. :param file: Opened binary file in read mode :type file: BufferedReader :return: Parsed PPM header :rtype: PPMHeader """ header = file.readline().split() if len(header) != 4: raise ValueError("Invalid PPM Header") header_id = header[0].decode("utf-8") width = int(header[1]) height = int(header[2]) maxval = int(header[3]) depth = int(math.ceil(math.log2(maxval + 1))) if header_id != "P6": raise ValueError(f'Invalid header id "{header_id}"') if depth <= 0 or depth > 16: 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 PPMHeader(width, height, depth) def _read_data(self, file: BufferedReader, header: PPMHeader) -> 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: PPM header with dimensions and format :type header: PPMHeader :return: Image data as 3D numpy array :rtype: np.ndarray """ dtype = "u1" if header.depth > 8: dtype = "u2" image_array = np.frombuffer(file.read(), dtype) shape = (header.height, header.width, 3) image_array = image_array.reshape(shape) return image_array def _write_header(self, file: BufferedWriter, data: np.ndarray, depth: int = 10): height, width, channels = data.shape assert channels == 3 maxval = (2**depth) - 1 # The maximum color value (Maxval), again in ASCII decimal. Must be less than 65536 and more than zero. assert maxval < 65536 and maxval > 0 file.write(f"P6 {width} {height} {maxval}\n".encode("ascii")) def _write_data(self, file: BufferedWriter, data: np.ndarray) -> None: """Writes raw image data to the file (raster of Height rows, Width pixels, RGB triplets). :param file: Opened binary file in write mode :type file: BufferedWriter :param data: Image data to write (height, width, 3) :type data: np.ndarray :return: None :rtype: None """ if data.dtype.byteorder == "<" or ( sys.byteorder == "little" and data.dtype.byteorder == "=" ): data.byteswap(inplace=True) bytes_data = data.tobytes() file.write(bytes_data)