Update app.py

app.py

@@ -1,37 +1,93 @@
 import cv2
 import numpy as np
+import scipy as sp
+import scipy.sparse.linalg
 from numba import jit, prange
 import gradio as gr
 
+@jit(nopython=True)
+def neighbours(i, j, max_i, max_j):
+    pairs = []
+    for n in [-1, 1]:
+        if 0 <= i+n <= max_i:
+            pairs.append((i+n, j))
+        if 0 <= j+n <= max_j:
+            pairs.append((i, j+n))
+    return pairs
+
+@jit(nopython=True)
+def build_poisson_matrix(img_h, img_w, alpha):
+    im2var = np.arange(img_h * img_w).reshape(img_h, img_w)
+    A_data = []
+    A_row = []
+    A_col = []
+    b = np.zeros(img_h*img_w*5)
+    
+    e = 0
+    for y in range(img_h):
+        for x in range(img_w):
+            A_data.append(1)
+            A_row.append(e)
+            A_col.append(im2var[y, x])
+            e += 1
+            
+            for n_y, n_x in neighbours(y, x, img_h-1, img_w-1):
+                A_data.append(1)
+                A_row.append(e)
+                A_col.append(im2var[y, x])
+                
+                A_data.append(-1)
+                A_row.append(e)
+                A_col.append(im2var[n_y, n_x])
+                
+                e += 1
+    
+    return A_data, A_row, A_col, b, e
+
 @jit(nopython=True, parallel=True)
-def fast_poisson_sharpening(img, alpha):
-    img_h, img_w = img.shape
-    output = np.zeros_like(img)
+def fill_b_vector(b, img, alpha, img_h, img_w):
+    e = 0
+    for y in prange(img_h):
+        for x in range(img_w):
+            b[e] = img[y, x]
+            e += 1
+            
+            for n_y, n_x in neighbours(y, x, img_h-1, img_w-1):
+                b[e] = alpha * (img[y, x] - img[n_y, n_x])
+                e += 1
+
+def poisson_sharpening(img: np.ndarray, alpha: float) -> np.ndarray:
+    """
+    Returns a sharpened image with strength of alpha.
+    :param img: the image
+    :param alpha: edge threshold and gradient scaler
+    """
+    img_h, img_w = img.shape[:2]
     
-    for y in prange(1, img_h - 1):
-        for x in range(1, img_w - 1):
-            laplacian = (
-                img[y-1, x] + img[y+1, x] + img[y, x-1] + img[y, x+1] - 4 * img[y, x]
-            )
-            output[y, x] = img[y, x] + alpha * laplacian
+    A_data, A_row, A_col, b, e = build_poisson_matrix(img_h, img_w, alpha)
+    fill_b_vector(b, img, alpha, img_h, img_w)
     
-    return np.clip(output, 0, 1)
+    A = sp.sparse.csr_matrix((A_data, (A_row, A_col)), shape=(e, img_h*img_w))
+    v = sp.sparse.linalg.lsqr(A, b[:e])[0]
+
+    return np.clip(v.reshape(img_h, img_w), 0, 1)
+
+def get_image(img):
+    return cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype('float32') / 255.0
 
 def sharpen_image(input_img, alpha):
-    # Convert BGR to RGB
-    img = cv2.cvtColor(input_img, cv2.COLOR_BGR2RGB).astype('float32') / 255.0
+    img = get_image(input_img)
     
-    sharpen_img = np.zeros_like(img)
+    sharpen_img = np.zeros(img.shape)
     for b in range(3):
-        sharpen_img[:,:,b] = fast_poisson_sharpening(img[:,:,b], alpha)
+        sharpen_img[:,:,b] = poisson_sharpening(img[:,:,b], alpha)
     
-    # Convert back to BGR for output
-    return cv2.cvtColor((sharpen_img * 255).astype(np.uint8), cv2.COLOR_RGB2BGR)
+    return (sharpen_img * 255).astype(np.uint8)
 
 # Create examples list
 examples = [
-    ["img1.jpg", 0.5],
-    ["img2.PNG", 0.3],
+    ["img1.jpg", 9.0],
+    ["img2.PNG", 7.0],
 ]
 
 # Create the Gradio interface
@@ -39,10 +95,10 @@ iface = gr.Interface(
     fn=sharpen_image,
     inputs=[
         gr.Image(label="Input Image", type="numpy"),
-        gr.Slider(minimum=0.1, maximum=1.0, step=0.01, value=0.5, label="Sharpening Strength (alpha)")
+        gr.Slider(minimum=1.0, maximum=15.0, step=0.01, value=9.0, label="Sharpening Strength (alpha)")
     ],
     outputs=gr.Image(label="Sharpened Image"),
-    title="Fast Poisson Image Sharpening",
+    title="Poisson Image Sharpening",
     description="Upload an image or choose from the examples, then adjust the sharpening strength to enhance edges and details.",
     theme='bethecloud/storj_theme',
     examples=examples,