Convert_between_annotation_format

convert_matterport_json_to_coco.py

# Modified from labelme2coco repository
# Usage: python convert_matterport_json_to_coco.py path_to_json.json

import argparse
import glob
import json
import os

import numpy as np
import PIL.Image


class NpEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        else:
            return super(NpEncoder, self).default(obj)


class labelme2coco(object):
    def __init__(
        self,
        maskrcnn_json=[],
        save_json_path="./coco.json",
        save_category=[],
        height=500,
        width=500,
    ):
        """
        :param maskrcnn_json: the path to matterport 's maskrcnn_json
        :param save_json_path: the path to save new json
        """
        self.maskrcnn_json = maskrcnn_json
        self.save_json_path = save_json_path
        self.save_cat = save_category
        self.images = []
        self.categories = []
        self.annotations = []
        self.label = []
        self.annID = 1
        self.height = height
        self.width = width

        self.save_json()

    def data_transfer(self):
        # for num, json_file in enumerate(self.maskrcnn_json):
        include_shape = ['polyline', 'polygon']
        with open(self.maskrcnn_json, "r") as fp:
            data = json.load(fp)
            i = 0  # image_id
            for d in data:
                self.images.append(self.get_image(data, d, i))
                for region in data[d]["regions"]:
                    if (
                        region is None
                        or "category_id" not in region["region_attributes"]
                        or region["shape_attributes"]["name"] not in include_shape
                    ):
                        continue
                    label = region["region_attributes"]["category_id"]
                    if label not in self.save_cat:
                        continue
                    if label not in self.label:
                        self.label.append(label)
                    all_ptns_x = region["shape_attributes"]["all_points_x"]
                    all_ptns_y = region["shape_attributes"]["all_points_y"]
                    points = []
                    for j in range(len(all_ptns_x)):
                        points.append((all_ptns_x[j], all_ptns_y[j]))
                    # if data[d]["filename"] == "12_028.png":
                    #     import ipdb; ipdb.set_trace();
                    # all_points_x, all_points_y
                    #     points = shapes["points"]
                    self.annotations.append(self.annotation(points, label, i))
                    self.annID += 1

                i += 1

        # Sort all text labels so they are in the same order across data splits.
        self.label.sort()
        for label in self.label:
            self.categories.append(self.category(label))
        for annotation in self.annotations:
            annotation["category_id"] = self.getcatid(annotation["category_id"])

    def get_image(self, data, d, pos):
        image = {}
        img = None
        image["height"] = self.height
        image["width"] = self.width
        image["id"] = pos
        image["file_name"] = data[d]["filename"]

        return image

    def category(self, label):
        category = {}
        category["supercategory"] = label[0]
        category["id"] = len(self.categories)
        category["name"] = label[0]
        return category

    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)))
        annotation["segmentation"] = [list(np.asarray(points).flatten())]
        annotation["iscrowd"] = 0
        annotation["area"] = area
        annotation["image_id"] = num

        annotation["bbox"] = list(map(float, self.getbbox(points)))

        annotation["category_id"] = label[0]  # self.getcatid(label)
        annotation["id"] = self.annID
        return annotation

    def getcatid(self, label):
        for category in self.categories:
            if label == category["name"]:
                return category["id"]
        print("label: {} not in categories: {}.".format(label, self.categories))
        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]

        left_top_r = np.min(rows)  # y
        left_top_c = np.min(clos)  # x

        right_bottom_r = np.max(rows)
        right_bottom_c = np.max(clos)

        return [
            left_top_c,
            left_top_r,
            right_bottom_c - left_top_c,
            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["images"] = self.images
        data_coco["categories"] = self.categories
        data_coco["annotations"] = self.annotations
        return data_coco

    def save_json(self):
        print("save coco json")
        self.data_transfer()
        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
        )
        # import ipdb; ipdb.set_trace();
        json.dump(self.data_coco, open(self.save_json_path, "w"), indent=4, cls=NpEncoder)


if __name__ == "__main__":

    parser = argparse.ArgumentParser(
        description="matterport maskrcnn to coco data json file."
    )
    parser.add_argument(
        "mask_json",
        help="Directory to labelme images and annotation json files.",
        type=str,
    )
    parser.add_argument(
        "--output", help="Output json file path.", default="train_coco_debug1.json"
    )
    args = parser.parse_args()
    # labelme_json = glob.glob(os.path.join(args.mask_json, "*.json"))
    save_cat = ["7", "8"]  # 7: left hand, 8: right hand
    width, height = 1920, 1440
    labelme2coco(args.mask_json, args.output, save_cat, height, width)