mirror of
https://github.com/ciromattia/kcc
synced 2025-12-13 01:36:27 +00:00
0d967084a0
It was a mistake to add the perpendicular angles in the previous commit: I had the rainbow effect removal process called 2 times when I did this, for testing purposes (One before downscale and one after downscale). The proper way to call the process is only after the downscale. And in this case it is not necessary to remove frequencies along the perpendicular angles. In the mean time, attenuation_factor=0.15 has proven to work well along a collection of testing images. It should be my latest commit for this feature
124 lines
4.6 KiB
Python
124 lines
4.6 KiB
Python
import numpy as np
|
|
from PIL import Image
|
|
|
|
def fourier_transform_image(img):
|
|
"""
|
|
Memory-optimized version that modifies the array in place when possible.
|
|
"""
|
|
# Convert with minimal copy
|
|
img_array = np.asarray(img, dtype=np.float32)
|
|
|
|
# Use rfft2 if the image is real to save memory
|
|
# and computation time (approximately 2x faster)
|
|
fft_result = np.fft.rfft2(img_array)
|
|
|
|
return fft_result
|
|
|
|
def attenuate_diagonal_frequencies(fft_spectrum, freq_threshold=0.30, target_angle=135,
|
|
angle_tolerance=10, attenuation_factor=0.15):
|
|
"""
|
|
Attenuates specific frequencies in the Fourier domain (optimized version for rfft2).
|
|
|
|
Args:
|
|
fft_spectrum: Result of 2D real Fourier transform (from rfft2)
|
|
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 of the rfft2 result
|
|
height, width_rfft = fft_spectrum.shape
|
|
# For rfft2, the original width is (width_rfft - 1) * 2
|
|
width_original = (width_rfft - 1) * 2
|
|
|
|
# Create frequency grids for rfft2 format
|
|
freq_y = np.fft.fftfreq(height, d=1.0)
|
|
freq_x = np.fft.rfftfreq(width_original, d=1.0) # Use rfftfreq for the X dimension
|
|
|
|
# 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
|
|
|
|
# Calculation of complementary angle
|
|
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 optimized inverse Fourier transform to reconstruct a PIL image.
|
|
|
|
Args:
|
|
fft_spectrum: Fourier transform result (complex array from rfft2)
|
|
original_shape: Original image shape (height, width) for proper cropping
|
|
|
|
Returns:
|
|
PIL.Image: Reconstructed image
|
|
"""
|
|
# Perform inverse Fourier transform
|
|
img_reconstructed = np.fft.irfft2(fft_spectrum)
|
|
|
|
# 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
|
|
|
|
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
|