|
|
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) |
|
|
|
|
|
|
|
|
examples = [ |
|
|
["img1.jpg", 9.0], |
|
|
["img2.PNG", 7.0], |
|
|
] |
|
|
|
|
|
|
|
|
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() |
