-- usage: DATA_ROOT=/path/to/data/ name=expt1 which_direction=BtoA th test.lua

--

-- code derived from https://github.com/soumith/dcgan.torch and https://github.com/phillipi/pix2pix

require 'image'

require 'nn'

require 'nngraph'

require 'models.architectures'

util = paths.dofile('util/util.lua')

options = require 'options'

opt = options.parse_options('test')

-- initialize torch GPU/CPU mode

if opt.gpu > 0 then

require 'cutorch'

require 'cunn'

cutorch.setDevice(opt.gpu)

print ("GPU Mode")

torch.setdefaulttensortype('torch.CudaTensor')

else

torch.setdefaulttensortype('torch.FloatTensor')

print ("CPU Mode")

end

-- setup visualization

visualizer = require 'util/visualizer'

function TableConcat(t1,t2)

end

-- load data

local data_loader = nil

if opt.align_data > 0 then

require 'data.aligned_data_loader'

data_loader = AlignedDataLoader()

else

require 'data.unaligned_data_loader'

data_loader = UnalignedDataLoader()

end

print( "DataLoader " .. data_loader:name() .. " was created.")

data_loader:Initialize(opt)

if opt.how_many == 'all' then

opt.how_many = data_loader:size()

end

opt.how_many = math.min(opt.how_many, data_loader:size())

-- set batch/instance normalization

set_normalization(opt.norm)

-- load model

opt.continue_train = 1

-- define model

if opt.model == 'cycle_gan' then

require 'models.cycle_gan_model'

model = CycleGANModel()

elseif opt.model == 'one_direction_test' then

require 'models.one_direction_test_model'

model = OneDirectionTestModel()

elseif opt.model == 'pix2pix' then

require 'models.pix2pix_model'

model = Pix2PixModel()

elseif opt.model == 'bigan' then

require 'models.bigan_model'

model = BiGANModel()

elseif opt.model == 'content_gan' then

require 'models.content_gan_model'

model = ContentGANModel()

else

error('Please specify a correct model')

end

model:Initialize(opt)

local pathsA = {} -- paths to images A tested on

local pathsB = {} -- paths to images B tested on

local web_dir = paths.concat(opt.results_dir, opt.name .. '/' .. opt.which_epoch .. '_' .. opt.phase)

paths.mkdir(web_dir)

local image_dir = paths.concat(web_dir, 'images')

paths.mkdir(image_dir)

s1 = opt.fineSize

s2 = opt.fineSize / opt.aspect_ratio

visuals = {}

for n = 1, math.floor(opt.how_many) do

print('processing batch ' .. n)

local cur_dataA, cur_dataB, cur_pathsA, cur_pathsB = data_loader:GetNextBatch()

cur_pathsA = util.basename_batch(cur_pathsA)

cur_pathsB = util.basename_batch(cur_pathsB)

print('pathsA', cur_pathsA)

print('pathsB', cur_PathsB)

model:Forward({real_A=cur_dataA, real_B=cur_dataB}, opt)

visuals = model:GetCurrentVisuals(opt, opt.fineSize)

for i,visual in ipairs(visuals) do

if opt.resize_or_crop == 'scale_width' or opt.resize_or_crop == 'scale_height' then

s1 = nil

s2 = nil

end

visualizer.save_images(visual.img, paths.concat(image_dir, visual.label), {string.gsub(cur_pathsA[1],'.jpg','.png')}, s1, s2)

end

print('Saved images to: ', image_dir)

pathsA = TableConcat(pathsA, cur_pathsA)

pathsB = TableConcat(pathsB, cur_pathsB)

end

labels = {}

for i,visual in ipairs(visuals) do

table.insert(labels, visual.label)

end

-- make webpage

io.output(paths.concat(web_dir, 'index.html'))

io.write('<table style="text-align:center;">')

io.write('<tr><td> Image </td>')

for i = 1, #labels do

io.write('<td>' .. labels[i] .. '</td>')

end

io.write('</tr>')

for n = 1,math.floor(opt.how_many) do

io.write('<tr>')

io.write('<td>' .. tostring(n) .. '</td>')

for j = 1, #labels do

label = labels[j]

io.write('<td><img src="./images/' .. label .. '/' .. string.gsub(pathsA[n],'.jpg','.png') .. '"/></td>')

end

io.write('</tr>')

end

io.write('</table>')