code.py

# -*- coding: utf-8 -*-
"""
Created on Mon May  7 23:34:58 2018

@author: admin
"""

from sklearn import linear_model
import matplotlib.pyplot as plt
from matplotlib.colors import rgb_to_hsv
from matplotlib.colors import hsv_to_rgb
import numpy as np
import sys
import numpy.random as npr

class Processor:
    def __init__(self, h, filename):
        assert(h % 2 == 1)
        self.h = h
        self.load_img(filename)
        self.optimizer = linear_model.Lasso(alpha = 0.1)
        
    def load_img(self, filename):
        #print((plt.imread(filename)[ : , : , : 3 ]))
        self.img = rgb_to_hsv((plt.imread(filename)[ : , : , : 3 ] / 255)) - 0.5
        print(self.img)
        #print(hsv_to_rgb(np.array(self.img[0][0])) * 255)

        #print(self.img)
        
    def patch_as_vector(self, patch):
        print(patch)
        h, w, _ = patch.shape
        return np.reshape(patch, (h * w, 3))
    
    def vector_as_patch(self, vec):
        return np.reshape(vec, (self.h, self.h, 3))
    
    def create_void_patch(self):
        return np.full((self.h, self.h, 3), None)
    
    def replace_patch(self, i, j, patch):
        h_half = int(self.h / 2)
        self.img[i - h_half : i + h_half + 1, j - h_half : j + h_half + 1] = patch
        
    def noise_img(self, portion):
        h_half = int(self.h / 2)
        h, w, _ = self.img.shape
        
        h_ratio = int(h / self.h)
        w_ratio = int(h / self.h)
        
        total = h_ratio * w_ratio
        total_to_remove = int(total * portion)
        
        for _ in range(total_to_remove):
            i = int(npr.random() * h_ratio - sys.float_info.epsilon) * self.h + h_half
            j = int(npr.random() * w_ratio - sys.float_info.epsilon) * self.h + h_half

            print(i, j)
            #print(self.get_patch(i, j))
            
            while np.all(np.isnan(self.get_patch(i, j)[0][0])):
                print("redo")
                i = int(npr.random() * h_ratio - sys.float_info.epsilon) * self.h + h_half
                j = int(npr.random() * w_ratio - sys.float_info.epsilon) * self.h + h_half
             
            #self.img[i][j] = [None, None, None]
            self.replace_patch(i, j, self.create_void_patch())
            
    def retrieve_dico(self, stepping):
        complete = []
        h_half = int(self.h / 2)

        h, w, _ = self.img.shape
        
        h_ratio = int(h / stepping)
        w_ratio = int(w / stepping) 
        
        for i in range(w_ratio):
            for j in range(h_ratio):
                patch = self.get_patch(i + h_half, j + h_half)
                if not np.any(np.isnan(patch)):
                    #print(patch)
                    complete.append(patch)
                    
        return complete
        

    def show_img(self):
        plt.imshow((hsv_to_rgb(self.img + 0.5) * 255).astype(np.uint8))

        
    def get_patch(self, i, j):
        h_half = int(self.h / 2)
        #print(i, j)
        return self.img[i - h_half : i + h_half + 1, j - h_half : j + h_half + 1]
    
    def approximate_patch(self, dic, patch, alpha):
        if np.any(np.isnan(patch)):
            y = []
            x = []
            expressed_pixels = np.logical_not(np.isnan(patch))
            expressed_patch = patch[expressed_pixels]
                        #print(patch)
                        #print(expressed_patch)
    
            if expressed_patch.size == 0:
                return
            for p in expressed_patch:
                y.append(p)
                for elem in dic:
                    new_x = elem[expressed_pixels]
                            #x.append(np.reshape(new_x, (int(new_x.size / 3), 3)))
                    x.append(new_x)
                            #print(new_x.size)
                    model = linear_model.Lasso(alpha)
                    print(np.array(x).shape)
                        #print(x)
                        #print(len(x))
                    model.fit(np.array(x).reshape(len(y), len(x)), np.array(y))
                    c = model.coef_
                        
                        #dic = np.array(dic)
                        #dic = dic.reshape(int(dic.size / 3), 3)           
                        
                        
                    for ix in range(self.h):
                        for jx in range(self.h):
                            if np.any(expressed_pixels[ix, jx] == False):
                                new_pixel = [0, 0, 0]
                                    #print("fixing", ix, jx)
                                    #print(len(dic))
                                    #print(dic)
                                    #print(len(c))
                                for k in range(len(dic)):
                                        #print(y[ix * jx] - c[k*3] * dic[k][ix][jx][0])
                                    new_pixel[0] += c[k] * dic[k][ix][jx][0]
                                    new_pixel[1] += c[k] * dic[k][ix][jx][1]
                                    new_pixel[2] += c[k] * dic[k][ix][jx][2]
                                print(self.img[i][j])
                                    #print(dic)
                                    #print(c)
                                print(np.array(new_pixel))
                                print((hsv_to_rgb(np.array(new_pixel)) + 0.5) * 255)
            
        
    def process(self, stepping, alpha = 0.1):
        dic = self.retrieve_dico(stepping)
        h_half = int(self.h / 2)

        
        h, w, _ = self.img.shape
        
        for i in range(h - self.h):
            for j in range(w - self.h):
                patch = self.get_patch(i + h_half, j + h_half)
                if np.any(np.isnan(patch)):
                    y = []
                    x = []
                    print("frame : ", i, j)

                    print(w - h_half - 1)
                    expressed_pixels = np.logical_not(np.isnan(patch))
                    expressed_patch = patch[expressed_pixels]
                    #print(patch)
                    #print(expressed_patch)

                    if expressed_patch.size == 0:
                        continue
                    for p in expressed_patch:
                        y.append(p)
                    for elem in dic:
                        new_x = elem[expressed_pixels]
                        #x.append(np.reshape(new_x, (int(new_x.size / 3), 3)))
                        x.append(new_x)
                        #print(new_x.size)
                    model = linear_model.Lasso(alpha)
                    print(np.array(x).shape)
                    #print(x)
                    #print(len(x))
                    model.fit(np.array(x).reshape(len(y), len(x)), np.array(y))
                    c = model.coef_
                    
                    #dic = np.array(dic)
                    #dic = dic.reshape(int(dic.size / 3), 3)           
                    
                    
                    for ix in range(self.h):
                        for jx in range(self.h):
                            if np.any(expressed_pixels[ix, jx] == False):
                                new_pixel = [0, 0, 0]
                                #print("fixing", ix, jx)
                                #print(len(dic))
                                #print(dic)
                                #print(len(c))
                                for k in range(len(dic)):
                                    #print(y[ix * jx] - c[k*3] * dic[k][ix][jx][0])
                                    new_pixel[0] += c[k] * dic[k][ix][jx][0]
                                    new_pixel[1] += c[k] * dic[k][ix][jx][1]
                                    new_pixel[2] += c[k] * dic[k][ix][jx][2]
                                print(self.img[i][j])
                                #print(dic)
                                #print(c)
                                print(np.array(new_pixel))
                                print((hsv_to_rgb(np.array(new_pixel)) + 0.5) * 255)
                                
                                self.img[ix + i][jx + j] = np.array(new_pixel)
                    
                                    

np.set_printoptions(threshold=np.nan)        
pr = Processor(7, "Lenna.jpg")

pr.noise_img(0.01)
#pr.show_img()
pr.process(5, 0.005)
pr.show_img()