Support YOLO Ultralytics segmentation

[hugo082]

Hi,

It would be nice to have support for YOLO Ultralytics segmentation format. Currently when we export a mask via CVAT and YOLO 1.1 format, which use datumaro under the hood, the masks are ignored.

The segmentation format is pretty similar to the bbox one (with more points) and I identified the code responsible of the convertion.

I can take time to propose a PR, do you think it is possible to merge only the export part?

I already develop some code to convert Label studio brush labels to YOLO segmentation, and it seems we can access the bitwise 2d mask from datumaro.annotation.Mask so the code should stay pretty similar to the following:

def brush_label_annotation_to_yolo_polygon(annotation):
    """
    Brush label annotation result to YOLO segment format
    https://docs.ultralytics.com/datasets/segment/
    """
    if annotation['type'] != 'brushlabels':
        raise ValueError(
            f"Annotation type {annotation['type']} not supported.")

    width = annotation['original_width']
    height = annotation['original_height']
    rle_mask = annotation['value']['rle']

    # convert to bitwise mask 1d
    mask = brush.decode_rle(rle_mask)

    # reshape to 2d
    mask = np.reshape(mask, [height, width, 4])[:, :, 3]

    # here mask === datumaro.annotation.Mask.image

    # Find the contours of the mask
    contours, hierarchy = cv2.findContours(
        mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
    )

    # Get the new bounding box and segmentation mask
    largest_contour = max(contours, key=cv2.contourArea)
    # bbox = [int(x) for x in cv2.boundingRect(largest_contour)]
    segment_mask = np.array(largest_contour.flatten().tolist()).reshape(-1, 2)

    return segment_mask_to_yolo_polygon(segment_mask, width=width, height=height)

def segment_mask_to_yolo_polygon(segment_mask, width: int, height: int):
    """
    From segment mask [[x1, y1], [x2, y2], ...] to YOLO polygon format
    """
    # convert mask to numpy array of shape (N,2)
    mask = np.array(segment_mask).reshape(-1, 2)

    # normalize the pixel coordinates
    mask_norm = mask / np.array([width, height])

    # [[x, y], ...] to [x, y, ...]
    mask_1d = mask_norm.reshape(-1)

    return "".join(["{:.6f} ".format(value) for value in mask_1d])

[wonjuleee]

Hi @hugo082, sorry for the late response, sure, we are always open to add more features. Please feel free to create the PR. This will be helpful to be consistent to YOLO-Ultralytics.

[ryouchinsa]

Using the script general_json2yolo.py, you can convert the RLE mask with holes to the YOLO segmentation format.

The RLE mask is converted to a parent polygon and a child polygon using cv2.findContours(). The parent polygon points are sorted in clockwise order. The child polygon points are sorted in counterclockwise order. Detect the nearest point in the parent polygon and in the child polygon. Connect those 2 points with narrow 2 lines. So that the polygon with a hole is saved in the YOLO segmentation format.

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 mask2polygon(image):
    contours, hierarchies = cv2.findContours(image, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_TC89_KCOS)
    contours_approx = []
    polygons = []
    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 

def rle2polygon(segmentation):
    if isinstance(segmentation["counts"], list):
        segmentation = mask.frPyObjects(segmentation, *segmentation["size"])
    m = mask.decode(segmentation) 
    m[m > 0] = 255
    polygons = mask2polygon(m)
    return polygons

The RLE mask.

The converted YOLO segmentation format.

To run the script, put the COCO JSON file coco_train.json into datasets/coco/annotations. Run the script. python general_json2yolo.py The converted YOLO txt files are saved in new_dir/labels/coco_train.

Edit use_segments and use_keypoints in the script.

if __name__ == '__main__':
    source = 'COCO'

    if source == 'COCO':
        convert_coco_json('../datasets/coco/annotations',  # directory with *.json
                          use_segments=True,
                          use_keypoints=False,
                          cls91to80=False)

To convert the COCO bbox format to YOLO bbox format.

use_segments=False,
use_keypoints=False,

To convert the COCO segmentation format to YOLO segmentation format.

use_segments=True,
use_keypoints=False,

To convert the COCO keypoints format to YOLO keypoints format.

use_segments=False,
use_keypoints=True,

This script originates from Ultralytics JSON2YOLO repository. We hope this script would help your business.

[CourchesneA]

was a PR ever made out of this ? It would be useful to have conversion from datumaro to yolo_ultralytics for segmentation