gray灰度标签图转coco数据集

import cv2
import os
import numpy as np
import json

# / *0
# beijing
# *1
# yalibiaoyuanhuan
# *2
# yellow
# *3
# green
# *4
# red
# *5
# blue
# *6
# indicator
# * /

###########################################################################################
## LABEL: yalibiao
# my_label = {"background":0,
#             "yuanhuan":1,
#             "yellow":2,
#             "green":3,
#             "red":4,
#             "blue":5,
#             "indicator":6}

my_label = {"background":0,
            "0":1,
            "1":2,
            "2":3,
            "3":4,
            "4":5,
            "5":6,
            "6":7,
            "7":8,
            "8":9,
            "9":10,
            "A":11,
            "B":12,
            "C": 13,
            "D": 14,
            "E": 15,
            "F": 16,
            "G": 17,
            "H": 18,
            "J": 19,
            "K": 20,
            "L": 21,
            "M": 22,
            "N": 23,
            "P": 24,
            "R": 25,
            "S": 26,
            "T": 27,
            "U": 28,
            "V": 29,
            "W": 30,
            "X": 31,
            "Y": 32,
            "Z": 33
            }

def gray2src_name(gray_img_name):
    #gray_name:  gray_sz_ylb_0.png
    #src_name:   src_sz_ylb_0.png
    img_name = gray_img_name.replace("gray", "src")

    #img_name = img_name.replace("png", "jpg")
    return img_name
#############################################################################################



coco = dict()
coco['images'] = []
coco['type'] = 'instances'
coco['annotations'] = []
coco['categories'] = []

category_set = dict()
image_set = set()

category_item_id = 0
image_id = 20190000000
annotation_id = 0

def addLabelItem(my_label):
    my_label = sorted(my_label.items(),key=lambda item:item[1])
    for val in my_label:
        category_item = dict()
        category_item['supercategory'] = 'none'
        category_item_id = val[1]
        if 0 == category_item_id:
            continue
        category_item['id'] = category_item_id
        category_item['name'] = val[0]
        coco['categories'].append(category_item)


def addImgItem(file_name, img):
    global image_id
    if file_name is None:
        raise Exception('Could not find filename tag in xml file.')
    if img is None:
        raise Exception('Could not find img.')

    image_id += 1
    image_item = dict()
    image_item['id'] = image_id
    image_item['file_name'] = file_name
    image_item['width'] = img.shape[1] #size['width']
    image_item['height'] = img.shape[0]   #size['height']
    coco['images'].append(image_item)
    image_set.add(file_name)
    return image_id

def addAnnoItem( image_id, category_id, area, seg, bbox):
    bbox = list(bbox)
    global annotation_id
    annotation_item = dict()
    annotation_item['segmentation'] = []
    annotation_item['segmentation'].append(seg)

    annotation_item['area'] = area
    annotation_item['iscrowd'] = 0
    annotation_item['ignore'] = 0
    annotation_item['image_id'] = image_id
    annotation_item['bbox'] = bbox
    annotation_item['category_id'] = category_id
    annotation_id += 1
    annotation_item['id'] = annotation_id
    coco['annotations'].append(annotation_item)

def gray_img2coco(root):
    addLabelItem(my_label)
    list_file = os.listdir(root)
    num_class = len(my_label)
    cnt = 0
    for img_name in list_file:
        cnt += 1
        print("cnt=%d,  img_name=%s"%(cnt,img_name))
        img_path = root + img_name
        img_gray = cv2.imread(img_path,-1)
        if len(img_gray.shape) == 3:
            img_gray = cv2.cvtColor(img_gray, cv2.COLOR_BGR2GRAY)
        #points = np.array(np.where(6 == img_gray)).transpose((1, 0))[:, ::-1]

        height = img_gray.shape[0]
        width = img_gray.shape[1]
        #img_name = img_name.replace("gray","src")
        img_name = gray2src_name(img_name)
        current_image_id = addImgItem(img_name, img_gray)
        for current_category_id in range(1,num_class):
            img_bi = np.zeros(img_gray.shape,dtype='uint8')
            img_bi[img_gray == current_category_id] = 255
            my_contours,_ = cv2.findContours(img_bi,cv2.RETR_EXTERNAL,cv2.CHAIN_APPROX_NONE)
            for c in my_contours:
                area_t = cv2.contourArea(c)
                #print("area=%d"%(cv2.contourArea(c)))
                # print("L_pt_size=%d"%(len(L_pt)))
                if 0 == len(c) or area_t < 20:
                    continue
                L_pt = c
                # x,y,w,h
                bbox = cv2.boundingRect(L_pt)
                x1,y1,w1,h1 = bbox
                if x1<0 or y1<0 or x1+w1>width or y1+h1>height:
                    continue
                seg = []
                for val in L_pt:
                    #print (val)
                    x = val[0][0]
                    y = val[0][1]
                    seg.append(int(x)) ##zhongyao
                    seg.append(int(y))
                    # print("x=%d,y=%d"%(x,y))

                addAnnoItem(current_image_id, current_category_id, area_t, seg, bbox)

                # #m_t = np.ones(img_gray.shape,np.uint8)*255
                # m_t = np.zeros([img_gray.shape[0], img_gray.shape[1], 3], dtype='uint8')
                # cv2.drawContours(m_t,[L_pt],0,(255,0,0),1)


                # cv2.imshow("m_t",m_t)
                # cv2.imshow("img_bi", img_bi)
                # cv2.imshow("gray",img_gray*20)
                # cv2.waitKey(0)


if __name__ == '__main__':
    root_gray_img =  "/maskrcnn-benchmark/datasets/coco/tmp/val_gray/"
    #  instances_train2014.json    instances_val2014.json
    json_file = "maskrcnn-benchmark/datasets/coco/tmp/instances_val2014.json"
    gray_img2coco(root_gray_img)
    json.dump(coco,open(json_file,'w'))