-
Notifications
You must be signed in to change notification settings - Fork 31
/
flowFileLoader.lua
38 lines (35 loc) · 1.06 KB
/
flowFileLoader.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
require 'torch'
require 'image'
--[[
Reads a flow field from a binary flow file.
bytes contents
0-3 tag: "PIEH" in ASCII, which in little endian happens to be the float 202021.25
(just a sanity check that floats are represented correctly)
4-7 width as an integer
8-11 height as an integer
12-end data (width*height*2*4 bytes total)
--]]
local function flowFileLoader_load(fileName)
local flowFile = torch.DiskFile(fileName, 'r')
flowFile:binary()
flowFile:readFloat()
local W = flowFile:readInt()
local H = flowFile:readInt()
-- image.warp needs 2xHxW, and also expects (y, x) for some reason...
local flow = torch.Tensor(2, H, W)
local raw_flow = torch.data(flow)
local elems_in_dim = H * W
local storage = flowFile:readFloat(2 * elems_in_dim)
for y=0, H - 1 do
for x=0, W - 1 do
local shift = y * W + x
raw_flow[elems_in_dim + shift] = storage[2 * shift + 1]
raw_flow[shift] = storage[2 * shift + 2]
end
end
flowFile:close()
return flow
end
return {
load = flowFileLoader_load
}