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) | |
| 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): | |
| img = resize(img,VALUE) | |
| 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 = [] | |
| for np in [color_points[j] for j in knn_p]: | |
| t = img.getpixel((np[0],np[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): | |
| c_p, n_p = get_points(img) | |
| return restore(img,k,c_p,n_p) |