example-apps/PDD/pcb-defect-detection/data/lib_coco/get_coco_next_batch.py

   1 # -*- coding: utf-8 -*-
   2
   3 from __future__ import absolute_import, print_function, division
   4
   5 import sys, os
   6 # sys.path.insert(0, os.path.abspath('.'))
   7 sys.path.insert(0, './PythonAPI/')
   8 # sys.path.insert(0, os.path.abspath('data'))
   9 for _ in sys.path:
  10     print (_)
  11 from PythonAPI.pycocotools.coco import COCO
  12 import cv2
  13 import numpy as np
  14 import os
  15 from libs.label_name_dict import coco_dict
  16
  17
  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)
  21
  22 imgId_list = coco.getImgIds()
  23 imgId_list = np.array(imgId_list)
  24
  25 total_imgs = len(imgId_list)
  26
  27 # print (NAME_LABEL_DICT)
  28
  29
  30 def next_img(step):
  31
  32     if step % total_imgs == 0:
  33         np.random.shuffle(imgId_list)
  34     imgid = imgId_list[step % total_imgs]
  35
  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))
  39
  40     annotation = coco.imgToAnns[imgid]
  41     gtbox_and_label_list = []
  42     for ann in annotation:
  43         box = ann['bbox']
  44
  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)
  54     else:
  55         return imgid, img[:, :, ::-1], gtbox_and_label_list
  56
  57
  58 if __name__ == '__main__':
  59
  60     imgid, img,  gtbox = next_img(3234)
  61
  62     print("::")
  63     from libs.box_utils.draw_box_in_img import draw_boxes_with_label_and_scores
  64
  65     img = draw_boxes_with_label_and_scores(img_array=img, boxes=gtbox[:, :-1], labels=gtbox[:, -1],
  66                                            scores=np.ones(shape=(len(gtbox), )))
  67     print ("_----")
  68
  69
  70     cv2.imshow("test", img)
  71     cv2.waitKey(0)
  72
  73