Update app.py

app.py

@@ -1,59 +1,52 @@
 import numpy as np
-from numba import njit, prange, vectorize
+import cv2
+import matplotlib.pyplot as plt
+from numba import jit
 import gradio as gr
 
-@vectorize(['float64(float64, float64)'], nopython=True)
-def clip(a, max_value):
-    return min(max(a, 0), max_value)
+@jit(nopython=True, parallel=True)
+def poisson_sharpening_rgb(image, alpha):
+    height, width, channels = image.shape
+    sharpened = np.zeros_like(image, dtype=np.float32)
 
-@njit
-def neighbours(y, x, max_y, max_x):
-    n = []
-    if y > 0:
-        n.append((y-1, x))
-    if y < max_y:
-        n.append((y+1, x))
-    if x > 0:
-        n.append((y, x-1))
-    if x < max_x:
-        n.append((y, x+1))
-    return n
+    for c in range(channels):
+        for i in range(height):
+            for j in range(width):
+                # Compute indices for neighboring pixels
+                left = max(0, j - 1)
+                right = min(width - 1, j + 1)
+                top = max(0, i - 1)
+                bottom = min(height - 1, i + 1)
 
-@njit
-def poisson_sharpening(img, alpha, num_iterations=50):
-    img_h, img_w = img.shape
-    v = img.copy()
-    
-    for _ in range(num_iterations):
-        for y in range(img_h):
-            for x in range(img_w):
-                neighbors = neighbours(y, x, img_h-1, img_w-1)
-                num_neighbors = len(neighbors)
-                neighbor_sum = sum(v[ny, nx] for ny, nx in neighbors)
-                laplacian = neighbor_sum - num_neighbors * v[y, x]
-                v[y, x] += (laplacian + alpha * (img[y, x] - v[y, x])) / (num_neighbors + alpha)
-    
-    return clip(v, 1.0)
+                # Compute differences with neighboring pixels
+                diff_left = float(image[i, j, c]) - float(image[i, left, c])
+                diff_right = float(image[i, j, c]) - float(image[i, right, c])
+                diff_top = float(image[i, j, c]) - float(image[top, j, c])
+                diff_bottom = float(image[i, j, c]) - float(image[bottom, j, c])
 
-@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
+                # Compute sharpened pixel value
+                sharpened[i, j, c] = min(max(
+                    float(image[i, j, c]) + alpha * (diff_left + diff_right + diff_top + diff_bottom),
+                    0.0), 255.0)
+
+    return sharpened.astype(np.uint8)
+
+def sharpen_image(image, alpha):
+    # Ensure the image is in RGB format
+    if image.shape[2] == 4:  # If RGBA, convert to RGB
+        image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
+    elif len(image.shape) == 2:  # If grayscale, convert to RGB
+        image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
     
-    sharpen_img = np.zeros_like(img)
-    for b in prange(3):
-        sharpen_img[:,:,b] = poisson_sharpening(img[:,:,b], alpha)
+    # Apply sharpening
+    sharpened = poisson_sharpening_rgb(image, alpha)
     
-    # Convert back to uint8 for Gradio output
-    return (sharpen_img * 255).astype(np.uint8)
+    return sharpened
 
 # Create examples list
 examples = [
-    ["img1.jpg", 9.0],
-    ["img2.PNG", 7.0],
+    ["img1.jpg", 2.0],
+    ["img2.PNG", 2.0],
 ]
 
 # Create the Gradio interface
@@ -61,7 +54,7 @@ iface = gr.Interface(
     fn=sharpen_image,
     inputs=[
         gr.Image(label="Input Image", type="numpy"),
-        gr.Slider(minimum=1.0, maximum=15.0, step=0.01, value=9.0, label="Sharpening Strength (alpha)")
+        gr.Slider(minimum=1.0, maximum=15.0, step=0.01, value=2.0, label="Sharpening Strength (alpha)")
     ],
     outputs=gr.Image(label="Sharpened Image"),
     title="Poisson Image Sharpening",