Update app.py

app.py

@@ -1,38 +1,49 @@
-import cv2
 import numpy as np
 import scipy as sp
-import scipy.sparse.linalg
-from numba import njit, prange
+from numba import njit, prange, vectorize
 import gradio as gr
 
+@vectorize(['float64(float64, float64)'], nopython=True)
+def clip(a, max_value):
+    return min(max(a, 0), max_value)
+
 @njit
 def neighbours(y, x, max_y, max_x):
-    neighbors = []
+    n = []
     if y > 0:
-        neighbors.append((y-1, x))
+        n.append((y-1, x))
     if y < max_y:
-        neighbors.append((y+1, x))
+        n.append((y+1, x))
     if x > 0:
-        neighbors.append((y, x-1))
+        n.append((y, x-1))
     if x < max_x:
-        neighbors.append((y, x+1))
-    return neighbors
+        n.append((y, x+1))
+    return n
 
-@njit
-def build_poisson_matrix(img, alpha, img_h, img_w):
+@njit(parallel=True)
+def poisson_sharpening(img, alpha):
+    """
+    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
+    img_s = img.copy()
+    
     im2var = np.arange(img_h * img_w).reshape(img_h, img_w)
-    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)
+    
+    A_data = np.zeros(img_h * img_w * 4 * 2, dtype=np.float64)
+    A_row = np.zeros(img_h * img_w * 4 * 2, dtype=np.int32)
+    A_col = np.zeros(img_h * img_w * 4 * 2, dtype=np.int32)
+    b = np.zeros(img_h * img_w * 4 * 2, dtype=np.float64)
     
     e = 0
-    for y in range(img_h):
-        for x in range(img_w):
+    for y in prange(img_h):
+        for x in prange(img_w):
             A_data[e] = 1
             A_row[e] = e
             A_col[e] = im2var[y, x]
-            b[e] = img[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):
@@ -45,32 +56,27 @@ def build_poisson_matrix(img, alpha, img_h, img_w):
                 A_row[e] = e - 1
                 A_col[e] = im2var[n_y, n_x]
                 
-                b[e-1] = alpha * (img[y, x] - img[n_y, n_x])
+                b[e-1] = alpha * (img_s[y, x] - img_s[n_y, n_x])
                 e += 1
     
-    return A_data[:e], A_row[:e], A_col[:e], b[:e], e
+    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]
 
-def poisson_sharpening(img: np.ndarray, alpha: float) -> np.ndarray:
-    img_h, img_w = img.shape
-    A_data, A_row, A_col, b, e = build_poisson_matrix(img, alpha, img_h, img_w)
-    
-    A = sp.sparse.csr_matrix((A_data, (A_row, A_col)), shape=(e, img_h * img_w))
-    v = sp.sparse.linalg.lsqr(A, b)[0]
+    return clip(v.reshape(img_h, img_w), 1.0)
 
-    return np.clip(v.reshape(img_h, img_w), 0, 1)
-
-def sharpen_image_channels(img, alpha):
+@njit(parallel=True)
+def sharpen_image(img, alpha):
+    # Convert alpha to float
+    alpha = float(alpha)
+    
+    # Ensure the image is in the correct format
+    img = img.astype(np.float64) / 255.0
+    
     sharpen_img = np.zeros_like(img)
-    for b in range(3):
+    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 = sharpen_image_channels(img, alpha)
+    
+    # Convert back to uint8 for Gradio output
     return (sharpen_img * 255).astype(np.uint8)
 
 # Create examples list