Efficient inpaint with neighbouring pixels

I am implementing a simple algorithm to do in-painting on a "damaged" image. I have a predefined mask that specifies the area which needs to be fixed. My strategy is to start at the border of the masked area and in-paint each pixel with the central mean of its neighboring non-zero pixels, repeating until there's no unknown pixels left.

function R = inPainting(I, mask)
H = [1 2 1; 2 0 2; 1 2 1];
R = I;

n = 1;
[row,col,~] = find(~mask); %Find zeros in mask (area to be inpainted)
unknown = horzcat(row, col)';
while size(unknown,2) > 0
    new_unknown = [];
    new_R = R;
    for u = unknown
        r = u(1);
        c = u(2);
        nb = R(max((r-n), 1):min((r+n), end), max((c-n),1):min((c+n),end));
        nz = nb~=0;
        nzs = sum(nz(:));

        if nzs ~= 0 %We have non-zero neighbouring pixels. In-paint with average.
            new_R(r,c) = sum(nb(:)) / nzs;
        else
            new_unknown = horzcat(new_unknown, u);
        end
    end
    unknown = new_unknown;
    R = new_R;
end

This works well, but it's not very efficient. Is it possible to vectorize such an approach, using mostly matrix operations? Does someone know of a more efficient way to implement this algorithm?

If I understand your problem statement, you are given a mask and you wish to fill in these pixels in this mask with the mean of the neighbourhood pixels that surround each pixel in the mask. Another constraint is that the image is defined such that any pixels that belong to the mask in the same spatial locations are zero in this mask. You are starting from the border of the mask and are propagating information towards the innards of the mask. Given this algorithm, there is unfortunately no way you can do this with standard filtering techniques as the current time step is dependent on the previous time step.

Image filtering mechanisms, like imfilter or conv2 can't work here because of this dependency.

As such, what I can do is help you speed up what is going on inside your loop and hopefully this will give you some speed up overall. I'm going to introduce you to a function called im2col. This is from the image processing toolbox, and given that you can use imfilter, we can use this function.

im2col creates a 2D matrix such that each column is a pixel neighbourhood unrolled into a single vector. How it works is that each pixel neighbourhood in column major order is grabbed, so we get a pixel neighbourhood at the top left corner of the image, then move down one row, and another row and we keep going until we reach the last row. We then move one column over and repeat the same process. For each pixel neighbourhood that we have, it gets unrolled into a single vector, and the output would be a MN x K matrix where you have a neighbourhood size of M x N for each pixel neighbourhood and there are K neighbourhoods.

Therefore, at each iteration of your loop, we can unroll the current inpainted image's pixel neighbourhoods into single vectors, determine which pixel neighborhoods are non-zero and from there, determine how many zero values there are for each of these selected pixel neighbourhood. After, we compute the mean for these non-zero columns disregarding the zero elements. Once we're done, we update the image and move to the next iteration.

What we're going to need to do first is pad the image with a 1 pixel border so that we're able to grab neighbourhoods that extend beyond the borders of the image. You can use padarray, also from the image processing toolbox.

Therefore, we can simply do this:

function R = inPainting(I, mask)
R = double(I); %// For precision

n = 1;
%// Change - column major indices
unknown = find(~mask); %Find zeros in mask (area to be inpainted)  

%// Until we have searched all unknown pixels
while numel(unknown) ~= 0    
    new_R = R;

    %// Change - take image at current iteration and 
    %// create columns of pixel neighbourhoods
    padR = padarray(new_R, [n n], 'replicate');
    cols = im2col(padR, [2*n+1 2*n+1], 'sliding');  

    %// Change - Access the right pixel neighbourhoods
    %// denoted by unknown       
    nb = cols(:,unknown);

    %// Get total sum of each neighbourhood
    nbSum = sum(nb, 1);

    %// Get total number of non-zero elements per pixel neighbourhood
    nzs = sum(nb ~= 0, 1);

    %// Replace the right pixels in the image with the mean
    new_R(unknown(nzs ~= 0)) = nbSum(nzs ~= 0) ./ nzs(nzs ~= 0);

    %// Find new unknown pixels to look at
    unknown = unknown(nzs == 0);

    %// Update image for next iteration
    R = new_R;
end

%// Cast back to the right type
R = cast(R, class(I));