Update app.py

app.py

@@ -19,28 +19,20 @@ def neighbours(y, x, max_y, max_x):
     return neighbors
 
 @njit
-def poisson_sharpening(img, alpha):
-    img_h, img_w = img.shape
-    img_s = img.copy()
-    
+def build_poisson_matrix(img, alpha, img_h, img_w):
     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)
     
-    return _poisson_sharpening_inner(img_s, alpha, im2var, A_data, A_row, A_col, b, 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 y in range(img_h):
         for x in range(img_w):
             A_data[e] = 1
             A_row[e] = e
             A_col[e] = im2var[y, x]
-            b[e] = img_s[y, x]
+            b[e] = img[y, x]
             e += 1
             
             for n_y, n_x in neighbours(y, x, img_h-1, img_w-1):
@@ -53,11 +45,17 @@ def _poisson_sharpening_inner(img_s, alpha, im2var, A_data, A_row, A_col, b, img
                 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])
+                b[e-1] = alpha * (img[y, x] - img[n_y, n_x])
                 e += 1
     
-    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 A_data[:e], A_row[:e], A_col[:e], b[:e], e
+
+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 np.clip(v.reshape(img_h, img_w), 0, 1)