diff --git a/kindlecomicconverter/rainbow_artifacts_eraser.py b/kindlecomicconverter/rainbow_artifacts_eraser.py new file mode 100644 index 0000000..4b81611 --- /dev/null +++ b/kindlecomicconverter/rainbow_artifacts_eraser.py @@ -0,0 +1,131 @@ +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 \ No newline at end of file