photographic_mosaic.m

clear;
load('cifar-10-batches-mat/data_batch_1.mat')

%% ========================================================================
%% Step 1: Shrink tiny images

% Resize tiny images from 2D to 4D matrix
all_candidates = data;
all_candidates = reshape(all_candidates', 32, 32, 3, []);
all_candidates = permute(all_candidates, [2 1 3 4]);
all_candidates = im2double(all_candidates);

% Downsample tiny images
all_candidates = all_candidates(1:3:end, 1:3:end, :, :);
% all_candidates = all_candidates(1:5:end, 1:5:end, :, :);

% Size of a patch
tile_sz = size(all_candidates, 1);

%% ========================================================================
%% Step 2: Shave off borders of input image

im_src = imread('sunset.jpg');
src_height = size(im_src, 1) - mod(size(im_src, 1), tile_sz);
src_width = size(im_src, 2) - mod(size(im_src, 2), tile_sz);
%src_height = tile_sz * 50;
%src_width= tile_sz * 50;
im_src = im_src(1:src_height, 1:src_width, :);
im_src = im2double(im_src);

%% ========================================================================
%% Step 3: Unroll input image patches and tiny images

% Unroll tiny images
num_candidates = 10000;
X = zeros([num_candidates, tile_sz * tile_sz * 3], 'like', all_candidates);
for i = 1 : num_candidates
    X(i, :) = reshape(all_candidates(:, :, :, i), 1, []);
end

% Unroll input image patches
Y = zeros([src_height * src_width / tile_sz^2, tile_sz * tile_sz * 3], ...
    'like', im_src);
y_ind = 1;
for j = 1 : tile_sz : src_height
    for i = 1 : tile_sz : src_width
        Y(y_ind, :) = reshape(...
            im_src(j:j+tile_sz-1, i:i+tile_sz-1, :), ...
            1, []);
        y_ind = y_ind + 1;
    end
end

%% ========================================================================
%% Step 4: Run KNN

fprintf('Search neighbors of %i points from %i points\n', size(Y, 1), size(X, 1));
search_inds = knnsearch(X, Y);

%% ========================================================================
%% Step 5: Replace each patch in input image with its neighbor

im_dst = im_src;
y_ind = 1;
for j = 1 : tile_sz : src_height
    for i = 1 : tile_sz : src_width
        im_dst(j:j+tile_sz-1, i:i+tile_sz-1, :) = all_candidates(:, :, :, search_inds(y_ind));            
        y_ind = y_ind + 1;
    end
end

%% ========================================================================
%% Display results

fprintf('%i images loaded\n', size(all_candidates, 4));
figure;
montage(all_candidates(:, :, :, 1:100));
figure; 
subplot(121); imshow(im_src); title('Source');
subplot(122); imshow(im_dst); title('Mosaic');