Brush Mask for semantic segmentation Export to COCO

[lackray]

Hello. I have a problem with exporting the mask to coco. I got the json file from labelstudio that had the rle format. I tried decoding it with all the open source codes and nothing works.

After decoding and exporting it, the coco viewer result is not correct.

What is the correct way of changing the RLE to mask?

Thanks in advance

@hlomzik

[chanderlud]

I faced the same issue as you and put together this code to solve it. This uses code copied from the label studio conversion tool and some StackOverflow threads. You will need to implement some further logic to get the exact coco format you want.

usage

# result and value are from the label studio JSON format
binary_mask = rle_to_mask(value['rle'], result['original_height'], result['original_width'])
# coco uses polygons for segmentation
polygons = mask_to_polygon(binary_mask)

main code

import json
from typing import List

import cv2
import numpy as np

import matplotlib.pyplot as plt
import matplotlib.patches as patches
from pycocotools.mask import encode


def is_clockwise(contour):
    value = 0
    num = len(contour)
    for i, point in enumerate(contour):
        p1 = contour[i]
        if i < num - 1:
            p2 = contour[i + 1]
        else:
            p2 = contour[0]
        value += (p2[0][0] - p1[0][0]) * (p2[0][1] + p1[0][1])
    return value < 0


def get_merge_point_idx(contour1, contour2):
    idx1 = 0
    idx2 = 0
    distance_min = -1
    for i, p1 in enumerate(contour1):
        for j, p2 in enumerate(contour2):
            distance = pow(p2[0][0] - p1[0][0], 2) + pow(p2[0][1] - p1[0][1], 2)
            if distance_min < 0:
                distance_min = distance
                idx1 = i
                idx2 = j
            elif distance < distance_min:
                distance_min = distance
                idx1 = i
                idx2 = j
    return idx1, idx2


def merge_contours(contour1, contour2, idx1, idx2):
    contour = []
    for i in list(range(0, idx1 + 1)):
        contour.append(contour1[i])
    for i in list(range(idx2, len(contour2))):
        contour.append(contour2[i])
    for i in list(range(0, idx2 + 1)):
        contour.append(contour2[i])
    for i in list(range(idx1, len(contour1))):
        contour.append(contour1[i])
    contour = np.array(contour)
    return contour


def merge_with_parent(contour_parent, contour):
    if not is_clockwise(contour_parent):
        contour_parent = contour_parent[::-1]
    if is_clockwise(contour):
        contour = contour[::-1]
    idx1, idx2 = get_merge_point_idx(contour_parent, contour)
    return merge_contours(contour_parent, contour, idx1, idx2)


def mask_to_polygon(image):
    contours, hierarchies = cv2.findContours(image, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_TC89_KCOS)
    if len(contours) == 0:
        return []
    contours_approx = []
    for contour in contours:
        epsilon = 0.001 * cv2.arcLength(contour, True)
        contour_approx = cv2.approxPolyDP(contour, epsilon, True)
        contours_approx.append(contour_approx)
    contours_parent = []
    for i, contour in enumerate(contours_approx):
        parent_idx = hierarchies[0][i][3]
        if parent_idx < 0 and len(contour) >= 3:
            contours_parent.append(contour)
        else:
            contours_parent.append([])
    for i, contour in enumerate(contours_approx):
        parent_idx = hierarchies[0][i][3]
        if parent_idx >= 0 and len(contour) >= 3:
            contour_parent = contours_parent[parent_idx]
            if len(contour_parent) == 0:
                continue
            contours_parent[parent_idx] = merge_with_parent(contour_parent, contour)
    contours_parent_tmp = []
    for contour in contours_parent:
        if len(contour) == 0:
            continue
        contours_parent_tmp.append(contour)
    polygons = []
    for contour in contours_parent_tmp:
        polygon = contour.flatten().tolist()
        polygons.append(polygon)
    return polygons


class InputStream:
    def __init__(self, data):
        self.data = data
        self.i = 0

    def read(self, size):
        out = self.data[self.i:self.i + size]
        self.i += size
        return int(out, 2)


def access_bit(data, num):
    """ from bytes array to bits by num position"""
    base = int(num // 8)
    shift = 7 - int(num % 8)
    return (data[base] & (1 << shift)) >> shift


def bytes2bit(data):
    """ get bit string from bytes data"""
    return ''.join([str(access_bit(data, i)) for i in range(len(data) * 8)])


def rle_to_mask(rle: List[int], height: int, width: int) -> np.array:
    """
    Converts rle to image mask
    Args:
        rle: your long rle
        height: original_height
        width: original_width

    Returns: np.array
    """

    rle_input = InputStream(bytes2bit(rle))

    num = rle_input.read(32)
    word_size = rle_input.read(5) + 1
    rle_sizes = [rle_input.read(4) + 1 for _ in range(4)]
    # print('RLE params:', num, 'values,', word_size, 'word_size,', rle_sizes, 'rle_sizes')

    i = 0
    out = np.zeros(num, dtype=np.uint8)
    while i < num:
        x = rle_input.read(1)
        j = i + 1 + rle_input.read(rle_sizes[rle_input.read(2)])
        if x:
            val = rle_input.read(word_size)
            out[i:j] = val
            i = j
        else:
            while i < j:
                val = rle_input.read(word_size)
                out[i] = val
                i += 1

    image = np.reshape(out, [height, width, 4])[:, :, 3]
    return image

results visualization

def visualize_conversion(binary_mask, polygons):
    import PIL.Image

    pil_image = PIL.Image.fromarray(binary_mask)
    pil_image.show()

    fig, ax = plt.subplots()
    ax.set_xlim(0, 640)
    ax.set_ylim(0, 640)
    ax.invert_yaxis()

    # Convert and draw polygons
    for polygon in polygons:
        # Convert interleaved list to numpy array of shape (n, 2)
        points = np.array(polygon, dtype=np.int32).reshape(-1, 2)
        poly_patch = patches.Polygon(
            points,
            closed=True,
            edgecolor='white',
            fill=None,
            linewidth=2)
        ax.add_patch(poly_patch)

        for i, point in enumerate(points):
            plt.plot(point[0], point[1], 'ro', markersize=2)

        # Display the image
        plt.gca().set_facecolor('black')
        plt.show()

[lackray]

thx. this code worked for me.

[BadCoder2]

Once you get the polygons, how exactly do you convert that into a COCO json file? I was able to get the rest of the code working but couldn't figure that part out

[chanderlud]

As I mentioned in my post, you will need to implement further logic to output COCO-style annotations. The labelstudio converter codebase is a good thing to check out if you need help.