--[[----------------------------------------------------------------------------

Interface for manipulating masks stored in RLE format.

For an overview of RLE please see http://mscoco.org/dataset/#download.
Additionally, more detailed information can be found in the Matlab MaskApi.m:
 https://github.com/pdollar/coco/blob/master/MatlabAPI/MaskApi.m

The following API functions are defined:
 encode - Encode binary masks using RLE.
 decode - Decode binary masks encoded via RLE.
 merge - Compute union or intersection of encoded masks.
 iou - Compute intersection over union between masks.
 nms - Compute non-maximum suppression between ordered masks.
 area - Compute area of encoded masks.
 toBbox - Get bounding boxes surrounding encoded masks.
 frBbox - Convert bounding boxes to encoded masks.
 frPoly - Convert polygon to encoded mask.
 drawCirc - Draw circle into image (alters input).
 drawLine - Draw line into image (alters input).
 drawMasks - Draw masks into image (alters input).

Usage:
 Rs = MaskApi.encode( masks )
 masks = MaskApi.decode( Rs )
 R = MaskApi.merge( Rs, [intersect=false] )
 o = MaskApi.iou( dt, gt, [iscrowd=false] )
 keep = MaskApi.nms( dt, thr )
 a = MaskApi.area( Rs )
 bbs = MaskApi.toBbox( Rs )
 Rs = MaskApi.frBbox( bbs, h, w )
 R = MaskApi.frPoly( poly, h, w )
 MaskApi.drawCirc( img, x, y, rad, clr )
 MaskApi.drawLine( img, x0, y0, x1, y1, rad, clr )
 MaskApi.drawMasks( img, masks, [maxn=n], [alpha=.4], [clrs] )
For detailed usage information please see cocoDemo.lua.

In the API the following formats are used:
 R,Rs - [table] Run-length encoding of binary mask(s)
 masks - [nxhxw] Binary mask(s)
 bbs - [nx4] Bounding box(es) stored as [x y w h]
 poly - Polygon stored as {[x1 y1 x2 y2...],[x1 y1 ...],...}
 dt,gt - May be either bounding boxes or encoded masks
Both poly and bbs are 0-indexed (bbox=[0 0 1 1] encloses first pixel).

Common Objects in COntext (COCO) Toolbox. version 3.0
Data, paper, and tutorials available at: http://mscoco.org/
Code written by Pedro O. Pinheiro and Piotr Dollar, 2016.
Licensed under the Simplified BSD License [see coco/license.txt]

------------------------------------------------------------------------------]]

local ffi = require 'ffi'
local coco = require 'coco.env'

coco.MaskApi = {}
local MaskApi = coco.MaskApi

coco.libmaskapi = ffi.load(package.searchpath('libmaskapi',package.cpath))
local libmaskapi = coco.libmaskapi

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

MaskApi.encode = function( masks )
 local n, h, w = masks:size(1), masks:size(2), masks:size(3)
 masks = masks:type('torch.ByteTensor'):transpose(2,3)
 local data = masks:contiguous():data()
 local Qs = MaskApi._rlesInit(n)
 libmaskapi.rleEncode(Qs[0],data,h,w,n)
 return MaskApi._rlesToLua(Qs,n)
end

MaskApi.decode = function( Rs )
 local Qs, n, h, w = MaskApi._rlesFrLua(Rs)
 local masks = torch.ByteTensor(n,w,h):zero():contiguous()
 libmaskapi.rleDecode(Qs,masks:data(),n)
 MaskApi._rlesFree(Qs,n)
 return masks:transpose(2,3)
end

MaskApi.merge = function( Rs, intersect )
 intersect = intersect or 0
 local Qs, n, h, w = MaskApi._rlesFrLua(Rs)
 local Q = MaskApi._rlesInit(1)
 libmaskapi.rleMerge(Qs,Q,n,intersect)
 MaskApi._rlesFree(Qs,n)
 return MaskApi._rlesToLua(Q,1)[1]
end

MaskApi.iou = function( dt, gt, iscrowd )
 if not iscrowd then iscrowd = NULL else
 iscrowd = iscrowd:type('torch.ByteTensor'):contiguous():data()
 end
 if torch.isTensor(gt) and torch.isTensor(dt) then
 local nDt, k = dt:size(1), dt:size(2); assert(k==4)
 local nGt, k = gt:size(1), gt:size(2); assert(k==4)
 local dDt = dt:type('torch.DoubleTensor'):contiguous():data()
 local dGt = gt:type('torch.DoubleTensor'):contiguous():data()
 local o = torch.DoubleTensor(nGt,nDt):contiguous()
 libmaskapi.bbIou(dDt,dGt,nDt,nGt,iscrowd,o:data())
 return o:transpose(1,2)
 else
 local qDt, nDt = MaskApi._rlesFrLua(dt)
 local qGt, nGt = MaskApi._rlesFrLua(gt)
 local o = torch.DoubleTensor(nGt,nDt):contiguous()
 libmaskapi.rleIou(qDt,qGt,nDt,nGt,iscrowd,o:data())
 MaskApi._rlesFree(qDt,nDt); MaskApi._rlesFree(qGt,nGt)
 return o:transpose(1,2)
 end
end

MaskApi.nms = function( dt, thr )
 if torch.isTensor(dt) then
 local n, k = dt:size(1), dt:size(2); assert(k==4)
 local Q = dt:type('torch.DoubleTensor'):contiguous():data()
 local kp = torch.IntTensor(n):contiguous()
 libmaskapi.bbNms(Q,n,kp:data(),thr)
 return kp
 else
 local Q, n = MaskApi._rlesFrLua(dt)
 local kp = torch.IntTensor(n):contiguous()
 libmaskapi.rleNms(Q,n,kp:data(),thr)
 MaskApi._rlesFree(Q,n)
 return kp
 end
end

MaskApi.area = function( Rs )
 local Qs, n, h, w = MaskApi._rlesFrLua(Rs)
 local a = torch.IntTensor(n):contiguous()
 libmaskapi.rleArea(Qs,n,a:data())
 MaskApi._rlesFree(Qs,n)
 return a
end

MaskApi.toBbox = function( Rs )
 local Qs, n, h, w = MaskApi._rlesFrLua(Rs)
 local bb = torch.DoubleTensor(n,4):contiguous()
 libmaskapi.rleToBbox(Qs,bb:data(),n)
 MaskApi._rlesFree(Qs,n)
 return bb
end

MaskApi.frBbox = function( bbs, h, w )
 if bbs:dim()==1 then bbs=bbs:view(1,bbs:size(1)) end
 local n, k = bbs:size(1), bbs:size(2); assert(k==4)
 local data = bbs:type('torch.DoubleTensor'):contiguous():data()
 local Qs = MaskApi._rlesInit(n)
 libmaskapi.rleFrBbox(Qs[0],data,h,w,n)
 return MaskApi._rlesToLua(Qs,n)
end

MaskApi.frPoly = function( poly, h, w )
 local n = #poly
 local Qs, Q = MaskApi._rlesInit(n), MaskApi._rlesInit(1)
 for i,p in pairs(poly) do
 local xy = p:type('torch.DoubleTensor'):contiguous():data()
 libmaskapi.rleFrPoly(Qs[i-1],xy,p:size(1)/2,h,w)
 end
 libmaskapi.rleMerge(Qs,Q[0],n,0)
 MaskApi._rlesFree(Qs,n)
 return MaskApi._rlesToLua(Q,1)[1]
end

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

MaskApi.drawCirc = function( img, x, y, rad, clr )
 assert(img:isContiguous() and img:dim()==3)
 local k, h, w, data = img:size(1), img:size(2), img:size(3), img:data()
 for dx=-rad,rad do for dy=-rad,rad do
 local xi, yi = torch.round(x+dx), torch.round(y+dy)
 if dx*dx+dy*dy<=rad*rad and xi>=0 and yi>=0 and xi<w and yi<h then
 for c=1,k do data[(c-1)*h*w + yi*w + xi] = clr[c] end
 end
 end end
end

MaskApi.drawLine = function( img, x0, y0, x1, y1, rad, clr )
 assert(img:isContiguous() and img:dim()==3)
 local k, h, w, data = img:size(1), img:size(2), img:size(3), img:data()
 local dx,dy,d; dx,dy=x1-x0,y1-y0; d=torch.sqrt(dx*dx+dy*dy); dx,dy=dx/d,dy/d
 for i=0,d,.5 do for j=-rad,rad,.5 do
 local xi, yi = torch.round(x0+dx*i+j*dy), torch.round(y0+dy*i-j*dx)
 if xi>=0 and yi>=0 and xi<w and yi<h then
 for c=1,k do data[(c-1)*h*w + yi*w + xi] = clr[c] end
 end
 end end
end

MaskApi.drawMasks = function( img, masks, maxn, alpha, clrs )
 assert(img:isContiguous() and img:dim()==3)
 local n, h, w = masks:size(1), masks:size(2), masks:size(3)
 if not maxn then maxn=n end
 if not alpha then alpha=.4 end
 if not clrs then clrs=torch.rand(n,3)*.6+.4 end
 for i=1,math.min(maxn,n) do
 local M = masks[i]:contiguous():data()
 local B = torch.ByteTensor(h,w):zero():contiguous():data()
 -- get boundaries B in masks M quickly
 for y=0,h-2 do for x=0,w-2 do
 local k=y*w+x
 if M[k]~=M[k+1] then B[k],B[k+1]=1,1 end
 if M[k]~=M[k+w] then B[k],B[k+w]=1,1 end
 if M[k]~=M[k+1+w] then B[k],B[k+1+w]=1,1 end
 end end
 -- softly embed masks into image and add solid boundaries
 for j=1,3 do
 local O,c,a = img[j]:data(), clrs[i][j], alpha
 for k=0,w*h-1 do if M[k]==1 then O[k]=O[k]*(1-a)+c*a end end
 for k=0,w*h-1 do if B[k]==1 then O[k]=c end end
 end
 end
end

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

MaskApi._rlesToLua = function( Qs, n )
 local h, w, Rs = tonumber(Qs[0].h), tonumber(Qs[0].w), {}
 for i=1,n do Rs[i]={size={h,w}, counts={}} end
 for i=1,n do
 local s = libmaskapi.rleToString(Qs[i-1])
 Rs[i].counts=ffi.string(s)
 ffi.C.free(s)
 end
 MaskApi._rlesFree(Qs,n)
 return Rs
end

MaskApi._rlesFrLua = function( Rs )
 if #Rs==0 then Rs={Rs} end
 local n, h, w = #Rs, Rs[1].size[1], Rs[1].size[2]
 local Qs = MaskApi._rlesInit(n)
 for i=1,n do
 local c = Rs[i].counts
 if( torch.type(c)=='string' ) then
 local s=ffi.new("char[?]",#c+1); ffi.copy(s,c)
 libmaskapi.rleFrString(Qs[i-1],s,h,w)
 elseif( torch.type(c)=='torch.IntTensor' ) then
 libmaskapi.rleInit(Qs[i-1],h,w,c:size(1),c:contiguous():data())
 else
 assert(false,"invalid RLE")
 end
 end
 return Qs, n, h, w
end

MaskApi._rlesInit = function( n )
 local Qs = ffi.new("RLE[?]",n)
 for i=1,n do libmaskapi.rleInit(Qs[i-1],0,0,0,NULL) end
 return Qs
end

MaskApi._rlesFree = function( Qs, n )
 for i=1,n do libmaskapi.rleFree(Qs[i-1]) end
end

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

ffi.cdef[[
 void free(void *ptr);
 typedef unsigned int uint;
 typedef unsigned long siz;
 typedef unsigned char byte;
 typedef double* BB;
 typedef struct { siz h, w, m; uint *cnts; } RLE;
 void rleInit( RLE *R, siz h, siz w, siz m, uint *cnts );
 void rleFree( RLE *R );
 void rlesInit( RLE **R, siz n );
 void rlesFree( RLE **R, siz n );
 void rleEncode( RLE *R, const byte *mask, siz h, siz w, siz n );
 void rleDecode( const RLE *R, byte *mask, siz n );
 void rleMerge( const RLE *R, RLE *M, siz n, int intersect );
 void rleArea( const RLE *R, siz n, uint *a );
 void rleIou( RLE *dt, RLE *gt, siz m, siz n, byte *iscrowd, double *o );
 void rleNms( RLE *dt, siz n, uint *keep, double thr );
 void bbIou( BB dt, BB gt, siz m, siz n, byte *iscrowd, double *o );
 void bbNms( BB dt, siz n, uint *keep, double thr );
 void rleToBbox( const RLE *R, BB bb, siz n );
 void rleFrBbox( RLE *R, const BB bb, siz h, siz w, siz n );
 void rleFrPoly( RLE *R, const double *xy, siz k, siz h, siz w );
 char* rleToString( const RLE *R );
 void rleFrString( RLE *R, char *s, siz h, siz w );
]]