chromatic aberration with options (for colab or not)

chromatic_aberration.py

import cv2
import numpy as np
try:
  from google.colab.patches import cv2_imshow
except:
  pass

def chromatic_aberration(image_path, save_path, shift_params=None):
    # Load the image
    image = cv2.imread(image_path)

    # Split the image into RGB channels
    r, g, b = cv2.split(image)

    # Apply the shift to each channel based on the shift_params dictionary
    shifted_r = shift_channel(r, shift_params['red'])
    shifted_g = shift_channel(g, shift_params['green'])
    shifted_b = shift_channel(b, shift_params['blue'])

    # Merge the shifted channels back into an image
    shifted_image = cv2.merge((shifted_r, shifted_g, shifted_b))

    # Save the resulting image
    cv2.imwrite(save_path, shifted_image)

    # Display the resulting image in the Colab cell
    cv2_imshow(shifted_image)

def shift_channel(channel, shift_param):
    # Extract the shift direction and strength from the shift_param dictionary
    direction_x = shift_param['direction_x']
    direction_y = shift_param['direction_y']
    strength = shift_param['strength']

    # Define the amount of shift for the channel
    shift_matrix = np.float32([[1, 0, direction_x * strength], [0, 1, direction_y * strength]])

    # Apply the shift to the channel
    shifted_channel = cv2.warpAffine(channel, shift_matrix, (channel.shape[1], channel.shape[0]))

    return shifted_channel

# Path to the input image
image_path = 'path_to_input_image.jpg'

# Path to save the resulting image
save_path = 'path_to_save_resulting_image.jpg'

# Default shift parameters for uniform shift in each channel
default_shift_params = {
    'red': {
        'direction_x': 1.0,  # X-direction shift for the red channel
        'direction_y': 0.0,  # Y-direction shift for the red channel
        'strength': 10       # Strength of the red channel shift
    },
    'green': {
        'direction_x': 0.0,  # X-direction shift for the green channel
        'direction_y': 1.0,  # Y-direction shift for the green channel
        'strength': 10       # Strength of the green channel shift
    },
    'blue': {
        'direction_x': 0.0,  # X-direction shift for the blue channel
        'direction_y': 0.0,  # Y-direction shift for the blue channel
        'strength': 10       # Strength of the blue channel shift
    }
}

# Apply chromatic aberration effect, save the resulting image, and display it
chromatic_aberration(image_path, save_path, default_shift_params)