鵝要長大

# 20240612: parse jsonfile to labeled image.
import argparse
import json
import os
import os.path as osp
import cv2 as cv
import numpy as np

jsonfilename = "freespace_3Dbox_APA.json"
imagepath = "image_APA"
polygon = {'Chock block', 'Convex Mirror', 'Crosswalk', 'Curb', 'Dashed line', 'Directional arrow', 'Fire Extinguisher', 'Fire hydrant', 'Floating edge', 'Free space', 'Motor vehicle', 'Non-motorized vehicle', 'Other Obstacles', 'Parking line', 'Parking number', 'Pedestrian', 'Pylon', 'Solid line', 'Speed bump', 'Steel pipe', 'Traffic cone', 'Wheel lock', 'background'}
bbox2d = {'Chock block 2D', 'Convex Mirror 2D', 'Fire Extinguisher 2D', 'Fire hydrant 2D', 'Motor vehicle 2D', 'Non-motorized vehicle 2D', 'Other Obstacles 2D', 'Pedestrian 2D', 'Steel pipe 2D', 'Traffic cone 2D', 'Tyre', 'Wheel lock 2D'}
bbox3d = {'Motor vehicle 3D'}

colors = {# BGR
        'Chock block'              : (251, 56, 56), 
        'Chock block 2D'           : (247,136, 58), 
        'Convex Mirror'            : ( 48,132, 48), 
        'Convex Mirror 2D'         : ( 56,148,228), 
        'Crosswalk'                : (243, 27,243), 
        'Curb'                     : ( 22,220,220), 
        'Dashed line'              : (255,192,203), 
        'Directional arrow'        : (  0,  0,128), 
        'Fire Extinguisher'        : (  0,255,  0), 
        'Fire Extinguisher 2D'     : (255,255,  0), 
        'Fire hydrant'             : (247, 72,249), 
        'Fire hydrant 2D'          : (103,198,209), 
        'Floating edge'            : ( 88,139, 70), 
        'Free space'               : (172,227, 65), 
        'Motor vehicle'            : ( 76,213,174), 
        'Motor vehicle 2D'         : ( 50, 83, 83), 
        'Motor vehicle 3D'         : (159,156, 73), 
        'Non-motorized vehicle'    : ( 10,239,189), 
        'Non-motorized vehicle 2D' : (183, 68,121), 
        'Other Obstacles'          : (193, 52,240), 
        'Other Obstacles 2D'       : ( 53,223, 54), 
        'Parking line'             : (  4,178,106), 
        'Parking number'           : (104, 28,162), 
        'Pedestrian'               : (  8,197,115),
        'Pedestrian 2D'            : ( 61,205,183),
        'Pylon'                    : (  5,129,227), 
        'Solid line'               : (239, 18, 44), 
        'Speed bump'               : (107,251,229), 
        'Steel pipe'               : (101, 49,243), 
        'Steel pipe 2D'            : ( 26,130, 36), 
        'Traffic cone'             : ( 54,175,167), 
        'Traffic cone 2D'          : ( 39,219,102), 
        'Tyre'                     : (142, 16,198),
        'Wheel lock'               : (174,143,124), 
        'Wheel lock 2D'            : (229,208,251), 
        'background'               : (  0,  0,  0)
    }
def parse_data_polygon(tagtype, datastr):
    # print('datastr: ', datastr)
    if isinstance(datastr, str):
        data = eval(datastr)
    assert isinstance(data, list)
    # print('data: ', data)

    pts = []
    # list - list - float
    # print('data len: ', len(data))
    # print('data[-1] len: ', len(data[-1]))
    # print('data: ', data)
    # print('data type: ', type(data))
    if len(data)==4 and isinstance(data[0], float):
        print("BBOX2D ANNOtation ERRORRRRRR, error tagtype is: ", tagtype)
    else:
        for onedata in data:
            if len(onedata) == 1:
                break
            pts.append(list(map(int, onedata[1:]))) # [x,y]
    return pts

def parse_data_bbox2d(datastr):
    # print('datastr: ', datastr)
    if isinstance(datastr, str):
        data = eval(datastr)
    assert isinstance(data, list)
    # print('data: ', data)
    box = list(map(int, data)) # [x, y, w, h]
    return box

def parse_json(path, dataset):
    i = 0
    for imgdata in dataset:
        # one image anno data.
        result = imgdata['result']
        i = i + 1
        # if not result or (i>1):
        if not result:
            # break
            continue
        file_obj = imgdata['file_obj']
        imgname = osp.split(file_obj)[-1]
        imgdata= cv.imread(osp.join(path, imagepath, imgname))
        # print("file_obj: ", file_obj)
        print("imgname: ", imgname)
        sgname = imgname.replace('.png', '_sg.png')
        odname = imgname.replace('.png', '_od.png')
        sgfile = osp.join(path, 'sg', sgname)
        odfile = osp.join(path, 'od', odname)
        # sgimg = cv.imread(osp.join(path, imagepath, imgname))
        sgimg = np.zeros_like(imgdata)
        odimg = cv.imread(osp.join(path, imagepath, imgname))
        mask = np.zeros_like(imgdata)
        maskname = imgname.replace('.png', '_mask.png')
        maskfile = osp.join(path, 'mask', maskname)
        # print("result type: ", type(result))
        # print("result len: ", len(result))
        # print("result[0]: ", result[0])

        # draw Free space
        for res in result:
            tagtype = res['tagtype']
            # print("tagtype--------------------------------------------------------------: \n\n", tagtype)
            if tagtype != 'Free space':
                continue
            # Free space
            points= parse_data_polygon(tagtype, res['data'])
            # plot polygon in seg image
            points_np = np.array(points, np.int32)
            if points:
                cv.fillPoly(mask, [points_np], colors[tagtype])

        for res in result:
            tagtype = res['tagtype']
            if tagtype == 'Free space':
                continue
            if tagtype == 'Motor vehicle 3D': 
                # print("tagtype--------------------------------------------------------------: ", tagtype)
                points= parse_data_polygon(tagtype, res['data'])
                # plot bbox3d in od image
                points_np = np.array(points, np.int32)
                if points:
                    cv.polylines(odimg, [points_np], True, colors[tagtype])
            elif tagtype == 'Floating edge': 
                # print("tagtype--------------------------------------------------------------: ", tagtype)
                points= parse_data_polygon(tagtype, res['data'])
                # plot floating edge in seg image
                points_np = np.array(points, np.int32)
                if points:
                    cv.polylines(mask, [points_np], False, colors[tagtype], 5)
            elif tagtype in polygon: 
                # print("tagtype--------------------------------------------------------------: ", tagtype)
                points= parse_data_polygon(tagtype, res['data'])
                # plot polygon in seg image
                points_np = np.array(points, np.int32)
                # print('points_np: ', points_np)
                if points:
                    cv.fillPoly(mask, [points_np], colors[tagtype])
            elif tagtype in bbox2d: 
                x, y, w, h = parse_data_bbox2d(res['data']) # Rect(xywh)
                # rect = cv.Rect(x, y, w, h)
                # plot bbox2d in od image
                cv.rectangle(odimg, (x,y), (x+w, y+h), colors[tagtype]) # (leftup, rightdown) 
                # cv.rectangle(odimg, rect, color) #  
        cv.addWeighted(imgdata, 0.7, mask, 0.3, 0, sgimg)
        # save seg and od image.
        cv.imwrite(sgfile, sgimg)
        cv.imwrite(odfile, odimg)
        cv.imwrite(maskfile, mask)

def json2annoimg(path):
    jsonfile=os.path.join(path, jsonfilename)
    dataset=json.load(open(jsonfile, 'r'))
    num=len(dataset)
    parse_json(path, dataset)

if __name__ == "__main__":
    path = os.path.dirname(os.path.realpath(__file__))
    json2annoimg(path)