Recreational Image Reconstruction with Decision Trees

reconstruct.py

from skimage.io import imread, imsave
import matplotlib.pyplot as plt
import numpy as np
from sklearn.tree import DecisionTreeRegressor
from sklearn.ensemble import RandomForestRegressor
from skimage.transform import rescale


def blue(r, k):
    cell_size = r / np.sqrt(2)

    def calc_dist(x1, y1, x2, y2):
        return np.sqrt(np.power(x1 - x2, 2) + np.power(y1 - y2, 2))

    x0, y0 = np.random.random(2)
    x = [x0]
    y = [y0]
    I = -np.ones((int(np.ceil(1 / cell_size)), int(np.ceil(1 / cell_size))), dtype=np.int32)
    I[int(y0 / cell_size), int(x0 / cell_size)] = 0
    active = [0]

    while len(active):
        s = active.pop(0)
        sx = x[s]
        sy = y[s]
        i = 0
        while i < k:
            i += 1
            theta = np.random.random() * np.pi * 2
            r2 = np.random.random() * r + r
            x2 = sx + r2 * np.cos(theta)
            y2 = sy + r2 * np.sin(theta)

            if x2 < 0 or y2 < 0 or x2 > 1 or y2 > 1:
                continue

            xi = int(x2 / cell_size)
            yi = int(y2 / cell_size)
            if I[yi, xi] >= 0:
                continue

            too_close = False
            for j in range(max(0, yi - 2), min(I.shape[0] - 1, yi + 2) + 1):
                for l in range(max(0, xi - 2), min(I.shape[1] - 1, xi + 2) + 1):
                    if I[j, l] >= 0 and calc_dist(x2, y2, x[I[j, l]], y[I[j, l]]) <= r:
                        too_close = True
                        break

                if too_close:
                    break

            if not too_close:
                I[yi, xi] = len(x)
                active.append(len(x))
                x.append(x2)
                y.append(y2)
                break

        if i < k:
            active.insert(0, s)

    return x, y


def get_image(img):
    x = np.arange(img.shape[1])
    y = np.arange(img.shape[0])
    xx, yy = np.meshgrid(x, y)
    xy = np.stack([xx.ravel(), yy.ravel()], axis=1)

    beacon_x, beacon_y = blue(.12, 10)
    beacon_x = (np.array(beacon_x) * img.shape[1]).astype(int)
    beacon_y = (np.array(beacon_y) * img.shape[0]).astype(int)
    beacon_xy = np.stack([beacon_x, beacon_y], axis=1)

    D = np.linalg.norm(xy.reshape((-1, 1, 2)) - beacon_xy.reshape(1, -1, 2), axis=2)

    X = D
    Y = img.reshape((-1, 3))

    clf = [DecisionTreeRegressor(max_depth=12)
           for i in range(Y.shape[1])]

    # clf = [RandomForestRegressor(max_depth=7, n_estimators=11, max_samples=.05)
    #        for i in range(Y.shape[1])]

    sample_prob = .05
    res = []
    sample_idx = np.where(np.random.random(size=X.shape[0]) <= sample_prob)[0]
    for i in range(Y.shape[1]):
        clf[i].fit(X[sample_idx], Y[sample_idx][:, i])
        res.append(clf[i].predict(X))

    res = np.stack(res, axis=1).reshape(img.shape)

    return res


def main():
    print("Working on it, a few moments please...")
    img = imread('https://images.unsplash.com/photo-1512641406448-6574e777bec6?fm=jpg&q=60&w=3000')
    scale = 800 / img.shape[1]
    img = (rescale(img, scale, channel_axis=-1) * 255).astype(np.uint8)
    img = img.astype(float) / 255
    res_img = (get_image(img) * 255).astype(np.uint8)
    # imsave('result.png', res_img)
    plt.imshow(res_img)
    plt.axis('off')
    plt.show()


if __name__ == '__main__':
    main()

requirements.txt

numpy
matplotlib
scikit-learn
scikit-image