opticalflow.lua

require 'torch'
torch.setdefaulttensortype('torch.FloatTensor')
require 'xlua'
require 'nnx'
require 'image'
require 'optim'
require 'load_data'
require 'groundtruth_opticalflow'
require 'opticalflow_model'
require 'opticalflow_model_io'
require 'sys'
require 'openmp'
require 'score_opticalflow'

torch.manualSeed(1)

op = xlua.OptionParser('%prog [options]')
-- general
op:option{'-nt', '--num-threads', action='store', dest='nThreads', default=2,
	  help='Number of threads used'}

-- network
op:option{'-nf', '--n-features', action='store', dest='n_features',
          default=10, help='Number of features used for the matching'}
op:option{'-k1s', '--kernel1-size', action='store', dest='kernel1_size',
	  default=5, help='Kernel 1 size'}
op:option{'-k2s', '--kernel2-size', action='store', dest='kernel2_size',
	  default=16, help='Kernel 2 size'}
op:option{'-k3s', '--kernel3-size', action='store', dest='kernel3_size',
	  default=16, help='Kernel 3 size'}
op:option{'-ws', '--window-size', action='store', dest='win_size',
	  default=16, help='Window size maxw (and maxh if no -wsh)'}
op:option{'-wsh', '--window-size-height', action='store', dest='win_size_height',
	  default=nil, help='Window size height (maxh)'}
op:option{'-nl', '-num-layers', action='store', dest='num_layers',
	  default=2, help='Number of layers in the network (1 2 or 3)'}
op:option{'-s2', '--layer-two-size', action='store', dest='layer_two_size', default=8,
	  help='Second layer size, if nl >= 2'}
op:option{'-s2c', '--layer-two-connections', action='store', dest='layer_two_connections',
	  default=4, help='Number of connectons between layers 1 and 2'}
op:option{'-s3', '--layer-three-size', action='store', dest='layer_three_size', default=8,
	  help='Third layer size, if nl >= 3'}
op:option{'-s3c', '--layer-three-connections', action='store', dest='layer_three_connections',
	  default=4, help='Number of connectons between layers 2 and 3'}
op:option{'-l2', '--l2-pooling', action='store_true', dest='l2_pooling', default=false,
	  help='L2 pooling (experimental)'}
op:option{'-ms', '--multiscale', action='store', dest='multiscale', default=0,
	  help='Number of scales used (0 disables multiscale)'}
op:option{'-sf', '--share-filters', action='store_true', dest='share_filters', default=false,
	  help='Share multiscale filters'}
op:option{'-lw', '--load-weights', action='store', dest='load_weights', default = nil,
	  help = 'Load weights from previously trained model'}
op:option{'-mstw', '--multiscale-trainable-weights', action='store_true',
	  dest='ms_trainable_weights', default = false,
	  help='Allow the weights of CascadingAddTable to be trained'}
op:option{'-mssb', '--multiscale-single-beta', action='store_true', default=false,
	  dest='ms_single_beta', help='Single beta per scale in CascadingAddTable'}
op:option{'-oem', '--output-extraction-method', action='store', default='max',
	  dest='output_extraction_method', help="Output extraction method (max | mean)"}

-- learning
op:option{'-n', '--n-train-set', action='store', dest='n_train_set', default=2000,
	  help='Number of patches in the training set'}
op:option{'-m', '--n-test-set', action='store', dest='n_test_set', default=1000,
	  help='Number of patches in the test set'}
op:option{'-mni', '--n-images-test-set', action='store', dest='n_images_test_set', default=2,
	  help='Number of full images to compute epoch score'}
op:option{'-e', '--num-epochs', action='store', dest='n_epochs', default=10,
	  help='Number of epochs'}
op:option{'-r', '--learning-rate', action='store', dest='learning_rate',
          default=5e-3, help='Learning rate'}
op:option{'-lrd', '--learning-rate-decay', action='store', dest='learning_rate_decay',
          default=5e-7, help='Learning rate decay'}
op:option{'-wd', '--weight-decay', action='store', dest='weight_decay',
	  default=0, help='Weight decay'}
op:option{'-rn', '--renew-train-set', action='store_true', dest='renew_train_set',
	  default=false, help='Renew train set at each epoch'}
op:option{'-st', '--soft-targets', action='store', dest='soft_targets',
	  default=nil, help='Targets are gaussians, specify sigma2'}
op:option{'-gtws', '--gt-window-size', action='store', dest='gt_win_size',
	  default=16, help='Groundtruth window size maxw (and maxh)'}

-- input
op:option{'-rd', '--root-directory', action='store', dest='root_directory',
	  default='data/', help='Root dataset directory'}
op:option{'-fi', '--first-image', action='store', dest='first_image', default=0,
	  help='Index of first image used'}
op:option{'-d', '--delta', action='store', dest='delta', default=1,
	  help='Delta between two consecutive frames'}
op:option{'-ni', '--num-input-images', action='store', dest='num_input_images',
	  default=10, help='Number of annotated images used'}
op:option{'-mc', '--motion-correction', action='store', dest='motion_correction',default=false,
	  help='Eliminate panning, tilting and rotation camera movements (mc | sfm)'}
op:option{'-gt', '--grountruth', action='store', dest='groundtruth',
	  default='cross-correlation', help='Groundtruth : cross-correlation | liu | cvlibs'}
op:option{'-nci', '--n-channels-in', action='store', dest='n_channels_in',
	  default=3, help='Number of channels of the input images'}

-- output 
op:option{'-omd', '--output-model-dir', action='store', dest='output_models_dir',
	  default = 'models', help='Output model directory'}

opt=op:parse()
opt.nThreads = tonumber(opt.nThreads)

opt.multiscale = tonumber(opt.multiscale)

opt.n_train_set = tonumber(opt.n_train_set)
opt.n_test_set = tonumber(opt.n_test_set)
opt.n_images_test_set = tonumber(opt.n_images_test_set)
opt.n_epochs = tonumber(opt.n_epochs)
opt.learning_rate = tonumber(opt.learning_rate)
opt.learning_rate_decay = tonumber(opt.learning_rate_decay)
opt.weight_decay = tonumber(opt.weight_decay)

if opt.root_directory:sub(-1) ~= '/' then opt.root_directory = opt.root_directory .. '/' end

openmp.setDefaultNumThreads(opt.nThreads)

local correction = {}
correction.motion_correction = opt.motion_correction
correction.wImg = 640
correction.hImg = 480
correction.bad_image_threshold = 0.2
correction.K = torch.FloatTensor(3,3):zero()
correction.K[1][1] = 293.824707
correction.K[2][2] = 310.435730
correction.K[1][3] = 300.631012
correction.K[2][3] = 251.624924
correction.K[3][3] = 1.0
correction.distP = torch.FloatTensor(5)
correction.distP[1] = -0.379940
correction.distP[2] = 0.212737
correction.distP[3] = 0.003098
correction.distP[4] = 0.000870
correction.distP[5] = -0.069770

local geometry = {}
geometry.wImg = 320
geometry.hImg = 180
geometry.maxwHR = tonumber(opt.win_size) --high res in case of multiscale
geometry.maxhHR = tonumber(opt.win_size) --high res in case of multiscale
if opt.win_size_height then
   geometry.maxhHR = tonumber(opt.win_size_height)
end
geometry.maxwGT = tonumber(opt.gt_win_size)
geometry.maxhGT = tonumber(opt.gt_win_size)
geometry.wKernelGT = 16
geometry.hKernelGT = 16
geometry.layers = {}
if tonumber(opt.num_layers) == 1 then
   geometry.layers[1] = {tonumber(opt.n_channels_in), tonumber(opt.kernel1_size),
			 tonumber(opt.kernel1_size), tonumber(opt.n_features)}
   geometry.wKernel = tonumber(opt.kernel1_size)
   geometry.hKernel = tonumber(opt.kernel1_size)
elseif tonumber(opt.num_layers) == 2 then
   geometry.layers[1] = {tonumber(opt.n_channels_in), tonumber(opt.kernel1_size),
			 tonumber(opt.kernel1_size), tonumber(opt.layer_two_size)}
   geometry.layers[2] = {tonumber(opt.layer_two_connections), tonumber(opt.kernel2_size),
			 tonumber(opt.kernel2_size), tonumber(opt.n_features)}
   geometry.wKernel = tonumber(opt.kernel1_size) + tonumber(opt.kernel2_size) - 1
   geometry.hKernel = tonumber(opt.kernel1_size) + tonumber(opt.kernel2_size) - 1
elseif tonumber(opt.num_layers) == 3 then
   geometry.layers[1] = {tonumber(opt.n_channels_in), tonumber(opt.kernel1_size),
			 tonumber(opt.kernel1_size), tonumber(opt.layer_two_size)}
   geometry.layers[2] = {tonumber(opt.layer_two_connections), tonumber(opt.kernel2_size),
			 tonumber(opt.kernel2_size), tonumber(opt.layer_three_size)}
   geometry.layers[3] = {tonumber(opt.layer_three_connections), tonumber(opt.kernel3_size),
			 tonumber(opt.kernel3_size), tonumber(opt.n_features)}
   geometry.wKernel = tonumber(opt.kernel1_size) + tonumber(opt.kernel2_size) + tonumber(opt.kernel3_size) - 2
   geometry.hKernel = tonumber(opt.kernel1_size) + tonumber(opt.kernel2_size) + tonumber(opt.kernel3_size) - 2
else
   assert(false)
end
geometry.L2Pooling = opt.l2_pooling
if opt.multiscale == 0 then
   geometry.multiscale = false
   geometry.ratios = {1}
   geometry.maxw = geometry.maxwHR
   geometry.maxh = geometry.maxhHR
else
   geometry.multiscale = true
   geometry.ratios = {}
   for i = 1,opt.multiscale do table.insert(geometry.ratios, math.pow(2, i-1)) end
   geometry.maxw = math.ceil(geometry.maxwHR / geometry.ratios[#geometry.ratios])
   geometry.maxh = math.ceil(geometry.maxhHR / geometry.ratios[#geometry.ratios])
end
geometry.wPatch2 = geometry.maxw + geometry.wKernel - 1
geometry.hPatch2 = geometry.maxh + geometry.hKernel - 1
geometry.motion_correction = opt.motion_correction
geometry.share_filters = opt.share_filters
geometry.training_mode = true
geometry.prefilter = false
if geometry.multiscale then
   geometry.cascad_trainable_weights = opt.ms_trainable_weights
   geometry.single_beta = opt.ms_single_beta
end
geometry.output_extraction_method = opt.output_extraction_method

local learning = {}
learning.first_image = tonumber(opt.first_image)
learning.delta = tonumber(opt.delta)
learning.num_images = tonumber(opt.num_input_images)
learning.rate = opt.learning_rate
learning.rate_decay = opt.learning_rate_decay
learning.weight_decay = opt.weight_decay
learning.renew_train_set = opt.renew_train_set
learning.soft_targets = opt.soft_targets ~= nil
learning.st_sigma2 = tonumber(opt.soft_targets)
learning.groundtruth = opt.groundtruth

if learning.groundtruth == 'liu' then
   geometry.hKernelGT = geometry.hKernel
   geometry.wKernelGT = geometry.wKernel
   --geometry.maxhGT = geometry.maxh
   --geometry.maxwGT = geometry.maxw
else
   assert(geometry.maxwGT >= geometry.maxw)
   assert(geometry.maxhGT >= geometry.maxh)
end


local summary = describeModel(geometry, learning)

--local model
if geometry.multiscale then
   model = getModelMultiscale(geometry, false)
else
   model = getModel(geometry, false)
end
if opt.load_weights then
   loadWeightsFrom(model, opt.load_weights)
end
local parameters, gradParameters = model:getParameters()

local criterion
if geometry.output_extraction_method == 'mean' then
   criterion = nn.MSECriterion()
else
   if learning.soft_targets then
      criterion = nn.DistNLLCriterion()
      --criterion.targetIsProbability = true
   else
      criterion = nn.ClassNLLCriterion()
   end
end

print('Loading images...')
print(opt.root_directory)
local raw_data = loadDataOpticalFlow(correction, geometry, learning, opt.root_directory)
print('Generating training set...')
local trainData = generateDataOpticalFlow(correction, geometry, learning,
					  raw_data, opt.n_train_set)
print('Generating test set...')
local testData = generateDataOpticalFlow(correction, geometry, learning,
					 raw_data, opt.n_test_set)
print("a")
local score = score_epoch(geometry, learning, model, criterion, testData,
			  raw_data, opt.n_images_test_set)
saveModel(opt.output_models_dir, 'model_of_', geometry, learning, model, 0, score)
print("b")
config = {learningRate = learning.rate,
	  weightDecay = learning.weight_decay,
	  momentum = 0,
	  learningRateDecay = learning.rate_decay}

for iEpoch = 1,opt.n_epochs do
   print('Epoch ' .. iEpoch .. ' over ' .. opt.n_epochs)
   print(summary)

   local nGood = 0
   local nBad = 0
   local meanErr = 0

   if learning.renew_train_set then
      trainData = generateDataOpticalFlow(correction, geometry, learning,
					  raw_data, opt.n_train_set)
   end
   
   for t = 1,trainData:size() do
      modProgress(t, trainData:size(), 100)

      local input, itarget, target
      if geometry.multiscale then
	 local sample = trainData:getElemFovea(t)
	 input = sample[1][1]
	 model:focus(sample[1][2][2], sample[1][2][1])
	 itarget, target = prepareTarget(geometry, learning, sample[2])
      else
	 local sample = trainData[t]
	 input = prepareInput(geometry, sample[1][1], sample[1][2])
	 itarget, target = prepareTarget(geometry, learning, sample[2])
      end

      
      local feval = function(x)
		       if x ~= parameters then
			  parameters:copy(x)
		       end
		       gradParameters:zero()
		       local output = model:forward(input)
		       if geometry.output_extraction_method == 'mean' then
			  local output_crit = torch.Tensor(2)
			  output_crit[1] = output[1]:squeeze()
			  output_crit[2] = output[2]:squeeze()
			  local target_crit = torch.Tensor(2)
			  target_crit[1], target_crit[2] = x2yx(geometry, target)
			  local err = criterion:forward(output_crit, target_crit)
			  local df_do = criterion:backward(output_crit, target_crit)
			  local df_do2 = {torch.Tensor(1,1):fill(df_do[1]),
					  torch.Tensor(1,1):fill(df_do[2])}
			  model:backward(input, df_do2)
			  meanErr = meanErr + err
			  if (output_crit-target_crit):norm() < 1 then
			     nGood = nGood + 1
			  else
			     nBad = nBad + 1
			  end
		       else
			  local err = criterion:forward(output:squeeze(), target)
			  local df_do = criterion:backward(output:squeeze(), target)
			  model:backward(input, df_do)
		       
			  meanErr = meanErr + err
			  local outputp = processOutput(geometry, output, false)
			  if outputp.index == itarget then
			     nGood = nGood + 1
			  else
			     nBad = nBad + 1
			  end
		       end
		       if sys.isNaN(gradParameters:sum()) then
			  error('stopped in main')
		       end
		       return err, gradParameters
		    end

      optim.sgd(feval, parameters, config)
      if model.cascad then --optim.sgd sucks. Why doesn't it use updateParameters ?????
	 model.cascad:updateNormalizers()
      end
   end
   collectgarbage()

   if geometry.cascad_trainable_weights then
      local p, gp = model.cascad:parameters()
      for i = 1,#p do
	 print(p[i][1])
      end
      print('--')
      for i = 1,#model.cascad.muls_normalizers do
	 print(model.cascad.muls_normalizers[i][1].weight[1])
      end
   end
      
   meanErr = meanErr / (trainData:size())
   print('train: nGood = ' .. nGood .. ' nBad = ' .. nBad .. ' (' .. 100.0*nGood/(nGood+nBad) .. '%) meanErr = ' .. meanErr)

   local score = score_epoch(geometry, learning, model, criterion, testData,
   			     raw_data, opt.n_images_test_set)
   saveModel(opt.output_models_dir, 'model_of_', geometry, learning, model, iEpoch, score)

end