infer.lua

--require 'cutorch';
require 'cunn';
require 'cudnn'
require 'image'
require 'nn'
require 'nnx'
require 'optim'
require 'hdf5'
require './util.lua'
require './DataLoader.lua'
require './model.lua'
require 'paths'

local opts = paths.dofile('./opts.lua')
opt = opts.parse(arg)

-- Default user options
options = {
  imgColorPath = opt.imgColorPath,
  imgDepthPath = opt.imgDepthPath,
  modelPath = '../affordance_model/model.t7',
  resultsPath = opt.resultPath,
  outputScale = 1/8,
  imgHeight =  424,
  imgWidth = 512
}

-- Parse user options from command line (i.e. imgColorPath=<image.png> imgDepthPath=<image.png> modelPath=<model.t7> th infer.lua)
for k,v in pairs(options) do options[k] = tonumber(os.getenv(k)) or os.getenv(k) or options[k] end

-- Set RNG seed
math.randomseed(os.time())

-- Load trained model and set to testing (evaluation) mode
-- print('Loading model: '..options.modelPath)
local model = torch.load(options.modelPath)
model:add(cudnn.SpatialSoftMax())
model = model:cuda()
model:evaluate()

-- Define dataset average constants (r,g,b)
local mean = {0.485,0.456,0.406}
local std = {0.229,0.224,0.225}

-- Initialize empty tensors for input and output
local input  = {torch.Tensor(1, 3, options.imgHeight, options.imgWidth),torch.Tensor(1, 3, options.imgHeight, options.imgWidth)}
local results = torch.Tensor(1,3,options.imgHeight*options.outputScale,options.imgWidth*options.outputScale):float()

-- Load and pre-process color image (24-bit RGB PNG)
-- print('Pre-processing color image: '..options.imgColorPath)
local colorImg = image.load(options.imgColorPath)
for c=1,3 do
    colorImg[c]:add(-mean[c])
    colorImg[c]:div(std[c])
end

-- Load and pre-process depth image (16-bit PNG depth in deci-millimeters)
-- print('Pre-processing depth image: '..options.imgDepthPath)
local depth = image.load(options.imgDepthPath) -- When load the image it is divided by 65536
depth = depth*65536/10000 -- same scale with postprocess
depth = depth:clamp(0.0,1.2) -- Depth range of Intel RealSense SR300 (>1.2 ->1.2; <0 ->0)
local depthImg = depth:cat(depth,1):cat(depth,1)
for c=1,3 do
    depthImg[c]:add(-mean[c])
    depthImg[c]:div(std[c])
end

-- Copy images into input tensors and convert them to CUDA
input[1]:copy(colorImg:reshape(1, 3, options.imgHeight, options.imgWidth))
input[2]:copy(depthImg:reshape(1, 3, options.imgHeight, options.imgWidth))
input[1] = input[1]:cuda()
input[2] = input[2]:cuda()

-- -- Compute forward pass
-- print('Computing forward pass...')
local output = model:forward(input)

-- -- Save output test results
-- print('Saving results to: '..options.resultsPath)
results = output:float()
local resultsFile = hdf5.open(options.resultsPath, 'w')
resultsFile:write('results', results:float())
resultsFile:close()
-- print('Done.')