app.py

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 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]
    
    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))
    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):
    img = get_image(input_img)
    
    sharpen_img = np.zeros(img.shape)
    for b in range(3):
        sharpen_img[:,:,b] = poisson_sharpening(img[:,:,b], alpha)
    
    return (sharpen_img * 255).astype(np.uint8)

# Create examples list
examples = [
    ["img1.jpg", 9.0],
    ["img2.PNG", 7.0],
]

# Create the Gradio interface
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)")
    ],
    outputs=gr.Image(label="Sharpened Image"),
    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,
    cache_examples=True
)

iface.launch()