generated from Hazel/python-project
252 lines
8.1 KiB
Python
252 lines
8.1 KiB
Python
import sys
|
||
import numpy as np
|
||
import cv2
|
||
from PyQt5.QtWidgets import (
|
||
QApplication, QWidget, QLabel, QSlider, QVBoxLayout,
|
||
QHBoxLayout, QGridLayout, QPushButton, QFileDialog
|
||
)
|
||
from PyQt5.QtCore import Qt
|
||
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg as FigureCanvas
|
||
from matplotlib.figure import Figure
|
||
import scipy.signal
|
||
from scipy.signal import convolve2d
|
||
|
||
import os
|
||
os.environ.pop("QT_QPA_PLATFORM_PLUGIN_PATH", None)
|
||
|
||
|
||
def generate_box_kernel(size):
|
||
return np.ones((size, size), dtype=np.float32) / (size * size)
|
||
|
||
def generate_disk_kernel(radius):
|
||
size = 2 * radius + 1
|
||
y, x = np.ogrid[-radius:radius+1, -radius:radius+1]
|
||
mask = x**2 + y**2 <= radius**2
|
||
kernel = np.zeros((size, size), dtype=np.float32)
|
||
kernel[mask] = 1
|
||
kernel /= kernel.sum()
|
||
return kernel
|
||
|
||
def generate_kernel(radius, sigma=None):
|
||
"""
|
||
Generate a 2D Gaussian kernel with a given radius.
|
||
|
||
Parameters:
|
||
- radius: int, the radius of the kernel (size will be 2*radius + 1)
|
||
- sigma: float (optional), standard deviation of the Gaussian. If None, sigma = radius / 3
|
||
|
||
Returns:
|
||
- kernel: 2D numpy array of shape (2*radius+1, 2*radius+1)
|
||
"""
|
||
size = 2 * radius + 1
|
||
if sigma is None:
|
||
sigma = radius / 3.0 # Common default choice
|
||
print(f"radius: {radius}, sigma: {sigma}")
|
||
|
||
# Create a grid of (x,y) coordinates
|
||
ax = np.arange(-radius, radius + 1)
|
||
xx, yy = np.meshgrid(ax, ax)
|
||
|
||
# Apply the 2D Gaussian formula
|
||
kernel = np.exp(-(xx**2 + yy**2) / (2 * sigma**2))
|
||
kernel /= 2 * np.pi * sigma**2 # Normalize based on Gaussian PDF
|
||
kernel /= kernel.sum() # Normalize to sum to 1
|
||
|
||
return kernel
|
||
|
||
|
||
def wiener_deconvolution(blurred, kernel, K=0.1):
|
||
"""
|
||
Perform Wiener deconvolution on a 2D image.
|
||
|
||
Parameters:
|
||
- blurred: 2D numpy array (blurred image)
|
||
- kernel: 2D numpy array (PSF / blur kernel)
|
||
- K: float, estimated noise-to-signal ratio
|
||
|
||
Returns:
|
||
- deconvolved: 2D numpy array (deblurred image)
|
||
"""
|
||
# Pad kernel to image size
|
||
kernel /= np.sum(kernel)
|
||
pad = [(0, blurred.shape[0] - kernel.shape[0]),
|
||
(0, blurred.shape[1] - kernel.shape[1])]
|
||
kernel_padded = np.pad(kernel, pad, 'constant')
|
||
|
||
# FFT of image and kernel
|
||
H = np.fft.fft2(kernel_padded)
|
||
G = np.fft.fft2(blurred)
|
||
|
||
# Avoid division by zero
|
||
H_conj = np.conj(H)
|
||
denominator = H_conj * H + K
|
||
F_hat = H_conj / denominator * G
|
||
|
||
# Inverse FFT to get result
|
||
deconvolved = np.fft.ifft2(F_hat)
|
||
deconvolved = np.abs(deconvolved)
|
||
deconvolved = np.clip(deconvolved, 0, 255)
|
||
return deconvolved.astype(np.uint8)
|
||
|
||
|
||
def richardson_lucy(image, psf, iterations=30, clip=True):
|
||
image = image.astype(np.float32) + 1e-6
|
||
psf = psf / psf.sum()
|
||
estimate = np.full(image.shape, 0.5, dtype=np.float32)
|
||
|
||
psf_mirror = psf[::-1, ::-1]
|
||
|
||
for _ in range(iterations):
|
||
conv = convolve2d(estimate, psf, mode='same', boundary='wrap')
|
||
relative_blur = image / (conv + 1e-6)
|
||
estimate *= convolve2d(relative_blur, psf_mirror, mode='same', boundary='wrap')
|
||
|
||
if clip:
|
||
estimate = np.clip(estimate, 0, 255)
|
||
|
||
return estimate
|
||
|
||
|
||
class KernelVisualizer(QWidget):
|
||
def __init__(self, image_path=None):
|
||
super().__init__()
|
||
self.setWindowTitle("Gaussian Kernel Visualizer")
|
||
self.image = None
|
||
self.deconvolved = None
|
||
|
||
self.load_button = QPushButton("Load Image")
|
||
self.load_button.clicked.connect(self.load_image)
|
||
|
||
self.radius_slider = QSlider(Qt.Horizontal)
|
||
self.radius_slider.setRange(1, 100)
|
||
self.radius_slider.setValue(5)
|
||
|
||
self.sigma_slider = QSlider(Qt.Horizontal)
|
||
self.sigma_slider.setRange(1, 300)
|
||
self.sigma_slider.setValue(15)
|
||
|
||
self.radius_slider.valueChanged.connect(self.update_visualization)
|
||
self.sigma_slider.valueChanged.connect(self.update_visualization)
|
||
|
||
self.kernel_fig = Figure(figsize=(3, 3))
|
||
self.kernel_canvas = FigureCanvas(self.kernel_fig)
|
||
|
||
self.image_fig = Figure(figsize=(6, 3))
|
||
self.image_canvas = FigureCanvas(self.image_fig)
|
||
|
||
self.iter_slider = QSlider(Qt.Horizontal)
|
||
self.iter_slider.setRange(1, 50)
|
||
self.iter_slider.setValue(10)
|
||
|
||
self.apply_button = QPushButton("Do Deconvolution.")
|
||
self.apply_button.clicked.connect(self.apply_kernel)
|
||
|
||
layout = QVBoxLayout()
|
||
layout.addWidget(self.load_button)
|
||
|
||
|
||
sliders_layout = QGridLayout()
|
||
sliders_layout.addWidget(QLabel("Radius:"), 0, 0)
|
||
sliders_layout.addWidget(self.radius_slider, 0, 1)
|
||
sliders_layout.addWidget(QLabel("Sigma:"), 1, 0)
|
||
sliders_layout.addWidget(self.sigma_slider, 1, 1)
|
||
|
||
sliders_layout.addWidget(QLabel("Iterations:"), 2, 0)
|
||
sliders_layout.addWidget(self.iter_slider, 2, 1)
|
||
sliders_layout.addWidget(self.apply_button, 3, 1)
|
||
|
||
layout.addLayout(sliders_layout)
|
||
layout.addWidget(QLabel("Kernel Visualization:"))
|
||
layout.addWidget(self.kernel_canvas)
|
||
layout.addWidget(QLabel("Original and Deconvolved Image:"))
|
||
layout.addWidget(self.image_canvas)
|
||
|
||
self.setLayout(layout)
|
||
|
||
if image_path:
|
||
self.load_image(image_path)
|
||
else:
|
||
self.update_visualization()
|
||
|
||
def load_image(self, image_path=None):
|
||
if not image_path:
|
||
fname, _ = QFileDialog.getOpenFileName(self, "Open Image", "", "Images (*.png *.jpg *.bmp *.jpeg)")
|
||
image_path = fname
|
||
|
||
if image_path:
|
||
img = cv2.imread(image_path)
|
||
if img is not None:
|
||
self.image = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||
self.update_visualization()
|
||
|
||
def load_image(self, image_path=None):
|
||
if not image_path:
|
||
fname, _ = QFileDialog.getOpenFileName(self, "Open Image", "", "Images (*.png *.jpg *.bmp *.jpeg)")
|
||
image_path = fname
|
||
|
||
if image_path:
|
||
img = cv2.imread(image_path)
|
||
if img is not None:
|
||
self.image = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
|
||
self.image = cv2.resize(self.image, (200, 200))
|
||
self.update_visualization()
|
||
|
||
def apply_kernel(self):
|
||
radius = self.radius_slider.value()
|
||
sigma = self.sigma_slider.value() / 10.0
|
||
iterations = self.iter_slider.value()
|
||
|
||
kernel = generate_kernel(radius, sigma)
|
||
|
||
self.deconvolved = richardson_lucy(self.image, kernel, iterations=iterations)
|
||
|
||
self.update_visualization()
|
||
|
||
def update_visualization(self):
|
||
radius = self.radius_slider.value()
|
||
sigma = self.sigma_slider.value() / 10.0 * (radius / 3)
|
||
kernel = generate_kernel(radius, sigma)
|
||
iterations = self.iter_slider.value()
|
||
|
||
|
||
# Kernel Visualization
|
||
self.kernel_fig.clear()
|
||
ax = self.kernel_fig.add_subplot(111)
|
||
cax = ax.imshow(kernel, cmap='hot')
|
||
self.kernel_fig.colorbar(cax, ax=ax)
|
||
ax.set_title(f"Kernel (r={radius}, σ={sigma:.2f})")
|
||
self.kernel_canvas.draw()
|
||
|
||
if self.image is not None:
|
||
self.image_fig.clear()
|
||
ax1 = self.image_fig.add_subplot(131)
|
||
ax1.imshow(self.image, cmap='gray')
|
||
ax1.set_title("Original")
|
||
ax1.axis('off')
|
||
|
||
if self.deconvolved is not None:
|
||
ax3 = self.image_fig.add_subplot(133)
|
||
ax3.imshow(self.deconvolved, cmap='gray')
|
||
ax3.set_title(f"Deconvolved (RL, {iterations} iter)")
|
||
ax3.axis('off')
|
||
|
||
self.image_canvas.draw()
|
||
else:
|
||
self.image_fig.clear()
|
||
ax = self.image_fig.add_subplot(111)
|
||
ax.text(0.5, 0.5, "No image loaded", fontsize=14, ha='center', va='center')
|
||
ax.axis('off')
|
||
self.image_canvas.draw()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
image_path = None
|
||
if len(sys.argv) > 1:
|
||
image_path = sys.argv[1] # Get image path from command-line argument
|
||
print(image_path)
|
||
|
||
app = QApplication(sys.argv)
|
||
viewer = KernelVisualizer(image_path=image_path)
|
||
viewer.show()
|
||
sys.exit(app.exec_())
|