require 'nn'

require 'cunn'

require 'cudnn'

require 'optim'

require 'xlua'

require 'torch'

require 'hdf5'

opt_string = [[

--model (default "logs/model.net") torch model file path

--h5_list_path (default "data/volume_data0.h5") h5 data path

--gpu_index (default 0) GPU index

--output_file (default "feat.txt") Ouput filename

--output_name_file (default "names.txt")

--file_label_file (default "train_file_label.txt")

--partial_data use partial data as input

--classes_file (default "") limit classes

]]

opt = lapp(opt_string)

-- print help or chosen options

if opt.help == true then

print('Usage: th train.lua')

print('Options:')

print(opt_string)

os.exit()

else

print(opt)

end

nchannels = 1

if opt.partial_data then nchannels = 2 end

print('#channels = ' .. nchannels)

function getClassesSet(file)

assert(paths.filep(file))

classes = {}

for line in io.lines(file) do

--split line by whitespace

parts = {}

for p in line:gmatch("%w+") do table.insert(parts, p) end

classes[tonumber(parts[1])+1] = true

end

return classes

end

-- find set of classes

class_set = getClassesSet(opt.classes_file)

-- set gpu

cutorch.setDevice(opt.gpu_index+1)

-- output file

outfile = assert(io.open(opt.output_file, "w"))

outnamefile = assert(io.open(opt.output_name_file, "w"))

-- specify which layer's output we would use as feature

OUTPUT_LAYER_INDEX = 33

print('Loading model...')

model = torch.load(opt.model):cuda()

model:evaluate()

print(model)

function getLinesFromFile(file)

assert(paths.filep(file))

lines = {}

for line in io.lines(file) do

lines[#lines + 1] = line

end

return lines

end

-- load h5 file data into memory

function loadDataFile(file_name)

local current_file = hdf5.open(file_name,'r')

local current_data = current_file:read('data'):all():float()

local current_label = torch.squeeze(current_file:read('label'):all():add(1))

current_file:close()

return current_data, current_label

end

constant_ones = torch.ones(1,1,30,30,30):float()

filenames = getLinesFromFile(opt.h5_list_path)

print('#filenames = ' .. #filenames)

instancenames = getLinesFromFile(opt.file_label_file)

print('first instance:')

print(instancenames[1])

fidx = 1

local count = 0

local total = 0

for fn = 1,#filenames do

print('Loading data...')

local current_data, current_label = loadDataFile(filenames[fn])

print(#current_data)

print('Starting testing...')

for t = 1,current_data:size(1) do

local inputs = current_data[t][{{1,nchannels},{},{},{}}]:reshape(1,nchannels,30,30,30)

--print(inputs:size())

--print(inputs:sum())

if not opt.partial_data then inputs = torch.cat(inputs, constant_ones, 2) end -- all voxels are known

--print(inputs:size())

--print(inputs:sum())

local target = current_label[t]

if class_set[target] then

total = total + 1

local outputs = model:forward(inputs:cuda())

val, idx = torch.max(outputs:double(), 1)

if idx[1] == target then count = count + 1 end

--print('pred ' .. idx[1] .. ', target ' .. target)

--print(instancenames[fidx])

--io.read()

--print(outputs)

feat = model:get(OUTPUT_LAYER_INDEX).output:double()

splitter = ','

for i=1,feat:size(1) do

outfile:write(string.format("%.6f", feat[i]))

if i < feat:size(1) then

outfile:write(splitter)

end

end

outfile:write('\n')

outnamefile:write(string.format("%s\n", instancenames[fidx]))

--else

-- print('warning: ignoring target ' .. target .. ' not in class set')

-- io.read()

end

fidx = fidx + 1

end

print('\tcur acc = ' .. count/total .. '\t(' .. count .. '/' .. total .. ')')

end