require 'torch'

require 'dok'

require 'image'

require 'xlua'

local help_desc = [[

OpenCV 2.4 wrapper

]]

opencv24 = {}

-- load C lib

require 'libopencv24'

--------------------------------------------------------------------------------

-- Image conversions

--

function opencv24.TH2CVImage(im)

if (im:type() == 'torch.ByteTensor') and ((im:nDimension() == 2) or (im:size(3) == 3)) then

-- TODO: in the unlikely case of a 3-channels 3xh byte image, this fails

return im

else

local im_cv = torch.ByteTensor()

im.libopencv24.TH2CVImage(im, im_cv)

return im_cv

end

end

function opencv24.CV2THImage(im_cv)

local im = torch.Tensor()

im.libopencv24.CV2THImage(im_cv, im)

return im

end

--------------------------------------------------------------------------------

-- Tracking

--

function opencv24.TrackPointsLK(...)

local self = {}

xlua.unpack_class(

self, {...}, 'opencv24.TrackPointsLK', help_desc,

{arg='im1', type='torch.Tensor', help='image 1'},

{arg='im2', type='torch.Tensor', help='image 2'},

{arg='maxPoints', type='number', help='Maximum number of tracked points', default=500},

{arg='pointsQuality',type='number',help='Minimum quality of trackedpoints',default=0.02},

{arg='pointsMinDistance', type='number',

help='Minumum distance between two tracked points', default=3.0},

{arg='featuresBlockSize', type='number',

help='opencv GoodFeaturesToTrack block size', default=20},

{arg='trackerWinSize', type='number',

help='opencv calcOpticalFlowPyrLK block size', default=11},

{arg='trackerMaxLevel', type='number',

help='opencv GoodFeaturesToTrack pyramid depth', default=5},

{arg='useHarris', type='bool', default = false, help = 'Use Harris detector'})

if self.im1:size(1) == 3 then

self.im1 = opencv24.TH2CVImage(self.im1)

end

if self.im2:size(1) == 3 then

self.im2 = opencv24.TH2CVImage(self.im2)

end

local corresps = torch.FloatTensor(self.maxPoints, 4)

libopencv24.TrackPoints(self.im1, self.im2, corresps, self.maxPoints, self.pointsQuality,

self.pointsMinDistance, self.featuresBlockSize,

self.trackerWinSize, self.useHarris)

return corresps

end

function opencv24.TrackPointsFREAK(...)

local self = {}

xlua.unpack_class(

self, {...}, 'opencv24.TrackPointsFREAK', help_desc,

{arg='im1', type='torch.Tensor', help='image 1'},

{arg='im2', type='torch.Tensor', help='image 2'},

{arg='im1Freaks', type='torch.ByteTensor', default=nil,

help='Precomputed FREAKS in image 1 (replaces im1)'},

{arg='im2Freaks', type='torch.ByteTensor', default=nil,

help='Precomputed FREAKS in image 2 (replaces im2)'},

{arg='iFREAK', type='number', default = nil,

help='FREAK object index (cf. opencv24.CreateFREAK)'},

{arg='detectionThres', type='number', default=40,

help='FAST detector threshold'},

{arg='matchingThres', type='number', default=100,

help='FREAK matching threshold (Hamming distance)'})

if not self.im1Freaks then

self.iFREAK = self.iFREAK or opencv24.CreateFREAK()

self.im1Freaks = opencv24.ComputeFREAK(self.im1, self.detectionThres, self.iFREAK)

end

if not self.im2Freaks then

self.iFREAK = self.iFREAK or opencv24.CreateFREAK()

self.im2Freaks = opencv24.ComputeFREAK(self.im2, self.detectionThres, self.iFREAK)

end

local matches = opencv24.MatchFREAK(self.im1Freaks, self.im2Freaks, self.matchingThres)

local tracked = torch.Tensor(matches:size(1), 4)

for i = 1,matches:size(1) do

tracked[{i, {1,2}}]:copy(self.im1Freaks.pos[{matches[i][1], {1,2}}])

tracked[{i, {3,4}}]:copy(self.im2Freaks.pos[{matches[i][2], {1,2}}])

end

return tracked

end

--------------------------------------------------------------------------------

-- Dense Optical Flow

--

function opencv24.DenseOpticalFlow(...)

local self = {}

xlua.unpack_class(

self, {...}, 'opencv24.DenseOpticalFlow', help_desc,

{arg='im1', type='torch.Tensor', help='image 1'},

{arg='im2', type='torch.Tensor', help='image 2'},

{arg='mode', type='string', default='farnebach',

help='mode = farnebach | block'},

{arg='pyr_scale', type='number', default=0.5,

help='Ratio between 2 successive pyramid scales (farnebach)'},

{arg='levels', type='number', default=5, help='Pyramid depth (farnebach)'},

{arg='winsize', type='number', default=11,

help='Window size'},

{arg='iterations', type='number', default=20,

help='Number of iteration at each level (farnebach)'},

{arg='poly_n', type='number', default=5,

help='Size of the pixel neighborhood used to find polynomial expansion in each pixel (farnebach)'},

{arg='poly_sigma', type='number', default=1.1,

help='Size of the pixel neighborhood used to find polynomial expansion in each pixel. For poly_n=5 , you can set poly_sigma=1.1 . For poly_n=7 , a good value would be poly_sigma=1.5 (farnebach)'},

{arg='shiftsize', type='number', default=1,

help='Block coordinate increments (block)'},

{arg='maxrange', type='number', default=11,

help='Size of the scanned neighborhood in pixels around the block (block)'},

{arg='flowguess', type='torch.Tensor', default=nil,

help="Initial guess for the initialization of the flow (doesn't seem to work too well with farnebach)"}

)

if self.im1:nDimension() == 3 then

self.im1 = image.rgb2y(self.im1)[1]

end

if self.im2:nDimension() == 3 then

self.im2 = image.rgb2y(self.im2)[1]

end

local im1_cv = opencv24.TH2CVImage(self.im1)

local im2_cv = opencv24.TH2CVImage(self.im2)

winr = image.display{image={im1_cv, im2_cv}, win=winr}

local flow = torch.FloatTensor(2, im1_cv:size(1), im1_cv:size(2))

if self.flowguess ~= nil then

flow:copy(self.flowguess)

else

flow:zero()

end

if self.mode == 'farnebach' then

flow.libopencv24.DenseOpticalFlowFarnebach(im1_cv, im2_cv, flow,

self.pyr_scale, self.levels,

self.winsize, self.iterations,

self.poly_n, self.poly_sigma,

self.flowguess ~= nil)

elseif self.mode == 'block' then

local h2 = math.floor(im1_cv:size(1) - self.winsize

+self.shiftsize) / self.shiftsize

local w2 = math.floor(im1_cv:size(2) - self.winsize

+ self.shiftsize) / self.shiftsize

local dh = math.floor((im1_cv:size(1)-h2)/2)

local dw = math.floor((im1_cv:size(2)-w2)/2)

local flow2 = flow:narrow(2, dh, h2):narrow(3, dw, w2)

libopencv24.DenseOpticalFlowBlockMatching(im1_cv, im2_cv, flow2,

self.winsize, self.shiftsize,

self.maxrange,

self.flowguess ~= nil)

end

return flow:real()

end

--------------------------------------------------------------------------------

-- CornerHarris

--

function opencv24.CornerHarris(...)

local self = {}

xlua.unpack_class(

self, {...}, 'opencv24.CornerHarris', help_desc,

{arg='im', type='torch.Tensor', help='image'},

{arg='blocksize', type='number', default=9,

help='Neighborhood size (See. opencv cornerEigenValsAndVecs())'},

{arg='ksize', type='number', default=3,

help='Aperture parameter for the Sobel() operator.'},

{arg='k', type='number', default=0.04,

help='Harris detector free parameter.'})

local out = torch.Tensor(self.im:size(2), self.im:size(3))

local im_cv = opencv24.TH2CVImage(self.im)

out.libopencv24.CornerHarris(im_cv, out,

self.blocksize,self.ksize,self.k)

return out

end

--------------------------------------------------------------------------------

-- CornerHarris

--

function opencv24.DetectExtract(...)

local self = {}

xlua.unpack_class(

self, {...}, 'opencv24.DetectExtract', help_desc,

{arg='im', type='torch.Tensor', help='image'},

{arg='mask', type='torch.Tensor',

help='mask areas where not to compute.',

default=torch.Tensor()},

{arg='detectorType', type="string",

help="GFTT etc.",default="FAST"},

{arg='extractorType', type="string",

help="FREAK etc.",default="SURF"},

{arg='maxPoints', type='number',

help='Maximum number of tracked points', default=0},

{arg='pointsQuality',type='number',

help='Minimum quality of trackedpoints',default=0.02},

{arg='pointsMinDistance', type='number',

help='Minumum distance between two tracked points', default=3.0},

{arg='blocksize', type='number', default=9,

help='Neighborhood size (See. opencv cornerEigenValsAndVecs())'},

{arg='useHarris', type='bool', default = false,

help = 'Use Harris detector'},

{arg='k', type='number', default=0.04,

help='Harris detector free parameter.'})

local positions = torch.Tensor(self.maxPoints, 2)

local feat = torch.Tensor(self.maxPoints, 128)

local im_cv = opencv24.TH2CVImage(self.im)

feat.libopencv24.DetectExtract(im_cv, self.mask, positions, feat,

self.detectorType, self.extractorType,

self.maxPoints)

return positions,feat

end

--------------------------------------------------------------------------------

-- FREAK

--

function opencv24.CreateFREAK(orientedNormalization, scaleNormalization,

patternSize, nOctave, trainedPairs)

orientedNormalization = orientedNormalization or true

scaleNormalization = scaleNormalization or true

patternSize = patternSize or 22

nOctave = nOctave or 4

trainedPairs = trainedPairs or torch.IntTensor()

local iFREAK = libopencv24.CreateFREAK(orientedNormalization, scaleNormalization,

patternSize, nOctave, trainedPairs)

opencv24.ComputeFREAK(image.lena(), 40, iFREAK)

return iFREAK

end

function opencv24.DeleteFREAK(iFREAK)

libopencv24.DeleteFREAK(iFREAK)

end

function opencv24.ComputeFREAKfromKeyPoints(im, kp, iFREAK)

local freaks = {}

freaks.descs = torch.ByteTensor()

freaks.pos = kp

libopencv24.ComputeFREAKfromKeyPoints(opencv24.TH2CVImage(im),

freaks.descs, freaks.pos,

detection_threshold, iFREAK);

return freaks

end

function opencv24.ComputeFREAK(im, detection_threshold, iFREAK)

local freaks = {}

freaks.descs = torch.ByteTensor()

freaks.pos = torch.FloatTensor()

libopencv24.ComputeFREAK(opencv24.TH2CVImage(im), freaks.descs, freaks.pos,

detection_threshold, iFREAK);

return freaks

end

function opencv24.DrawFREAK(im, freaks, r, g, b)

r = r or 1

g = g or 0

b = b or 0

require 'draw'

for i = 1,freaks.pos:size(1) do

local x = freaks.pos[i][1]

local y = freaks.pos[i][2]

local rad = freaks.pos[i][3]

local ang = freaks.pos[i][4]

draw.circle(im, x, y, rad, r, g, b)

draw.line(im, x, y, x+rad*math.cos(ang), y+rad*math.sin(ang), r, g, b)

end

end

function opencv24.MatchFREAK(freaks1, freaks2, threshold)

local matches = torch.LongTensor()

local nMatches = libopencv24.MatchFREAK(freaks1.descs, freaks2.descs, matches, threshold)

if nMatches == 0 then

return torch.Tensor()

else

matches:add(1) -- one-based lua

return matches:narrow(1,1,nMatches)

end

end

function opencv24.TrainFREAK(images, iFREAK, keypoints_threshold, correlation_threshold)

if #images < 1 then

error("opencv24.TrainFREAK : there must be at least one image...")

end

local images_cv = {}

for i = 1,#images do

images_cv[i] = opencv24.TH2CVImage(images[i])

end

local pairs = torch.IntTensor()

libopencv24.TrainFREAK(images_cv, pairs, iFREAK, keypoints_threshold,

correlation_threshold)

return pairs

end

function opencv24.ComputeFAST(im, detection_threshold)

local pos = torch.FloatTensor()

libopencv24.ComputeFAST(opencv24.TH2CVImage(im), pos,

detection_threshold);

return pos

end

function opencv24.DrawFAST(im, pos, r, g, b)

r = r or 1

g = g or 0

b = b or 0

require 'draw'

for i = 1,pos:size(1) do

local x = pos[i][1]

local y = pos[i][2]

local rad = pos[i][3]

local ang = pos[i][4]

draw.circle(im, x, y, rad, r, g, b)

draw.line(im, x, y,

x+rad*math.cos(ang), y+rad*math.sin(ang),

r, g, b)

end

end

function opencv24.DrawPos(im, pos, size, r, g, b)

r = r or 1

g = g or 0

b = b or 0

require 'draw'

for i = 1,pos:size(1) do

local x = pos[i][1]

local y = pos[i][2]

draw.circle(im, x, y, size , r, g, b)

end

end

function opencv24.Version()

libopencv24.Version()

end

--------------------------------------------------------------------------------

-- Test/Example

--

function opencv24.ImageConversion_testme()

local im = image.scale(image.lena(), 123, 242)

local eps = 1/(255*2)

local im_cv = opencv24.TH2CVImage(im)

local im2 = opencv24.CV2THImage(im_cv)

local diff = (im-im2):abs():gt(eps):sum()

assert(diff == 0)

local im3 = opencv24.CV2THImage(opencv24.TH2CVImage(im_cv))

diff = (im-im3):abs():gt(eps):sum()

assert(diff == 0)

im = im[1]

im_cv = opencv24.TH2CVImage(im)

im2 = opencv24.CV2THImage(im_cv)

diff = (im-im2):abs():gt(eps):sum()

assert(diff == 0)

im3 = opencv24.CV2THImage(opencv24.TH2CVImage(im_cv))

diff = (im-im3):abs():gt(eps):sum()

assert(diff == 0)

end

function opencv24.TrackPointsLK_testme()

require 'draw'

local im = image.lena()

local im2 = image.rotate(im, 0.1)

local timer = torch.Timer()

local corresps = opencv24.TrackPointsLK{im1=im, im2=im2, maxPoints = 100}

print("Track : ", timer:time().real)

local disp = torch.Tensor(3, im:size(2), im:size(3)*2)

disp[{{},{},{1,im:size(3)}}]:copy(im)

disp[{{},{},{im:size(3)+1,im:size(3)*2}}]:copy(im2)

for i = 1,corresps:size(1) do

local x1 = corresps[i][1]

local y1 = corresps[i][2]

local x2 = corresps[i][3]

local y2 = corresps[i][4]

draw.line(disp, x1, y1, x2+im:size(3), y2, 0, 0, 1)

end

image.display{image=disp, zoom=1}

end

function opencv24.TrackPointsFREAK_testme()

require 'draw'

local im = image.lena()

local im2 = image.rotate(im, 0.1)

local corresps = opencv24.TrackPointsFREAK{im1=im, im2=im2}

local disp = torch.Tensor(3, im:size(2), im:size(3)*2)

disp[{{},{},{1,im:size(3)}}]:copy(im)

disp[{{},{},{im:size(3)+1,im:size(3)*2}}]:copy(im2)

for i = 1,corresps:size(1) do

local x1 = corresps[i][1]

local y1 = corresps[i][2]

local x2 = corresps[i][3]

local y2 = corresps[i][4]

draw.line(disp, x1, y1, x2+im:size(3), y2, 0, 0, 1)

end

image.display{image=disp, zoom=1}

end

function opencv24.FREAK_testme()

local size = 22

local iFREAK = opencv24.CreateFREAK(true, true, size, 4)

local im = image.lena()

local im2 = image.rotate(im, 0.1)

--local im = image.load('/home/myrhev/NYU/depth-estimation/radial/data/no-risk/part1/images/000000001.jpg')

--local im2 = image.load('/home/myrhev/NYU/depth-estimation/radial/data/no-risk/part1/images/000000002.jpg')

local trainedPairs = opencv24.TrainFREAK({im}, iFREAK, 40, 0.7)

local iFREAK = opencv24.CreateFREAK(true, true, size, 4, trainedPairs)

local timer = torch.Timer()

local freaks = opencv24.ComputeFREAK(im, 40, iFREAK)

print("Freak 1 : ", timer:time().real)

local freaks2 = opencv24.ComputeFREAK(im2, 40, iFREAK)

print("Freak 2 : ", timer:time().real)

local matches = opencv24.MatchFREAK(freaks, freaks2, 100)

print("Matches : ", timer:time().real)

local imb = im:clone()

local im2b = im2:clone()

opencv24.DrawFREAK(imb, freaks)

opencv24.DrawFREAK(im2b, freaks2)

local disp = torch.Tensor(3, imb:size(2), imb:size(3)*2)

disp[{{},{},{1,imb:size(3)}}]:copy(imb)

disp[{{},{},{imb:size(3)+1,imb:size(3)*2}}]:copy(im2b)

for i = 1,matches:size(1) do

local x1 = freaks.pos[matches[i][1]][1]

local y1 = freaks.pos[matches[i][1]][2]

local x2 = freaks2.pos[matches[i][2]][1]

local y2 = freaks2.pos[matches[i][2]][2]

draw.line(disp, x1, y1, x2+imb:size(3), y2, 0, 0, 1)

end

image.display{image=disp, zoom=1}

opencv24.DeleteFREAK(iFREAK)

end

function opencv24.FAST_testme()

local im = image.lena()

local timer = torch.Timer()

local t0 = timer:time().real

local pos = opencv24.ComputeFAST(im,40)

local t1 = timer:time().real

print("FAST : ", t1-t0)

opencv24.DrawFAST(im,pos)

image.display{image=im, zoom=1}

end

function opencv24.CornerHarris_testme()

local im = image.lena()

local timer = torch.Timer()

local t0 = timer:time().real

local cmap = opencv24.CornerHarris(im)

local t1 = timer:time().real

print("CornerHarris : ", t1-t0)

image.display{image=cmap, zoom=1}

return cmap

end

function opencv24.DetectExtract_testme(dtype,etype)

if not dtype then

dtype = "FAST"

end

if not etype then

etype = "SURF"

end

-- require 'draw'

local im = image.lena()

-- make a simple mask

local m = torch.Tensor(im:size(2),im:size(3)):fill(0)

m:narrow(1,150,250):narrow(2,100,100):fill(1)

local timer = torch.Timer()

local pos, feat =

opencv24.DetectExtract{im=im[1], maxPoints = 100,

mask=m,

detectorType=dtype,

extractorType=etype}

local d1 = timer:time().real

print("DetectExtract (mask) : ", d1)

local pos2, feat2 =

opencv24.DetectExtract{im=im[1], maxPoints = 100,

detectorType=dtype,

extractorType=etype}

local d2 = timer:time().real

print("DetectExtract (no mask): ", d2-d1)

-- remove red channel from areas outside mask

im[1]:cmul(m)

-- masked points are red

for i = 1,pos:size(1) do

im[1][pos[i][2]][pos[i][1]] = 1

end

-- non-masked points are green

for i = 1,pos2:size(1) do

im[2][pos2[i][2]][pos2[i][1]] = 1

end

image.display{image=im, zoom=1}

return pos, feat

end