test.lua

require 'cutorch'
require 'nn'
require 'cunn'
require 'image'
require 'nngraph'
require 'paths'
require 'src/InstanceNormalization'
util = paths.dofile('src/util.lua')

local cmd = torch.CmdLine()

cmd:option('-input_dir', 'test_img');
cmd:option('-output_dir', 'test_output', 'Path to save stylized image.')
cmd:option('-load_size', 450)
cmd:option('-gpu', 0, '-1 for CPU mode')
cmd:option('-model_path', './pretrained_model/')
cmd:option('-style', 'Hosoda')

opt = cmd:parse(arg)

if paths.dirp(opt.output_dir) then
else
    paths.mkdir(opt.output_dir)
end

if opt.gpu > -1 then
  cutorch.setDevice(opt.gpu+1)
end

-- Define model
local model = torch.load(paths.concat(opt.model_path .. opt.style .. '_net_G_float.t7'))
model:evaluate()
if opt.gpu > -1 then
  print('GPU mode')
  model:cuda()
else
  print('CPU mode')
  model:float()
end

contentPaths = {}
if opt.input_dir ~= '' then 
  contentPaths = util.extractImageNamesRecursive(opt.input_dir)
else
  print('Please specify the input dierectory')
end

for i=1, #contentPaths do
  local contentPath = contentPaths[i]
  local contentExt = paths.extname(contentPath)
  local contentName = paths.basename(contentPath, contentExt)
  -- load image
	local img = image.load(contentPath, 3, 'float')
  -- resize image, keep aspect ratio
	img = image.scale(img, opt.load_size, 'bilinear')
	sg = img:size()
	local input = nil
  if opt.gpu > -1 then
    input = torch.zeros(1, sg[1], sg[2], sg[3]):cuda()
    input[1] = img
  else
    input = torch.zeros(1, sg[1], sg[2], sg[3]):float()
    input[1] = img
  end
  -- forward
	local out = util.deprocess_batch(model:forward(util.preprocess_batch(input)))
  -- save
	local savePath = paths.concat(opt.output_dir, contentName .. '_' .. opt.style .. '.' .. contentExt)
  image.save(savePath, out[1])
	collectgarbage()
end
print('Done!')