import torch
from PIL import Image
from matplotlib import pyplot as plt
from pycocotools.coco import COCO
KEYPOINTS = ['nose', 'left_eye', 'right_eye', 'left_ear', 'right_ear',
'left_shoulder', 'right_shoulder', 'left_elbow',
'right_elbow', 'left_wrist', 'right_wrist', 'left_hip',
'right_hip', 'left_knee', 'right_knee', 'left_ankle',
'right_ankle']
_coco_helper = COCO()
_coco_helper.cats = {1: {'id': 1,
'keypoints': KEYPOINTS,
'name': 'person',
'skeleton': [[16, 14], [14, 12], [17, 15], [15, 13], [12, 13], [6, 12],
[7, 13], [6, 7], [6, 8], [7, 9], [8, 10], [9, 11], [2, 3],
[1, 2], [1, 3], [2, 4], [3, 5], [4, 6], [5, 7]],
'supercategory': 'person'}}
class Coco(COCO):
def __init__(self, img_dir, ann_file):
self.image_dir = Path(img_dir)
super().__init__(ann_file)
def show_by_id(self, img_id):
img = self.get_img(img_id)
annotations = self.get_annotations(img_id)
def _init_coco(self):
coco = COCO()
# self.anns, self.cats, self.imgs self.dataset, self.imgToAnns, self.catToImgs
# generate each member respectively
det_anno_list = OrderedDict()
images = OrderedDict()
img2anns = OrderedDict()
cat2imgs = defaultdict(list)
for each_anno in self.vg_annotations:
img_id = each_anno['id']
images[img_id] = {
'file_name': each_anno['path'],
'height': each_anno['height'],
'width': each_anno['width'],
'id': img_id
}
# each image contains many annotation
# add sub id in single image
obj_list = []
for each_object in each_anno['objects']:
xywh = xyxy2xywh(each_object['box'])
obj_list.append({
'category_id':
self.obj_cls_ind_q[each_object['class']],
'bbox':
xywh,
'image_id':
img_id,
'area':
xywh[2] * xywh[3],
'iscrowd':
0,
'segmentation': [],
'id':
-1 # left for last step accumulate
})
# accumulate the sub_id to overall id
last_seg_idx = len(det_anno_list.keys())
for indx, each in enumerate(obj_list):
anno_id = last_seg_idx + indx
obj_list[indx]['id'] = anno_id
det_anno_list[anno_id] = each
cat2imgs[each['category_id']].append(img_id)
img2anns[img_id] = obj_list
coco.anns = det_anno_list
coco.imgs = images
coco.imgToAnns = img2anns
coco.catToImgs = cat2imgs
# cats
cats = OrderedDict()
for id_, each_cat in enumerate(self.obj_cls_list):
if id_ == 0:
continue
cats[id_] = {
'id': id_,
'name': each_cat,
}
coco.cats = cats
# other info
coco.dataset['licenses'] = ''
coco.dataset['annotations'] = list(coco.anns.values())
coco.dataset['info'] = {}
coco.dataset['images'] = list(coco.imgs.values())
coco.dataset['categories'] = list(coco.cats.values())
self.coco = coco
self.ids = list(self.coco.imgs.keys())
