Reconstruct RAW images.

recons_raw.py

"""
Reconstruct RAW images.
"""

from pathlib import Path

import tyro
import cv2
import imageio
import numpy as np


def apply_white_balance(
    image: np.ndarray, gains: tuple[float, float, float]
) -> np.ndarray:
    """
    Apply white balance to an image.

    Args:
        image (np.ndarray): the input image
        gains (tuple[float, float, float]): the white balance gains

    Returns:
        np.ndarray: the balanced image
    """
    balanced_image = image.astype(np.float32)
    balanced_image = balanced_image * np.array(gains)[None, None, :]
    balanced_image = np.clip(balanced_image, 0, 4095)
    balanced_image = balanced_image.astype(np.uint16)
    return balanced_image


def process_raw_image(
    input_path: Path,
    output_path: Path,
    gains: tuple[float, float, float] = (1.5, 1.0, 2.0),
):
    """
    Process a RAW image.

    Args:
        input_path (Path): the input path of the RAW image
        output_path (Path): the output path of the processed image
    """
    with open(input_path, "rb") as f:
        raw_data = np.fromfile(f, dtype=np.uint16)
    width, height = 3840, 2344
    raw_image = raw_data.reshape((height, width))
    rgb_image = cv2.cvtColor(raw_image, cv2.COLOR_BayerRG2RGB_EA)  # pylint: disable=no-member
    balanced_rgb_image = apply_white_balance(rgb_image, gains)
    image_max_value = 4095
    finals = (balanced_rgb_image / image_max_value * 255).astype(np.uint8)
    imageio.imwrite(output_path, finals)


def main():
    """
    Main entry point.
    """
    tyro.cli(process_raw_image)


if __name__ == "__main__":
    main()

wow