| import numpy as np | |
| import matplotlib.pyplot as plt | |
| import pywt | |
| import gradio as gr | |
| LEVEL_THRESHOLDS = { | |
| 1: 20, | |
| 2: 40, | |
| 3: 60, | |
| 4: 80 | |
| } | |
| def threshold_single_level(coeffs_level, threshold): | |
| LH, HL, HH = coeffs_level | |
| new_LH = np.where(np.abs(LH) < threshold, 0, LH) | |
| new_HL = np.where(np.abs(HL) < threshold, 0, HL) | |
| new_HH = np.where(np.abs(HH) < threshold, 0, HH) | |
| return (new_LH, new_HL, new_HH) | |
| def wavelet_compress(img, level): | |
| if img.ndim == 3: | |
| img = img[..., :3] | |
| img = np.dot(img, [0.299, 0.587, 0.114]) | |
| img = img.astype(np.uint8) | |
| h, w = img.shape | |
| threshold = LEVEL_THRESHOLDS[level] | |
| coeffs = pywt.wavedec2(img, "haar", level=level) | |
| new_coeffs = [coeffs[0]] | |
| for level_tuple in coeffs[1:]: | |
| new_tuple = threshold_single_level(level_tuple, threshold) | |
| new_coeffs.append(new_tuple) | |
| recon = pywt.waverec2(new_coeffs, "haar") | |
| recon = recon[:h, :w] | |
| recon = np.clip(recon, 0, 255).astype(np.uint8) | |
| return recon | |
| def interface(img, level): | |
| return wavelet_compress(img, level) | |
| demo = gr.Interface( | |
| fn=interface, | |
| inputs=[ | |
| gr.Image(type="numpy", label="Upload Image"), | |
| gr.Slider(1, 4, step=1, value=1, label="Wavelet Level") | |
| ], | |
| outputs=gr.Image(label="Reconstructed Image"), | |
| title="Wavelet Image Compression (A-Version)", | |
| description="Higher wavelet levels use larger thresholds, resulting in stronger distortion." | |
| ) | |
| demo.launch() |