Add files via upload

scripts/cocoapi/LuaAPI/CocoApi.lua

@@ -0,0 +1,261 @@
+--[[----------------------------------------------------------------------------
+
+Interface for accessing the Common Objects in COntext (COCO) dataset.
+
+For an overview of the API please see http://mscoco.org/dataset/#download.
+CocoApi.lua (this file) is modeled after the Matlab CocoApi.m:
+  https://github.com/pdollar/coco/blob/master/MatlabAPI/CocoApi.m
+
+The following API functions are defined in the Lua API:
+  CocoApi    - Load COCO annotation file and prepare data structures.
+  getAnnIds  - Get ann ids that satisfy given filter conditions.
+  getCatIds  - Get cat ids that satisfy given filter conditions.
+  getImgIds  - Get img ids that satisfy given filter conditions.
+  loadAnns   - Load anns with the specified ids.
+  loadCats   - Load cats with the specified ids.
+  loadImgs   - Load imgs with the specified ids.
+  showAnns   - Display the specified annotations.
+Throughout the API "ann"=annotation, "cat"=category, and "img"=image.
+For detailed usage information please see cocoDemo.lua.
+
+LIMITATIONS: the following API functions are NOT defined in the Lua API:
+  loadRes    - Load algorithm results and create API for accessing them.
+  download   - Download COCO images from mscoco.org server.
+In addition, currently the getCatIds() and getImgIds() do not accept filters.
+getAnnIds() can be called using getAnnIds({imgId=id}) and getAnnIds({catId=id}).
+
+Note: loading COCO JSON annotations to Lua tables is quite slow. Hence, a call
+to CocApi(annFile) converts the annotations to a custom 'flattened' format that
+is more efficient. The first time a COCO JSON is loaded, the conversion is
+invoked (this may take up to a minute). The converted data is then stored in a
+t7 file (the code must have write permission to the dir of the JSON file).
+Future calls of cocoApi=CocApi(annFile) take a fraction of a second. To view the
+created data just inspect cocoApi.data of a created instance of the CocoApi.
+
+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 json = require 'cjson'
+local coco = require 'coco.env'
+
+local TensorTable = torch.class('TensorTable',coco)
+local CocoSeg = torch.class('CocoSeg',coco)
+local CocoApi = torch.class('CocoApi',coco)
+
+--------------------------------------------------------------------------------
+
+--[[ TensorTable is a lightweight data structure for storing variable size 1D
+tensors. Tables of tensors are slow to save/load to disk. Instead, TensorTable
+stores all the data in a single long tensor (along with indices into the tensor)
+making serialization fast. A TensorTable may only contain 1D same-type torch
+tensors or strings. It supports only creation from a table and indexing. ]]
+
+function TensorTable:__init( T )
+  local n = #T; assert(n>0)
+  local isStr = torch.type(T[1])=='string'
+  assert(isStr or torch.isTensor(T[1]))
+  local c=function(s) return torch.CharTensor(torch.CharStorage():string(s)) end
+  if isStr then local S=T; T={}; for i=1,n do T[i]=c(S[i]) end end
+  local ms, idx = torch.LongTensor(n), torch.LongTensor(n+1)
+  for i=1,n do ms[i]=T[i]:numel() end
+  idx[1]=1; idx:narrow(1,2,n):copy(ms); idx=idx:cumsum()
+  local type = string.sub(torch.type(T[1]),7,-1)
+  local data = torch[type](idx[n+1]-1)
+  if isStr then type='string' end
+  for i=1,n do if ms[i]>0 then data:sub(idx[i],idx[i+1]-1):copy(T[i]) end end
+  if ms:eq(ms[1]):all() and ms[1]>0 then data=data:view(n,ms[1]); idx=nil end
+  self.data, self.idx, self.type = data, idx, type
+end
+
+function TensorTable:__index__( i )
+  if torch.type(i)~='number' then return false end
+  local d, idx, type = self.data, self.idx, self.type
+  if idx and idx[i]==idx[i+1] then
+    if type=='string' then d='' else d=torch[type]() end
+  else
+    if idx then d=d:sub(idx[i],idx[i+1]-1) else d=d[i] end
+    if type=='string' then d=d:clone():storage():string() end
+  end
+  return d, true
+end
+
+--------------------------------------------------------------------------------
+
+--[[ CocoSeg is an efficient data structure for storing COCO segmentations. ]]
+
+function CocoSeg:__init( segs )
+  local polys, pIdx, sizes, rles, p, isStr = {}, {}, {}, {}, 0, 0
+  for i,seg in pairs(segs) do if seg.size then isStr=seg.counts break end end
+  isStr = torch.type(isStr)=='string'
+  for i,seg in pairs(segs) do
+    pIdx[i], sizes[i] = {}, {}
+    if seg.size then
+      sizes[i],rles[i] = seg.size,seg.counts
+    else
+      if isStr then rles[i]='' else rles[i]={} end
+      for j=1,#seg do p=p+1; pIdx[i][j],polys[p] = p,seg[j] end
+    end
+    pIdx[i],sizes[i] = torch.LongTensor(pIdx[i]),torch.IntTensor(sizes[i])
+    if not isStr then rles[i]=torch.IntTensor(rles[i]) end
+  end
+  for i=1,p do polys[i]=torch.DoubleTensor(polys[i]) end
+  self.polys, self.pIdx = coco.TensorTable(polys), coco.TensorTable(pIdx)
+  self.sizes, self.rles = coco.TensorTable(sizes), coco.TensorTable(rles)
+end
+
+function CocoSeg:__index__( i )
+  if torch.type(i)~='number' then return false end
+  if self.sizes[i]:numel()>0 then
+    return {size=self.sizes[i],counts=self.rles[i]}, true
+  else
+    local ids, polys = self.pIdx[i], {}
+    for i=1,ids:numel() do polys[i]=self.polys[ids[i]] end
+    return polys, true
+  end
+end
+
+--------------------------------------------------------------------------------
+
+--[[ CocoApi is the API to the COCO dataset, see main comment for details. ]]
+
+function CocoApi:__init( annFile )
+  assert( string.sub(annFile,-4,-1)=='json' and paths.filep(annFile) )
+  local torchFile = string.sub(annFile,1,-6) .. '.t7'
+  if not paths.filep(torchFile) then self:__convert(annFile,torchFile) end
+  local data = torch.load(torchFile)
+  self.data, self.inds = data, {}
+  for k,v in pairs({images='img',categories='cat',annotations='ann'}) do
+    local M = {}; self.inds[v..'IdsMap']=M
+    if data[k] then for i=1,data[k].id:size(1) do M[data[k].id[i]]=i end end
+  end
+end
+
+function CocoApi:__convert( annFile, torchFile )
+  print('convert: '..annFile..' --> .t7 [please be patient]')
+  local tic = torch.tic()
+  -- load data and decode json
+  local data = torch.CharStorage(annFile):string()
+  data = json.decode(data); collectgarbage()
+  -- transpose and flatten each field in the coco data struct
+  local convert = {images=true, categories=true, annotations=true}
+  for field, d in pairs(data) do if convert[field] then
+    print('converting: '..field)
+    local n, out = #d, {}
+    if n==0 then d,n={d},1 end
+    for k,v in pairs(d[1]) do
+      local t, isReg = torch.type(v), true
+      for i=1,n do isReg=isReg and torch.type(d[i][k])==t end
+      if t=='number' and isReg then
+        out[k] = torch.DoubleTensor(n)
+        for i=1,n do out[k][i]=d[i][k] end
+      elseif t=='string' and isReg then
+        out[k]={}; for i=1,n do out[k][i]=d[i][k] end
+        out[k] = coco.TensorTable(out[k])
+      elseif t=='table' and isReg and torch.type(v[1])=='number' then
+        out[k]={}; for i=1,n do out[k][i]=torch.DoubleTensor(d[i][k]) end
+        out[k] = coco.TensorTable(out[k])
+        if not out[k].idx then out[k]=out[k].data end
+      else
+        out[k]={}; for i=1,n do out[k][i]=d[i][k] end
+        if k=='segmentation' then out[k] = coco.CocoSeg(out[k]) end
+      end
+      collectgarbage()
+    end
+    if out.id then out.idx=torch.range(1,out.id:size(1)) end
+    data[field] = out
+    collectgarbage()
+  end end
+  -- create mapping from cat/img index to anns indices for that cat/img
+  print('convert: building indices')
+  local makeMap = function( type, type_id )
+    if not data[type] or not data.annotations then return nil end
+    local invmap, n = {}, data[type].id:size(1)
+    for i=1,n do invmap[data[type].id[i]]=i end
+    local map = {}; for i=1,n do map[i]={} end
+    data.annotations[type_id..'x'] = data.annotations[type_id]:clone()
+    for i=1,data.annotations.id:size(1) do
+      local id = invmap[data.annotations[type_id][i]]
+      data.annotations[type_id..'x'][i] = id
+      table.insert(map[id],data.annotations.id[i])
+    end
+    for i=1,n do map[i]=torch.LongTensor(map[i]) end
+    return coco.TensorTable(map)
+  end
+  data.annIdsPerImg = makeMap('images','image_id')
+  data.annIdsPerCat = makeMap('categories','category_id')
+  -- save to disk
+  torch.save( torchFile, data )
+  print(('convert: complete [%.2f s]'):format(torch.toc(tic)))
+end
+
+function CocoApi:getAnnIds( filters )
+  if not filters then filters = {} end
+  if filters.imgId then
+    return self.data.annIdsPerImg[self.inds.imgIdsMap[filters.imgId]] or {}
+  elseif filters.catId then
+    return self.data.annIdsPerCat[self.inds.catIdsMap[filters.catId]] or {}
+  else
+    return self.data.annotations.id
+  end
+end
+
+function CocoApi:getCatIds()
+  return self.data.categories.id
+end
+
+function CocoApi:getImgIds()
+  return self.data.images.id
+end
+
+function CocoApi:loadAnns( ids )
+  return self:__load(self.data.annotations,self.inds.annIdsMap,ids)
+end
+
+function CocoApi:loadCats( ids )
+  return self:__load(self.data.categories,self.inds.catIdsMap,ids)
+end
+
+function CocoApi:loadImgs( ids )
+  return self:__load(self.data.images,self.inds.imgIdsMap,ids)
+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 Rs, alpha = {}, anns[1].keypoints and .25 or .4
+    for i=1,n do
+      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
+    MaskApi.drawMasks(O,MaskApi.decode(Rs),nil,alpha,clrs)
+  end
+  return O
+end
+
+function CocoApi:__load( data, map, ids )
+  if not torch.isTensor(ids) then ids=torch.LongTensor({ids}) end
+  local out, idx = {}, nil
+  for i=1,ids:numel() do
+    out[i], idx = {}, map[ids[i]]
+    for k,v in pairs(data) do out[i][k]=v[idx] end
+  end
+  return out
+end