Spaces:
Running
Running
File size: 2,468 Bytes
53b8c3a fa0d071 53b8c3a aeda751 fa0d071 53b8c3a e53dbad 53b8c3a aeda751 fa0d071 aeda751 fa0d071 53b8c3a e53dbad |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 |
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) |