1 # -*- coding: utf-8 -*-
3 from __future__ import absolute_import, print_function, division
6 # sys.path.insert(0, os.path.abspath('.'))
7 sys.path.insert(0, './PythonAPI/')
8 # sys.path.insert(0, os.path.abspath('data'))
11 from PythonAPI.pycocotools.coco import COCO
15 from libs.label_name_dict import coco_dict
18 annotation_path = '/home/yjr/DataSet/COCO/2017/annotations/instances_train2017.json'
19 print ("load coco .... it will cost about 17s..")
20 coco = COCO(annotation_path)
22 imgId_list = coco.getImgIds()
23 imgId_list = np.array(imgId_list)
25 total_imgs = len(imgId_list)
27 # print (NAME_LABEL_DICT)
32 if step % total_imgs == 0:
33 np.random.shuffle(imgId_list)
34 imgid = imgId_list[step % total_imgs]
36 imgname = coco.loadImgs(ids=[imgid])[0]['file_name']
37 # print (type(imgname), imgname)
38 img = cv2.imread(os.path.join("/home/yjr/DataSet/COCO/2017/train2017", imgname))
40 annotation = coco.imgToAnns[imgid]
41 gtbox_and_label_list = []
42 for ann in annotation:
45 box = [box[0], box[1], box[0]+box[2], box[1]+box[3]] # [xmin, ymin, xmax, ymax]
46 cat_id = ann['category_id']
47 cat_name = coco_dict.originID_classes[cat_id] #ID_NAME_DICT[cat_id]
48 label = coco_dict.NAME_LABEL_MAP[cat_name]
49 gtbox_and_label_list.append(box + [label])
50 gtbox_and_label_list = np.array(gtbox_and_label_list, dtype=np.int32)
51 # print (img.shape, gtbox_and_label_list.shape)
52 if gtbox_and_label_list.shape[0] == 0:
53 return next_img(step+1)
55 return imgid, img[:, :, ::-1], gtbox_and_label_list
58 if __name__ == '__main__':
60 imgid, img, gtbox = next_img(3234)
63 from libs.box_utils.draw_box_in_img import draw_boxes_with_label_and_scores
65 img = draw_boxes_with_label_and_scores(img_array=img, boxes=gtbox[:, :-1], labels=gtbox[:, -1],
66 scores=np.ones(shape=(len(gtbox), )))
70 cv2.imshow("test", img)