Spaces:
Running
Running
import numpy as np | |
import scipy | |
from PIL import Image | |
VALUE = 512 | |
def resize(value,img): | |
img = Image.open(img) | |
#img = img.resize((value,value), Image.Resampling.LANCZOS) | |
img.thumbnail((VALUE,VALUE), Image.Resampling.LANCZOS) | |
return img | |
def get_mask(img,p): | |
w,h=img.size | |
return np.random.choice(a=[False, True], size=(w, h), p=[p, 1-p]) | |
def generate_points(mask): | |
(w,h) = mask.shape | |
noise_points = [] | |
color_points = [] | |
for x in range(w): | |
for y in range(h): | |
if mask[x,y]: | |
color_points.append(np.array([x,y])) | |
else: | |
noise_points.append(np.array([x,y])) | |
return color_points, noise_points | |
def encoder_cp(img,color_points): | |
w,h=img.size | |
img2=Image.new('RGB',(w,h)) | |
for p in color_points: | |
t = img.getpixel((p[0],p[1])) | |
img2.putpixel((p[0],p[1]),(t[0],t[1],t[2])) | |
return img2 | |
def encoder(img,p=0.95): | |
img = resize(VALUE,img) | |
mask = get_mask(img,p) | |
c_p, n_p = generate_points(mask) | |
return encoder_cp(img, c_p) | |
def get_points(img): | |
w,h=img.size | |
noise_points = [] | |
color_points = [] | |
for x in range(w): | |
for y in range(h): | |
t = img.getpixel((x,y)) | |
if np.sum(t[:3]) > 0 : | |
color_points.append(np.array([x,y])) | |
else: | |
noise_points.append(np.array([x,y])) | |
return color_points, noise_points | |
def restore(img, k, color_points, noise_points): | |
kdtree = scipy.spatial.KDTree(color_points) | |
for p in noise_points: | |
_, knn_p = kdtree.query(p, k) | |
r_m = [] | |
v_m = [] | |
b_m = [] | |
if k == 1: | |
c_p = color_points[knn_p] | |
t = img.getpixel((c_p[0],c_p[1])) | |
img.putpixel((p[0],p[1]),(t[0],t[1],t[2])) | |
else: | |
for c_p in [color_points[j] for j in list(knn_p)]: | |
t = img.getpixel((c_p[0],c_p[1])) | |
r_m.append(t[0]) | |
v_m.append(t[1]) | |
b_m.append(t[2]) | |
r_m = int(sum(r_m)/k) | |
v_m = int(sum(v_m)/k) | |
b_m = int(sum(b_m)/k) | |
img.putpixel((p[0],p[1]),(r_m,v_m,b_m)) | |
return img | |
def decoder(img,k=1): | |
img = resize(VALUE,img) | |
c_p, n_p = get_points(img) | |
return restore(img,int(k),c_p,n_p) | |
def decoder_noise(img,k=1): | |
img = Image.fromarray(img) | |
#img = resize(VALUE,img) | |
c_p, n_p = get_points(img) | |
return restore(img,int(k),c_p,n_p) |