Torch class to read video files by piping from ffmpeg binary output

ffmpeg.lua

cjson = require 'cjson'

local Ffmpeg = torch.class('Ffmpeg')

function Ffmpeg:__init(path)
    local ffprobeCmd = 'ffprobe -v quiet -print_format json -show_streams ' .. path
    local output = io.popen(ffprobeCmd, 'r'):read('*a')
    j = cjson.decode(output)
    video_st = j['streams'][1]

    self.w, self.h = tonumber(video_st['width']), tonumber(video_st['height'])
    self.depth = 3
    self.frameSize = self.depth * self.w * self.h

    self.proc = torch.PipeFile('ffmpeg -i ' .. path .. ' -f image2pipe -pix_fmt rgb24 -vcodec rawvideo - 2>/dev/null', 'r')
    self.frameIndex = 1
end

function Ffmpeg:frameIter(reshape)
    reshape = reshape or true

    return function()
        local res = self.proc:quiet():readByte(self.frameSize)
        if self.proc:hasError() then
            return nil
        end

        if reshape then
            res = torch.ByteTensor(res, 1, torch.LongStorage({self.h, self.w, self.depth})):transpose(2, 3):transpose(1, 2):float() / 255
        end

        self.frameIndex = self.frameIndex + 1
        return self.frameIndex - 1, res
    end
end

function Ffmpeg:close()
    self.proc:close()
end