Find file
2934299 Jul 17, 2016
262 lines (232 sloc) 10.5 KB
Interface for accessing the Common Objects in COntext (COCO) dataset.
For an overview of the API please see
CocoApi.lua (this file) is modeled after the Matlab 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 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 of a created instance of the CocoApi.
Common Objects in COntext (COCO) Toolbox. version 3.0
Data, paper, and tutorials available at:
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.idx, self.type = data, idx, type
function TensorTable:__index__( i )
if torch.type(i)~='number' then return false end
local d, idx, type =, self.idx, self.type
if idx and idx[i]==idx[i+1] then
if type=='string' then d='' else d=torch[type]() end
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
return d, true
--[[ 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
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
pIdx[i],sizes[i] = torch.LongTensor(pIdx[i]),torch.IntTensor(sizes[i])
if not isStr then rles[i]=torch.IntTensor(rles[i]) 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)
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
local ids, polys = self.pIdx[i], {}
for i=1,ids:numel() do polys[i]=self.polys[ids[i]] end
return polys, true
--[[ 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.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
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
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
if then out.idx=torch.range(1, end
data[field] = out
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, do
local id = invmap[data.annotations[type_id][i]]
data.annotations[type_id..'x'][i] = id
for i=1,n do map[i]=torch.LongTensor(map[i]) end
return coco.TensorTable(map)
data.annIdsPerImg = makeMap('images','image_id')
data.annIdsPerCat = makeMap('categories','category_id')
-- save to disk torchFile, data )
print(('convert: complete [%.2f s]'):format(torch.toc(tic)))
function CocoApi:getAnnIds( filters )
if not filters then filters = {} end
if filters.imgId then
return[self.inds.imgIdsMap[filters.imgId]] or {}
elseif filters.catId then
return[self.inds.catIdsMap[filters.catId]] or {}
function CocoApi:getCatIds()
function CocoApi:getImgIds()
function CocoApi:loadAnns( ids )
return self:__load(,self.inds.annIdsMap,ids)
function CocoApi:loadCats( ids )
return self:__load(,self.inds.catIdsMap,ids)
function CocoApi:loadImgs( ids )
return self:__load(,self.inds.imgIdsMap,ids)
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
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
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
return O
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
return out