mirror of
https://github.com/ciromattia/kcc
synced 2025-12-13 01:36:27 +00:00
e7ee8bed9d
This helper file contains the methods necessary to perform a fourier transform on the picture, to remove frequencies responsible for rainbow artifacts on Kaleido screens, and performe the reverse fourier transform
131 lines
4.5 KiB
Python
131 lines
4.5 KiB
Python
import numpy as np
|
|
from PIL import Image
|
|
|
|
def fourier_transform_image(img):
|
|
"""
|
|
Performs a 2D Fourier transform on a PIL image.
|
|
|
|
Args:
|
|
img: PIL Image (can be color or grayscale)
|
|
|
|
Returns:
|
|
fft_result: Complex result of the 2D FFT
|
|
"""
|
|
# Convert PIL image to NumPy array
|
|
img_array = np.array(img)
|
|
|
|
# Perform 2D Fourier transform
|
|
fft_result = np.fft.fft2(img_array)
|
|
|
|
return fft_result
|
|
|
|
def attenuate_diagonal_frequencies(fft_spectrum, freq_threshold=0.3, target_angle=135,
|
|
angle_tolerance=15, attenuation_factor=0.1):
|
|
"""
|
|
Attenuates specific frequencies in the Fourier domain (optimized version).
|
|
|
|
Args:
|
|
fft_spectrum: Result of 2D Fourier transform (non-centered)
|
|
freq_threshold: Frequency threshold in cycles/pixel (default: 0.3, theoretical max: 0.5)
|
|
target_angle: Target angle in degrees (default: 135)
|
|
angle_tolerance: Angular tolerance in degrees (default: 15)
|
|
attenuation_factor: Attenuation factor (0.1 = 90% attenuation)
|
|
|
|
Returns:
|
|
np.ndarray: Modified FFT with applied attenuation (same format as input)
|
|
"""
|
|
|
|
# Get dimensions
|
|
height, width = fft_spectrum.shape
|
|
|
|
# Create frequency grids in an optimized way
|
|
freq_y = np.fft.fftfreq(height, d=1.0)
|
|
freq_x = np.fft.fftfreq(width, d=1.0)
|
|
|
|
# Use broadcasting to create grids without meshgrid (more efficient)
|
|
freq_y_grid = freq_y.reshape(-1, 1) # Column
|
|
freq_x_grid = freq_x.reshape(1, -1) # Row
|
|
|
|
# Calculate squared radial frequencies (avoid sqrt)
|
|
freq_radial_sq = freq_x_grid**2 + freq_y_grid**2
|
|
freq_threshold_sq = freq_threshold**2
|
|
|
|
# Frequency condition
|
|
freq_condition = freq_radial_sq >= freq_threshold_sq
|
|
|
|
# Early exit if no frequency satisfies the condition
|
|
if not np.any(freq_condition):
|
|
return fft_spectrum
|
|
|
|
# Calculate angles only where necessary
|
|
# Use atan2 directly with broadcasting
|
|
angles_rad = np.arctan2(freq_y_grid, freq_x_grid)
|
|
|
|
# Convert to degrees and normalize in a single operation
|
|
angles_deg = np.rad2deg(angles_rad) % 360
|
|
|
|
# Optimize angular condition
|
|
# Calculate both target angles at once
|
|
target_angle_2 = (target_angle + 180) % 360
|
|
|
|
# Create angular conditions in a vectorized way
|
|
angle_condition = np.zeros_like(angles_deg, dtype=bool)
|
|
|
|
# Process both angles simultaneously
|
|
for angle in [target_angle, target_angle_2]:
|
|
min_angle = (angle - angle_tolerance) % 360
|
|
max_angle = (angle + angle_tolerance) % 360
|
|
|
|
if min_angle > max_angle: # Interval crosses 0°
|
|
angle_condition |= (angles_deg >= min_angle) | (angles_deg <= max_angle)
|
|
else: # Normal interval
|
|
angle_condition |= (angles_deg >= min_angle) & (angles_deg <= max_angle)
|
|
|
|
# Combine conditions
|
|
combined_condition = freq_condition & angle_condition
|
|
|
|
# Apply attenuation directly (avoid creating a full mask)
|
|
if attenuation_factor == 0:
|
|
# Special case: complete suppression
|
|
fft_spectrum[combined_condition] = 0
|
|
return fft_spectrum
|
|
elif attenuation_factor == 1:
|
|
# Special case: no attenuation
|
|
return fft_spectrum
|
|
else:
|
|
# General case: partial attenuation
|
|
fft_spectrum[combined_condition] *= attenuation_factor
|
|
return fft_spectrum
|
|
|
|
def inverse_fourier_transform_image(fft_spectrum):
|
|
"""
|
|
Performs an inverse Fourier transform to reconstruct a PIL image.
|
|
|
|
Args:
|
|
fft_spectrum: Fourier transform result (complex array)
|
|
|
|
Returns:
|
|
PIL.Image: Reconstructed image
|
|
"""
|
|
# Perform inverse Fourier transform
|
|
img_reconstructed = np.fft.ifft2(fft_spectrum)
|
|
|
|
# Take real part (eliminate imaginary artifacts due to numerical errors)
|
|
img_reconstructed = np.real(img_reconstructed)
|
|
|
|
# Normalize values between 0 and 255
|
|
img_reconstructed = np.clip(img_reconstructed, 0, 255)
|
|
img_reconstructed = img_reconstructed.astype(np.uint8)
|
|
|
|
# Convert to PIL image
|
|
pil_image = Image.fromarray(img_reconstructed, mode='L')
|
|
|
|
return pil_image
|
|
|
|
import numpy as np
|
|
|
|
def erase_rainbow_artifacts(img):
|
|
fft_spectrum = fourier_transform_image(img)
|
|
clean_spectrum = attenuate_diagonal_frequencies(fft_spectrum)
|
|
clean_image = inverse_fourier_transform_image(clean_spectrum)
|
|
return clean_image |