forked from AaronJackson/vrn
-
Notifications
You must be signed in to change notification settings - Fork 0
/
process.lua
executable file
·41 lines (31 loc) · 1.03 KB
/
process.lua
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
#!/usr/bin/env th
local nn = require 'nn'
local image = require 'image'
cmd = torch.CmdLine()
cmd:option('--model', 'model.t7', 'Trained network model')
cmd:option('--input', 'examples/', 'Directory of images to process')
cmd:option('--output', 'output/', 'Location to dump output')
cmd:option('--device', 'gpu', 'GPU or CPU')
opt = cmd:parse(arg or {})
net = torch.load(opt.model)
if opt.device == 'gpu' then
local cunn = require 'cunn'
local cudnn = require 'cudnn'
net = net:cuda()
end
for fname in paths.iterfiles(opt.input .. '/', '*.jpg') do
fname = string.sub(fname, 1, -5)
local img = image.load(opt.input .. '/' .. fname .. '.jpg')
img = image.scale(img, 192, 192)
img = img:view(1,3,192,192):float()
if opt.device == 'gpu' then
img = img:cuda()
end
local output = net:forward(img)
local vol = (output[1]*255):byte()
local out = torch.DiskFile(opt.output .. '/' .. fname .. '.raw', 'w')
out:binary()
out:writeByte(vol:storage())
out:close()
io.write('Processed ' .. fname .. '.\n')
end