-- this program takes an original image, such as a photo,
-- and a generated image, such as generated by jcjohnson/fast-neural-style
-- and copies the original colors to the generated image
-- like when using the original_colors param in jcjohnson/neural-style
--
-- by hannu töyrylä @htoyryla 30 oct 2016
--

require 'torch'
require 'image'

local cmd = torch.CmdLine()

cmd:option('-generated_image', '',
           'Image to copy colors to')
cmd:option('-original_image', '',
           'Image to copy colors from')
cmd:option('-output_image', 'out.png')


local function main(params)
  local generated_image = image.load(params.generated_image, 3)
  local oimg = image.load(params.original_image, 3)
  local original_image = image.scale(oimg, generated_image:size(3), generated_image:size(2))  
  disp = original_colors(original_image, generated_image)
  image.save(params.output_image, disp)  
end

function original_colors(content, generated)
  local generated_y = image.rgb2yuv(generated)[{{1, 1}}]
  local content_uv = image.rgb2yuv(content)[{{2, 3}}]
  return image.yuv2rgb(torch.cat(generated_y, content_uv, 1))
end

local params = cmd:parse(arg)
main(params)