-
Notifications
You must be signed in to change notification settings - Fork 0
/
pipelineTest.lua
executable file
·65 lines (55 loc) · 1.4 KB
/
pipelineTest.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
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
#! /usr/local/torch/install/bin/th
require "cutorch"
require "image"
require "lfs"
require "pipeline"
local function traverse(fn, pngs)
print(fn)
local ll = lfs.attributes(fn)
if ll.mode == "directory" then
local children = { }
for child in lfs.dir(fn) do
if child ~= "." and child ~= ".." then
local cfn = string.format("%s/%s", fn, child)
table.insert(children, cfn)
end
end
for i, cfn in ipairs(children) do
traverse(cfn, pngs)
end
elseif fn:sub(-4) == ".png" then
table.insert(pngs, fn)
end
end
function main(argv)
local input = nil
local output = nil
local files = { }
for i, fn in ipairs(argv) do
traverse(string.format("%s/input", fn), files)
end
print(files)
for i, fn in ipairs(files) do
local img = image.load(fn, "float")
if input == nil then
input = torch.CudaTensor(#files, img:size(1), img:size(2), img:size(3))
end
input[{i,{},{},{}}]:copy(img)
end
input = input:permute(1, 3, 4, 2):contiguous()
local pp = pipeline.new()
for j = 1, 100 do
output = torch.CudaTensor(input:size(1), 227, 227, 3):fill(1.0):contiguous()
local tt = torch.Timer()
pp:run(true, input, output)
print(tt:time())
output = output:permute(1,4,2,3)
print(j)
for i, fn in ipairs(files) do
local ofn, _ = fn:gsub("/input/", "/output/")
image.save(string.format("/tmp/qball/output/%04d.png", j) , output[i])
end
end
end
main({"/tmp/qball"})
collectgarbage()