LuaApi/: updating demo and draw capabilities for keypoints

LuaAPI/CocoApi.lua

@@ -225,19 +225,29 @@ end
 
 function CocoApi:showAnns( img, anns )
   local n, h, w = #anns, img:size(2), img:size(3)
+  local MaskApi, clrs = coco.MaskApi, torch.rand(n,3)*.6+.4
+  local O = img:clone():contiguous():float()
   if n==0 then anns,n={anns},1 end
+  if anns[1].keypoints then for i=1,n do if anns[i].iscrowd==0 then
+    local sk, kp, j, k = self:loadCats(anns[i].category_id)[1].skeleton
+    kp=anns[i].keypoints; k=kp:size(1); j=torch.range(1,k,3):long(); k=k/3;
+    local x,y,v = kp:index(1,j), kp:index(1,j+1), kp:index(1,j+2)
+    for _,s in pairs(sk) do if v[s[1]]>0 and v[s[2]]>0 then
+      MaskApi.drawLine(O,x[s[1]],y[s[1]],x[s[2]],y[s[2]],.75,clrs[i])
+    end end
+    for j=1,k do if v[j]==1 then MaskApi.drawCirc(O,x[j],y[j],4,{0,0,0}) end end
+    for j=1,k do if v[j]>0 then MaskApi.drawCirc(O,x[j],y[j],3,clrs[i]) end end
+  end end end
   if anns[1].segmentation or anns[1].bbox then
-    local MaskApi, Rs = coco.MaskApi, {}
+    local Rs, alpha = {}, anns[1].keypoints and .25 or .4
     for i=1,n do
-      if anns[i].segmentation then
-        Rs[i]=anns[i].segmentation
-        if #Rs[i]>0 then Rs[i]=MaskApi.frPoly(Rs[i],h,w) end
-      elseif anns[i].bbox then
-        Rs[i]=MaskApi.frBbox(anns[i].bbox,h,w)[1]
-      end
+      Rs[i]=anns[i].segmentation
+      if Rs[i] and #Rs[i]>0 then Rs[i]=MaskApi.frPoly(Rs[i],h,w) end
+      if not Rs[i] then Rs[i]=MaskApi.frBbox(anns[i].bbox,h,w)[1] end
     end
-    return MaskApi.drawMasks(img,MaskApi.decode(Rs))
+    MaskApi.drawMasks(O,MaskApi.decode(Rs),nil,alpha,clrs)
   end
+  return O
 end
 
 function CocoApi:__load( data, map, ids )

LuaAPI/MaskApi.lua

@@ -15,6 +15,9 @@ The following API functions are defined:
   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 )
@@ -25,7 +28,9 @@ Usage:
   bbs    = MaskApi.toBbox( Rs )
   Rs     = MaskApi.frBbox( bbs, h, w )
   R      = MaskApi.frPoly( poly, h, w )
-  img    = MaskApi.drawMasks( img, masks, [maxn=n], [alpha=.4], [clrs] )
+  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:
@@ -139,12 +144,37 @@ MaskApi.frPoly = function( poly, h, w )
   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)*1.4 end
-  local out = img:clone():contiguous():float()
+  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()
@@ -157,12 +187,11 @@ MaskApi.drawMasks = function( img, masks, maxn, alpha, clrs )
     end end
     -- softly embed masks into image and add solid boundaries
     for j=1,3 do
-      local O,c,a = out[j]:data(), math.min(clrs[i][j],1), alpha
+      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
-  return out
 end
 
 --------------------------------------------------------------------------------

LuaAPI/cocoDemo.lua

@@ -2,18 +2,19 @@
 coco = require 'coco'
 image = require 'image'
 
--- initialize COCO api for instance annotations
-dataDir, dataType = '../', 'val2014'
-annFile = dataDir..'/annotations/instances_'..dataType..'.json'
-if not cocoApi then cocoApi=coco.CocoApi(annFile) end
+-- initialize COCO api (please specify dataType/annType below)
+annTypes = { 'instances', 'captions', 'person_keypoints' }
+dataType, annType = 'val2014', annTypes[1]; -- specify dataType/annType
+annFile = '../annotations/'..annType..'_'..dataType..'.json'
+cocoApi=coco.CocoApi(annFile)
 
 -- get all image ids, select one at random
 imgIds = cocoApi:getImgIds()
 imgId = imgIds[torch.random(imgIds:numel())]
 
 -- load image
 img = cocoApi:loadImgs(imgId)[1]
-I = image.load(dataDir..'/images/'..dataType..'/'..img.file_name,3)
+I = image.load('../images/'..dataType..'/'..img.file_name,3)
 
 -- load and display instance annotations
 annIds = cocoApi:getAnnIds({imgId=imgId})