ISP_pipeline.py

import numpy as np
import cv2
import imageio

def simplest_color_balance(img, percent):
    out_channels = []
    channels = cv2.split(img)
    for channel in channels:
        total_pixels = img.shape[0] * img.shape[1]
        low_val, high_val = np.percentile(channel, [percent, 100 - percent])
        channel = np.clip(channel, low_val, high_val)
        channel = (channel - low_val) * (255 / (high_val - low_val))
        out_channels.append(channel)
    return cv2.merge(out_channels)

def gray_world(img):
    img_float = img.astype(np.float32)
    average_color_per_row = np.average(img_float, axis=0)
    average_color = np.average(average_color_per_row, axis=0)
    average_color = np.uint8(average_color)
    gray_avg = np.average(average_color)
    scale = gray_avg / average_color
    scale = scale.astype(np.float32).reshape(1, 1, 3)  # Convert scale to float32 and reshape to match the image dimensions
    balanced_img = np.multiply(img_float, scale)
    return np.clip(balanced_img, 0, 255).astype(np.uint8)

def tone_mapping(image, gamma=2.2):
    return np.power(image, 1.0 / gamma)

def load_image(self, image_index):
    if self.image_names[image_index].endswith('.jpg') or self.image_names[image_index].endswith('.png'):
        raw_image = imageio.imread(self.image_names[image_index])

        if raw_image.ndim == 2:
            # Raw image processing
            rgb_image = cv2.cvtColor(raw_image, cv2.COLOR_BayerGR2BGR)
            rgb_image_uint8 = (rgb_image * 255).astype(np.uint8)
            # Denoising
            denoised_image = cv2.fastNlMeansDenoisingColored(rgb_image_uint8)

            # Automatic white balancing using simplest color balance
            wb_image = simplest_color_balance(denoised_image, 1)

            # Automatic color correction using gray world assumption
            color_corrected_image = gray_world(wb_image)

            normalized_rgb_image = color_corrected_image.astype(np.float32) / 2**12
        else:
            # Non-raw image processing (JPEG or RGB PNG)
            normalized_rgb_image = raw_image.astype(np.float32) / 255.0

        # Clip the values of the normalized image to the range [0, 1]
        clipped_normalized_rgb_image = np.clip(normalized_rgb_image, 0, 1)

        # Apply tone mapping
        tone_mapped_image = tone_mapping(clipped_normalized_rgb_image)

        # Convert the image to the range [0, 255] and dtype uint8
        img = (tone_mapped_image * 255).astype(np.uint8)
    return img

if __name__ == '__main__':
    load_image()