1
0
mirror of https://github.com/ciromattia/kcc synced 2025-12-13 09:46:25 +00:00

Update rainbow_artifacts_eraser.py

Enable odd width for BW image when using Rainbow Eraser by passing original image dimensions to irfft2
This commit is contained in:
Its-my-right
2025-07-20 20:08:22 +02:00
committed by Alex Xu
parent 26327728d0
commit f1ffb2c4e8

View File

@@ -108,7 +108,7 @@ def attenuate_diagonal_frequencies(fft_spectrum, freq_threshold=0.30, target_ang
fft_spectrum[combined_condition, :] *= attenuation_factor fft_spectrum[combined_condition, :] *= attenuation_factor
return fft_spectrum return fft_spectrum
def inverse_fourier_transform_image(fft_spectrum, is_color): def inverse_fourier_transform_image(fft_spectrum, is_color, original_shape=None):
""" """
Performs an optimized inverse Fourier transform to reconstruct a PIL image. Performs an optimized inverse Fourier transform to reconstruct a PIL image.
@@ -119,8 +119,11 @@ def inverse_fourier_transform_image(fft_spectrum, is_color):
Returns: Returns:
PIL.Image: Reconstructed image PIL.Image: Reconstructed image
""" """
# Perform inverse Fourier transform # Perform inverse Fourier transform with original shape if provided
img_reconstructed = np.fft.irfft2(fft_spectrum) if original_shape is not None:
img_reconstructed = np.fft.irfft2(fft_spectrum, s=original_shape)
else:
img_reconstructed = np.fft.irfft2(fft_spectrum)
# Normalize values between 0 and 255 # Normalize values between 0 and 255
img_reconstructed = np.clip(img_reconstructed, 0, 255) img_reconstructed = np.clip(img_reconstructed, 0, 255)
@@ -232,9 +235,12 @@ def erase_rainbow_artifacts(img, is_color):
# Grayscale processing (original behavior) # Grayscale processing (original behavior)
if img.mode != 'L': if img.mode != 'L':
img = img.convert('L') img = img.convert('L')
# Get original image dimensions
original_shape = (img.height, img.width)
fft_spectrum = fourier_transform_image(img) fft_spectrum = fourier_transform_image(img)
clean_spectrum = attenuate_diagonal_frequencies(fft_spectrum) clean_spectrum = attenuate_diagonal_frequencies(fft_spectrum)
clean_image = inverse_fourier_transform_image(clean_spectrum, is_color) clean_image = inverse_fourier_transform_image(clean_spectrum, is_color, original_shape)
return clean_image return clean_image