## 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 filter mosaic image into color. Input image must use a
## @var{sensorAlignment} of type 'gbrg', 'grbg', 'bggr', or 'rggb'.
## Note that this function uses bilinear interpolation, which may generate
## slightly different results than the Malvar-He-Cutler algorithm
## implemented in Matlab.
## @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) < 2 || size (I,2) < 2)
  error ('demosaic: Image should have at least 2 rows and 2 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 (strcmp ('str', class (sensorAlignment)))
  error ('demosaic: Invalid value for sensorAlignment')
endif

sensorAlignment = tolower(sensorAlignment);
if (strcmp (sensorAlignment, 'gbrg'))
  ind = [2 1    # Red
         1 1    # Green 1
         2 2    # Green 2
         1 2 ]; # Blue
elseif (strcmp (sensorAlignment, 'grbg'))
  ind = [1 2    # Red
         1 1    # Green 1
         2 2    # Green 2
         2 1 ]; # Blue
elseif (strcmp (sensorAlignment, 'bggr'))
  ind = [2 2    # Red
         2 1    # Green 1
         1 2    # Green 2
         1 1 ]; # Blue
elseif (strcmp (sensorAlignment, 'rggb'))
  ind = [1 1    # Red
         2 1    # Green 1
         1 2    # Green 2
         2 2 ]; # Blue
else
  error ('demosaic: Invalid sensorAlignment "%s"', sensorAlignment)
endif

## Image processing kernels
kern_rb = [ 0.25 0.5  0.25
            0.5  1    0.5
            0.25 0.5  0.25  ];
kern_g  = [ 0    0.25 0
            0.25 1    0.25
            0    0.25 0    ];

# Pad input image by 1 pixel to avoid artifacts at image borders
I = padarray (I, [1 1], 'reflect');

# This also shifts the starting indices of R/G/B channels
ind = 3 - ind;

## Red Channel
channel = cast (zeros (size (I)), class(I));
channel(ind(1,1):2:end,ind(1,2):2:end) = ...
      I(ind(1,1):2:end,ind(1,2):2:end);
channel = imfilter (channel, kern_rb);
RGB(:,:,1) = channel(2:end-1, 2:end-1);

## Green Channel (both pixels)
channel(:) = 0;
channel(ind(2,1):2:end,ind(2,2):2:end) = ...
      I(ind(2,1):2:end,ind(2,2):2:end);
channel(ind(3,1):2:end,ind(3,2):2:end) = ...
      I(ind(3,1):2:end,ind(3,2):2:end);
channel = imfilter (channel, kern_g);
RGB(:,:,2) = channel(2:end-1, 2:end-1);

## Blue Channel
channel(:) = 0;
channel(ind(4,1):2:end,ind(4,2):2:end) = ...
      I(ind(4,1):2:end,ind(4,2):2:end);
channel = imfilter (channel, kern_rb);
RGB(:,:,3) = channel(2:end-1, 2:end-1);

endfunction
%!error <demosaic: Image must be two-dimensional> demosaic (zeros ([2 2 2]), 'rggb')
%!error <demosaic: Image should have at least 2 rows and 2 columns> demosaic ([], 'rggb')
%!error <demosaic: Image should be of class uint8 or uint16 or uint32> demosaic (zeros (2), 'rggb')
%!error <demosaic: Invalid sensorAlignment "foo"> demosaic (uint8 (zeros (2)), 'foo')

%!test # GBRG
%! ref(:,:,1) = uint8 (20 * ones(2));
%! ref(:,:,2) = uint8 (40 * ones(2));
%! ref(:,:,3) = uint8 (60 * ones(2));
%! raw        = uint8 ([ 40 60; 20 40 ]);
%! assert (ref, demosaic (raw, 'gbrg'))

%!test # GRBG
%! ref(:,:,1) = uint8 (20 * ones(2));
%! ref(:,:,2) = uint8 (40 * ones(2));
%! ref(:,:,3) = uint8 (60 * ones(2));
%! raw        = uint8 ([ 40 20; 60 40 ]);
%! assert (ref, demosaic (raw, 'grbg'))

%!test # BGGR
%! ref(:,:,1) = uint8 (20 * ones(2));
%! ref(:,:,2) = uint8 (40 * ones(2));
%! ref(:,:,3) = uint8 (60 * ones(2));
%! raw        = uint8 ([ 60 40; 40 20 ]);
%! assert (ref, demosaic (raw, 'bggr'))

%!test # RGGB
%! ref(:,:,1) = uint8 (20 * ones(2));
%! ref(:,:,2) = uint8 (40 * ones(2));
%! ref(:,:,3) = uint8 (60 * ones(2));
%! raw        = uint8 ([ 20 40; 40 60 ]);
%! assert (ref, demosaic (raw, 'rggb'))

%!test # 16 bit datatype
%! ref(:,:,1) = uint16 (20 * ones(2));
%! ref(:,:,2) = uint16 (40 * ones(2));
%! ref(:,:,3) = uint16 (60 * ones(2));
%! raw        = uint16 ([ 20 40; 40 60 ]);
%! assert (ref, demosaic (raw, 'rggb'))

%!test # 32 bit datatype
%! ref(:,:,1) = uint32 (20 * ones(2));
%! ref(:,:,2) = uint32 (40 * ones(2));
%! ref(:,:,3) = uint32 (60 * ones(2));
%! raw        = uint32 ([ 20 40; 40 60 ]);
%! assert (ref, demosaic (raw, 'rggb'))

%!test # Odd sized matrix
%! ref(:,:,1) = uint8 (20 * ones(3));
%! ref(:,:,2) = uint8 (40 * ones(3));
%! ref(:,:,3) = uint8 (60 * ones(3));
%! raw        = uint8 ([ 20 40 20; 40 60 40; 20 40 20 ]);
%! assert (ref, demosaic (raw, 'rggb'))

%!demo if 1
%!  [x y] = meshgrid(-1:.005:1);
%!  ang = atan2(x, y);
%!  true        = uint8(128 + 128 * sin(ang));
%!  true(:,:,2) = uint8(128 + 128 * sin(ang + (pi * 2 / 3)));
%!  true(:,:,3) = uint8(128 + 128 * sin(ang - (pi * 2 / 3)));
%!  true = true .* (hypot (x, y) < 1);
%!  figure
%!  # subplot(1, 3, 1);
%!  imshow(true);
%!  title('True Image');
%!  bayer = uint8 (zeros (size (true)(1:2)));
%!  bayer(1:2:end, 1:2:end) = true(1:2:end, 1:2:end, 1);
%!  bayer(2:2:end, 1:2:end) = true(2:2:end, 1:2:end, 2);
%!  bayer(1:2:end, 2:2:end) = true(1:2:end, 2:2:end, 2);
%!  bayer(2:2:end, 2:2:end) = true(2:2:end, 2:2:end, 3);
%!  figure # subplot (1, 3, 2);
%!  imshow (bayer);
%!  title ('Bayer Filter Image  (RGGB) showing Moiré Patterns');
%!  rgb = demosaic(bayer, 'rggb');
%!  figure # subplot (1, 3, 3);
%!  imshow (rgb);
%!  title ('Reconstructed Demosaic showing Zippering and False Color Artifacts');
%! endif
