Update app.py

app.py

@@ -2,86 +2,78 @@ import cv2
 import numpy as np
 import scipy as sp
 import scipy.sparse.linalg
-from numba import jit, prange
+from numba import njit, 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
+@njit
+def neighbours(y, x, max_y, max_x):
+    neighbors = []
+    if y > 0:
+        neighbors.append((y-1, x))
+    if y < max_y:
+        neighbors.append((y+1, x))
+    if x > 0:
+        neighbors.append((y, x-1))
+    if x < max_x:
+        neighbors.append((y, x+1))
+    return neighbors
 
-@jit(nopython=True)
-def build_poisson_matrix(img_h, img_w, alpha):
+@njit
+def poisson_sharpening(img, alpha):
+    img_h, img_w = img.shape
+    img_s = img.copy()
+    
     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
+    A_data = np.zeros(img_h * img_w * 5, dtype=np.float64)
+    A_row = np.zeros(img_h * img_w * 5, dtype=np.int32)
+    A_col = np.zeros(img_h * img_w * 5, dtype=np.int32)
+    b = np.zeros(img_h * img_w * 5, dtype=np.float64)
     
-    return A_data, A_row, A_col, b, e
+    return _poisson_sharpening_inner(img_s, alpha, im2var, A_data, A_row, A_col, b, img_h, img_w)
 
-@jit(nopython=True, parallel=True)
-def fill_b_vector(b, img, alpha, img_h, img_w):
+@njit(parallel=True)
+def _poisson_sharpening_inner(img_s, alpha, im2var, A_data, A_row, A_col, b, img_h, img_w):
     e = 0
     for y in prange(img_h):
         for x in range(img_w):
-            b[e] = img[y, x]
+            A_data[e] = 1
+            A_row[e] = e
+            A_col[e] = im2var[y, x]
+            b[e] = img_s[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])
+                A_data[e] = 1
+                A_row[e] = e
+                A_col[e] = im2var[y, x]
+                e += 1
+                
+                A_data[e] = -1
+                A_row[e] = e - 1
+                A_col[e] = im2var[n_y, n_x]
+                
+                b[e-1] = alpha * (img_s[y, x] - img_s[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]
-    
-    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)
     
-    A = sp.sparse.csr_matrix((A_data, (A_row, A_col)), shape=(e, img_h*img_w))
+    A = sp.sparse.csr_matrix((A_data[:e], (A_row[:e], A_col[:e])), 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)
 
+@njit(parallel=True)
+def sharpen_image_channels(img, alpha):
+    sharpen_img = np.zeros_like(img)
+    for b in prange(3):
+        sharpen_img[:,:,b] = poisson_sharpening(img[:,:,b], alpha)
+    return sharpen_img
+
 def get_image(img):
     return cv2.cvtColor(img, cv2.COLOR_BGR2RGB).astype('float32') / 255.0
 
 def sharpen_image(input_img, alpha):
     img = get_image(input_img)
-    
-    sharpen_img = np.zeros(img.shape)
-    for b in range(3):
-        sharpen_img[:,:,b] = poisson_sharpening(img[:,:,b], alpha)
-    
+    sharpen_img = sharpen_image_channels(img, alpha)
     return (sharpen_img * 255).astype(np.uint8)
 
 # Create examples list