segment_custom_angle.lua

require 'image'
require 'lfs'
require 'cunn'
require 'nngraph'

function segment(model, flow_mag_ang_file, minmax_file, output_file)

    local file = io.open(minmax_file)
    local minmaxes = {}
    local ind = 1;
    if file then
        for line in file:lines() do
            local mag_min, mag_max, ang_min, ang_max = unpack(line:split(" "))
            minmaxes[ind]  = {mag_min, mag_max, ang_min, ang_max}
            ind = ind + 1
        end
    else
        print(file)
        print('File not found!!!!!!!!')
    end
    io.close(file)

    -- README.txt:
    --  When evaluating on other datasets the inputs have to be resize so that the smallest dimension equals to 232.

    -- TODO: this fails when the video is in portrait orientation (!!)
    local resized_width  = 424
    local resized_height = 232
    local batch = torch.Tensor(1, 2, resized_height, resized_width);

    -- Read out the frame number from string
    -- NOTE: Lua indexing starts at 1, so this needs to be incremented by 1 (!!)
    local frameNum = tonumber(flow_mag_ang_file:match('(%d+).jpg'))

    -- min/max values
    local mm = minmaxes[frameNum+1]

    local f = io.open(flow_mag_ang_file, "r")
    if f~=nil then
        io.close(f)
    else
        print("unable to read mag-angle file from disk")
        return false
    end

    -- Load flow magnitude-angle image from disk
    local flowMagAngle = image.load(flow_mag_ang_file)

    flowMag   = flowMagAngle[{{1}, {}, {}}]
    flowAngle = flowMagAngle[{{2}, {}, {}}]

    -- Some image information
    local height, width = flowMag:size(2), flowMag:size(3)

    local flowFrame = torch.cat(flowAngle, flowMag, 1)
    flowFrame = image.scale(flowFrame, resized_width, resized_height, 'simple');

    print(mm[1], mm[2], mm[3], mm[4])
    -- NOTE: flowFrame has FIRST angle then magnitude, so denormalization requires attention with indices!
    flowFrame[{{1}, {}, {}}] = flowFrame[{{1}, {}, {}}] * (mm[4] - mm[3]) + mm[3]
    flowFrame[{{2}, {}, {}}] = flowFrame[{{2}, {}, {}}] * (mm[2] - mm[1]) + mm[1]
    --flowFrame[{{1}, {}, {}}]:cmul(flowFrame[{{2}, {}, {}}]:gt(1):double())
    flowFrame[{{2}, {}, {}}]:div(math.sqrt(math.pow(854, 2) + math.pow(480, 2)) / 6)

    batch[1] = flowFrame;
    batch = batch:float():cuda()
    local outputs = model:forward(batch)

    local preds
    preds = torch.Tensor(1, height, width)
    local pred = outputs[1];
    pred = nn.utils.recursiveType(pred, 'torch.DoubleTensor')
    pred = image.scale(pred, width, height)

    --image.save(output_file, pred)
    --image.save(output_file, pred);

    -- Apply objectivness score
    local predRaw = pred
    --objectness:div(objectness:max()):add(0.5)
    --pred:cmul(objectness)
    pred:clamp(0, 1)
    preds[1] = pred

    -- Save results
    -- local resultPath = string.gsub(rgbPath, 'JPEGImages', 'Results/Segmentations');
    -- local resultPath = string.gsub(resultPath, '480p', '480p/' .. setting);
    -- local resultPath = string.gsub(resultPath, 'jpg', 'png');
    -- local resultDir = string.gsub(resultPath, '%d+.png', '');
    -- if not path.exists(resultDir) then
    --     os.execute("mkdir " .. resultDir)
    -- end

    -- local resultPathRaw = string.gsub(resultPath, '(%d+).png', 'raw_%1.png');
    image.save(output_file, predRaw);
end