## Copyright (C) 2017   Timothy Parys
##
## This program is free software: you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation, either version 3 of the License, or
## (at your option) any later version.
##
## This program is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with this program.  If not, see <http://www.gnu.org/licenses/>.

## -*- texinfo -*-
## @deftypefn {Function File} {@var{RGB} =} demosaic (@var{I}, @var{sensorAlignment})
## Convert raw bayer image into color. Input image must use a
## @var{sensorAlignment} of type 'gbrg', 'grbg', 'bggr', or 'rggb'.
## @end deftypefn

## Author: Timothy Parys <tparys@gmail.com>
## Created: April 2017
## Version: 0.1

function RGB = demosaic(I,sensorAlignment)

## Check inputs
  if (nargin != 2)
    print_usage ();
  elseif (size(size(I))(2) > 2)
    error('demosaic: Image must be two-dimensional')
  elseif (size(I, 1) < 5 || size (I,2) < 5)
    error ('demosaic: Image should have at least 5 rows and 5 columns');
  elseif (!strcmp(class(I),'uint8') && !strcmp(class(I),'uint16') && !strcmp(class(I),'uint32'))
    error ('demosaic: Image should be of class uint8 or uint16 or uint32');    
  elseif (any(mod(size(I), 2) != [0 0]))
    error('demosaic: Image dimensions must be even')
  elseif (strcmp('str', class(sensorAlignment)))
    error('demosaic: Invalid value for sensorAlignment')
  elseif (strcmp(sensorAlignment,'gbrg'))
    pat        = [0 0; 1 0];
    pat(:,:,2) = [1 0; 0 1];
    pat(:,:,3) = [0 1; 0 0];
  elseif (strcmp(sensorAlignment, 'grbg'))
    pat        = [0 1; 0 0];
    pat(:,:,2) = [1 0; 0 1];
    pat(:,:,3) = [0 0; 1 0];
  elseif (strcmp(sensorAlignment, 'bggr'))
    pat        = [0 0; 0 1];
    pat(:,:,2) = [0 1; 1 0];
    pat(:,:,3) = [1 0; 0 0];
  elseif (strcmp(sensorAlignment, 'rggb'))
    pat        = [1 0; 0 0];
    pat(:,:,2) = [0 1; 1 0];
    pat(:,:,3) = [0 0; 0 1];
  else
    error('demosaic: Invalid value "%s" for sensorAlignment', sensorAlignment)
  end

## Image processing kernels
kern        = [ .25 .5 .25; .5 1 .5; .25 .5 .25 ];
kern(:,:,2) = [ 0 .25 0; .25 1 .25; 0 .25 0 ];
kern(:,:,3) = [ .25 .5 .25; .5 1 .5; .25 .5 .25 ];

## Loop for red, green, blue channels
s = size(I) / 2;
for i = 1:3
    
    # Extract all pixels in this channel
    channel = I .* repmat(pat(:,:,i), s);
    
    # Create border to avoid edge artifacts
    channel = padarray(channel, [1 1], 'reflect');
    
    # Run the kernel
    channel = imfilter(channel, kern(:,:,i)); 
    
    # Copy back everything but the border
    RGB(:,:,i) = channel(2:end-1, 2:end-1);
    
end
