自定义标注格式转coco

#！-*- coding:utf-8 -*-
import os
import sys
import json
import cv2
import numpy as np
import shutil
import io
import base64
import PIL.Image
import PIL.ImageDraw
import argparse


def mkdir_os(path):
    if not os.path.exists(path):
        os.makedirs(path)


class labelme2coco(object):
    def __init__(self, save_json_path="./coco.json", all_categories=None):
        """
        :param save_json_path: the path to save new json
        """
        self.save_json_path = save_json_path
        self.images = []
        if all_categories is not None:
            self.categories = all_categories
        else:
            print("预先设置所有的categories")
            exit()
        self.annotations = []
        self.annID = 0
        self.height = 0
        self.width = 0

    def data_transfer(self, input_data, num):

        self.images.append(self.image(input_data, num))
        for shapes in input_data["shapes"]:
            label = shapes["label"]
            points = shapes["points"]
            self.annotations.append(self.annotation(points, label, num))
            self.annID += 1

    def image(self, data, num):
        image = {}
        #img = utils.img_b64_to_arr(data["imageData"])
        #height, width = img.shape[:2]
        #img = None
        image["height"] = data["imageHeight"]
        image["width"] = data["imageWidth"]
        image["id"] = num
        image["file_name"] = data["imagePath"].split("/")[-1]

        self.height = data["imageHeight"]
        self.width = data["imageWidth"]

        return image

    def annotation(self, points, label, num):
        annotation = {}
        #contour = np.array(points)
        #x = contour[:, 0]
        #y = contour[:, 1]
        #area = 0.5 * np.abs(np.dot(x, np.roll(y, 1)) - np.dot(y, np.roll(x, 1)))
        #TypeError: Object of type 'int64' is not JSON serializable
        temp = [list(np.asarray(points).flatten())]
        annotation["segmentation"] = [list(map(int, key)) for key in temp]
        annotation["iscrowd"] = 0
        annotation["ignore"] = 0
        #annotation["area"] = area
        annotation["image_id"] = num

        # 如果是bbox，不走计算bbox的逻辑
        if len(points)==2:
            x,y,w,h = points[0][0], points[0][1], (points[1][0]-points[0][0]), (points[1][1]-points[0][1])
            annotation["bbox"] = [round(x,2), round(y,2), round(w,2), round(h,2)]
        else:
            annotation["bbox"] = list(map(float, self.getbbox(points)))

        annotation["area"] = int(annotation["bbox"][2]*annotation["bbox"][3])

        annotation["category_id"] = self.getcatid(label, num)
        annotation["id"] = self.annID
        return annotation

    def getcatid(self, label, img_id):
        for category in self.categories:
            if label == category["name"]:
                return category["id"]
        print("label: {} not in categories: {}.".format(label, self.categories))
        print(self.images[img_id]["file_name"])
        exit()
        return -1

    def getbbox(self, points):
        polygons = points
        mask = self.polygons_to_mask([self.height, self.width], polygons)
        return self.mask2box(mask)

    def mask2box(self, mask):

        index = np.argwhere(mask == 1)
        rows = index[:, 0]
        clos = index[:, 1]

        # if rows.size==0:
        #     left_top_r=0
        #     right_bottom_r=0
        # else:
        #     left_top_r = np.min(rows)  # y
        #     right_bottom_r = np.max(rows)

        # if clos.size==0:
        #     left_top_c=0
        #     right_bottom_c=0
        # else:
        #     left_top_c = np.min(clos)  # x
        #     right_bottom_c = np.max(clos)

        left_top_r = np.min(rows)
        left_top_c = np.min(clos)
        
        right_bottom_r = np.max(rows)
        right_bottom_c = np.max(clos)
        
        #TypeError: Object of type 'int64' is not JSON serializable
        return [
            int(left_top_c),
            int(left_top_r),
            int(right_bottom_c - left_top_c),
            int(right_bottom_r - left_top_r),
        ]

    def polygons_to_mask(self, img_shape, polygons):
        mask = np.zeros(img_shape, dtype=np.uint8)
        mask = PIL.Image.fromarray(mask)
        xy = list(map(tuple, polygons))
        PIL.ImageDraw.Draw(mask).polygon(xy=xy, outline=1, fill=1)
        mask = np.array(mask, dtype=bool)
        return mask

    def data2coco(self):
        data_coco = {}
        data_coco["info"] = 'spytensor created'
        data_coco["license"] = ['license']
        data_coco["categories"] = self.categories
        data_coco["images"] = self.images
        data_coco["annotations"] = self.annotations
        return data_coco

    def save_json(self):
        print("save coco json")
        self.data_coco = self.data2coco()

        print(self.save_json_path)
        os.makedirs(
            os.path.dirname(os.path.abspath(self.save_json_path)), exist_ok=True
        )
        #mkdir_os(os.path.dirname(os.path.abspath(self.save_json_path)))

        json.dump(self.data_coco, open(self.save_json_path, "w", encoding="utf-8"), ensure_ascii=False, indent=1)

        # import io
        # with io.open(self.save_json_path, 'w', encoding="utf-8") as outfile:
        #     my_json_str = json.dumps(self.data_coco, ensure_ascii=False, indent=1)
        #     outfile.write(my_json_str)


if __name__ == "__main__":

    dg_json_path = ".json"

    image_path = "ori_img"

    coco_json_path = ".json"


    # update更新的地方，需要更新这次的label
    all_categories = [
        {
            "name": "bolt",
            "id": 1
        }
    ]
    object_coco = labelme2coco(coco_json_path, all_categories)

    ori_suffix = [m.split('.')[-1] for m in os.listdir(image_path)]
    from collections import Counter
    count = Counter(ori_suffix)
    suffix = count.most_common(1)[0][0]

    data = []
    with open(dg_json_path) as f:
        for line in f:
            data.append(json.loads(line))

    print("原始数据的后缀是： ", suffix)
    for dg_ind, dg_val in enumerate(data):

        if dg_ind%50==0:
            print(dg_ind,'/',len(data))

        filename = dg_val["url_image"]
        filename = os.path.basename(filename).replace('jpg', suffix)


        img_arr = cv2.imread(os.path.join(image_path, filename), -1)
        if img_arr is None:
            print('img_arr is None:', filename)
            continue
    
        height, width = img_arr.shape[:2]

        if len(dg_val["result"])==0:
            print('len(onedate["result"])==0: \n', filename)
            continue

        if 'data' in dg_val["result"] or 'data' in dg_val["result"][0]:
                
                data_ann = dg_val["result"]
                points_list = []
                for index in range(len(data_ann)):
                    ann = data_ann[index]
                    shape = {}

                    # update更新的地方，需要更新这次的label
                    if ann['tagtype'] == 'bolt':
                        shape['label'] = 'bolt'
                    else:
                        print("tagtype 不是此次标注类型")
                        exit()

                    shape['points'] = []
                    if len(ann['data'])==4:
                        x1=min(ann['data'][0],ann['data'][2])
                        y1=min(ann['data'][1],ann['data'][3])
                        x2=max(ann['data'][0],ann['data'][2])
                        y2=max(ann['data'][1],ann['data'][3])
                        
                        if x2-x1 < 5 or y2-y1 < 5:
                            print(ann['data'], "数据标注错误", "\n")
                            continue

                        shape['points'].append([x1, y1])
                        shape['points'].append([x2, y2])

                        points_list.append(shape)
                    elif len(ann['data'])>4:
                        print("shape_type polygon 不是此次任务范围")
                        exit()

                result_dict = {
                    'imageHeight': height,
                    'imageWidth': width,
                    'imagePath': filename,
                    'shapes': points_list,
                }

                object_coco.data_transfer(result_dict, dg_ind)

    object_coco.save_json()