encoder

app.py

@@ -1,7 +1,14 @@
 import gradio as gr
+from utils import encoder, decoder
+
+source_img = gr.Image(source="upload", type="filepath", label="init_img | 512*512 px")
+encoded_img = gr.Image()
 
 def greet(name):
     return "Hello " + name + "!!"
 
 iface = gr.Interface(fn=greet, inputs="text", outputs="text")
-iface.launch()
+iface = gr.Interface(fn=encoder, inputs=[source_img,
+    gr.Slider(label='noise', minimum = 0, maximum = 1, step = .05, value = .95)],
+    outputs=encoded_img)
+iface.launch(share=True)

requirements.txt

@@ -0,0 +1 @@
+scipy

utils.py

@@ -0,0 +1,76 @@
+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)