feat: further tests

deblur/deblur_2d.py

@@ -209,40 +209,89 @@ def kernel_detection(blurred, mask, edge_threshold=30, profile_length=21):
 
     kernel = create_gaussian_kernel(sigma)
 
-    print(kernel)
+    # print(kernel)
 
     return kernel / kernel.sum()
 
 
 
+
+def kernel_detection_box(blurred, mask, edge_threshold=30, profile_length=21):
+    def box_function(x, amp, center, width):
+        """Simple box profile: flat region with sharp transitions."""
+        return amp * ((x >= (center - width / 2)) & (x <= (center + width / 2))).astype(float)
+
+    def fit_box(profile, x_vals):
+        # Initial guess: full amplitude, centered at 0, small width
+        p0 = [1.0, 0.0, 5.0]
+        bounds = ([0, -10, 1], [1.5, 10, len(x_vals)])  # reasonable bounds
+        popt, _ = curve_fit(box_function, x_vals, profile, p0=p0, bounds=bounds)
+        return popt  # amp, center, width
+
+    def create_box_kernel(width):
+        """Generate a normalized 2D box kernel."""
+        ksize = int(round(width))
+        if ksize < 1:
+            ksize = 1
+        if ksize % 2 == 0:
+            ksize += 1  # ensure odd size
+        kernel = np.ones((ksize, ksize), dtype=np.float32)
+        return kernel / kernel.sum()
+
+
+    edges, gradient_mag = color_edge_detection(blurred, threshold=edge_threshold)
+    edges = cv2.bitwise_and(edges, edges, mask=mask)
+
+    y_idxs, x_idxs = np.where(edges > 0)
+    if len(x_idxs) == 0:
+        raise RuntimeError("No edges found.")
+    idx = len(x_idxs) // 2
+    cx, cy = x_idxs[idx], y_idxs[idx]
+
+    gray = cv2.cvtColor(blurred, cv2.COLOR_BGR2GRAY)
+    profile, x_vals = extract_vertical_profile(gray, cx, cy, length=profile_length)
+    popt = fit_box(profile, x_vals)
+    amp, mu, width = popt
+
+    print(f"Estimated box width: {width:.2f} pixels")
+
+    kernel = create_box_kernel(width)
+    return kernel
+
+
+
 def deconvolution(image_file, edge_threshold=30, profile_length=21):
-    blurred = cv2.imread(image_file)
+    image = cv2.imread(image_file)
     mask = get_mask(image_file)
+    image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB).astype(np.float32) / 255.0
 
-    kernel = kernel_detection(blurred, mask, edge_threshold=edge_threshold, profile_length=profile_length)
-    
-    test_blurred = cv2.imread(image_file, cv2.IMREAD_GRAYSCALE).astype(np.float32) / 255.0
+    kernel = kernel_detection_box(image, mask, edge_threshold=edge_threshold, profile_length=profile_length)
+
+
+    # Apply Richardson-Lucy to each channel
+    num_iter = 30
+    deblurred_channels = []
+    for i in range(3):  # R, G, B
+        channel = image_rgb[..., i]
+        deconv = richardson_lucy(channel, kernel, num_iter=num_iter)
+        deblurred_channels.append(deconv)
+
+    # Stack back into an RGB image
+    deblurred_rgb = np.stack(deblurred_channels, axis=2)
+    deblurred_rgb = np.clip(deblurred_rgb, 0, 1)
 
-    deconvolved = richardson_lucy(test_blurred, kernel, num_iter=30)
     
-    # Display results
-    plt.figure(figsize=(12, 4))
-    plt.subplot(1, 3, 1)
+    # Show result
+    plt.figure(figsize=(10, 5))
+    plt.subplot(1, 2, 1)
+    plt.imshow(image_rgb)
     plt.title("Blurred Image")
-    plt.imshow(test_blurred, cmap='gray')
     plt.axis('off')
 
-    plt.subplot(1, 3, 2)
-    plt.title("Estimated Kernel")
-    plt.imshow(kernel, cmap='hot')
-    plt.axis('off')
-
-    plt.subplot(1, 3, 3)
+    plt.subplot(1, 2, 2)
+    plt.imshow(deblurred_rgb)
     plt.title("Deconvolved Image")
-    plt.imshow(deconvolved, cmap='gray')
     plt.axis('off')
-
-    plt.tight_layout()
     plt.show()
 
 
@@ -253,4 +302,4 @@ if __name__ == "__main__":
     img_file = "assets/real_test.jpg"
 
     #demo("assets/omas.png")
-    deconvolution(img_file, edge_threshold=10)
+    deconvolution(img_file, edge_threshold=5)