PixClean-enhancer / utils.py
Keramo's picture
Create utils.py
cf0f768 verified
Raw
History Blame Contribute Delete
3.37 kB
import cv2
import numpy as np
from PIL import Image, ImageEnhance, ImageFilter
import io
from scipy.signal import wiener
from skimage.restoration import denoise_nl_means, estimate_sigma
def auto_white_balance(image_bytes: bytes) -> bytes:
"""Gray-world white balance"""
nparr = np.frombuffer(image_bytes, np.uint8)
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
result = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
avg_a = np.average(result[:, :, 1])
avg_b = np.average(result[:, :, 2])
result[:, :, 1] = result[:, :, 1] - ((avg_a - 128) * (result[:, :, 1] / 255.0))
result[:, :, 2] = result[:, :, 2] - ((avg_b - 128) * (result[:, :, 2] / 255.0))
result = cv2.cvtColor(result, cv2.COLOR_LAB2BGR)
_, encoded = cv2.imencode('.png', result)
return encoded.tobytes()
def adjust_exposure(image_bytes: bytes, exposure: float = 1.0) -> bytes:
"""exposure factor >1 brightens, <1 darkens"""
img = Image.open(io.BytesIO(image_bytes))
enhancer = ImageEnhance.Brightness(img)
img = enhancer.enhance(exposure)
out = io.BytesIO()
img.save(out, format='PNG')
return out.getvalue()
def recover_shadows_highlights(image_bytes: bytes, shadow_boost: float = 1.2, highlight_reduce: float = 0.8) -> bytes:
"""Simple gamma correction in shadows/highlights"""
nparr = np.frombuffer(image_bytes, np.uint8)
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR).astype(np.float32) / 255.0
# Adaptive gamma: shadows ^0.7, highlights ^1.3
gamma = np.ones_like(img)
gamma[img < 0.3] = shadow_boost
gamma[img > 0.7] = highlight_reduce
img = np.power(img, gamma)
img = (img * 255).astype(np.uint8)
_, encoded = cv2.imencode('.png', img)
return encoded.tobytes()
def advanced_denoise(image_bytes: bytes, patch_size: int = 5, h: float = 1.0) -> bytes:
"""Non-local means denoising with auto sigma estimation"""
nparr = np.frombuffer(image_bytes, np.uint8)
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
# Estimate noise sigma
sigma_est = estimate_sigma(img, channel_axis=-1, average_sigmas=True)
denoised = denoise_nl_means(img, h=h * sigma_est, patch_size=patch_size, patch_distance=6, channel_axis=-1)
denoised = (denoised * 255).astype(np.uint8)
_, encoded = cv2.imencode('.png', denoised)
return encoded.tobytes()
def deblur(image_bytes: bytes, kernel_size: int = 5, noise_var: float = 0.01) -> bytes:
"""Wiener deconvolution – assumes Gaussian blur"""
nparr = np.frombuffer(image_bytes, np.uint8)
img = cv2.imdecode(nparr, cv2.IMREAD_COLOR).astype(np.float32) / 255.0
# Create a simple Gaussian kernel
kernel = cv2.getGaussianKernel(kernel_size, -1)
kernel = kernel @ kernel.T
# Apply Wiener filter to each channel
from scipy.signal import convolve2d
restored = np.zeros_like(img)
for c in range(3):
channel = img[:,:,c]
# Wiener in frequency domain (simplified)
from scipy.fft import fft2, ifft2, fftshift
H = fft2(kernel, s=channel.shape)
Y = fft2(channel)
H_conj = np.conj(H)
Wiener = H_conj / (np.abs(H)**2 + noise_var)
restored_channel = np.real(ifft2(Y * Wiener))
restored[:,:,c] = np.clip(restored_channel, 0, 1)
restored = (restored * 255).astype(np.uint8)
_, encoded = cv2.imencode('.png', restored)
return encoded.tobytes()