Unity-Technologies / com.unity.perception

Perception toolkit for sim2real training and validation in Unity
Other
895 stars 172 forks source link

Instance Segmentation annotations with contour points of mask #538

Closed dootiedoot closed 1 year ago

dootiedoot commented 1 year ago

Hi, I am trying to figure out a way to generate array of points along the contour of each object in an instance mask such as example below where every instance is assigned a unique color id:

What I'm looking for is points that run along the contour of the instance mask for example below where the white dots encompass the region defines a couch and another region defining the chair:

Currently, the SOLO output of the dataset includes path to segmentation image, instance id, and label color but no points unlike 2D bounding box XY coordinates unless I am missing this information somewhere.

My goal is to create dataset to train with Yolov5's latest v7.0 release that supports instance segmentation and inputs an array of contour points as ground truth data.

I'm thinking it may be possible to run post-process image difference and generate a polygon for each color or finding where the perception camera captures the instance mask and generating points from the mesh's bounding box?

I can see some big caveats with these approaches, so any advice is appreciated.

StevenBorkman commented 1 year ago

We create our segmentation masks to be pixel perfect. Our philosophy is to create the most accurate data as possible, and then use post processing (generally in python) to generalize the data to meet a user's needs. We have a new python package, pysolotools which allows you to parse a solo dataset and to do some conversions. We include a solo2coco converter that, as a part of the conversion, converts our instance png images to RLE (run length encoded) masks. We don't have any examples of converting to polyline masks, just because they are not as accurate and have to generalize the data too much.

I have found this example (I am not really sure if it works, because I haven't tried it) that will go from RLEs to polygons: https://www.kaggle.com/code/linrds/convert-rle-to-polygons

Anyway, I think your best bet is to post process the data in python, and go from images -> RLEs -> polygons. I hope this helps.

Steve

dootiedoot commented 1 year ago

Hi Steve,

Thank you for the advice, this was exactly what I was looking for and I believe this issue can be closed.

SOLVED: For anybody who is having the same issue, below is a python script I expanded from the kaggle example link that translates Unity's SOLO segmentation RLEs (after coco conversion) into YOLOv5 label files. Segmentation is post-processed with OpenCV's findContours and ConvexHull functions to get a simplified polygon line.

import cv2
import os
import sys
import glob
import pycocotools.mask
import json
import numpy as np
import math

PATH_TO_CONVERTED_COCO_DIR = 'datasets/coco/'

class clockwise_angle_and_distance():
    '''
    A class to tell if point is clockwise from origin or not.
    This helps if one wants to use sorted() on a list of points.

    Parameters
    ----------
    point : ndarray or list, like [x, y]. The point "to where" we g0
    self.origin : ndarray or list, like [x, y]. The center around which we go
    refvec : ndarray or list, like [x, y]. The direction of reference

    use: 
        instantiate with an origin, then call the instance during sort
    reference: 
    https://stackoverflow.com/questions/41855695/sorting-list-of-two-dimensional-coordinates-by-clockwise-angle-using-python

    Returns
    -------
    angle
    distance
    '''
    def __init__(self, origin):
        self.origin = origin

    def __call__(self, point, refvec = [0, 1]):
        if self.origin is None:
            raise NameError("clockwise sorting needs an origin. Please set origin.")
        # Vector between point and the origin: v = p - o
        vector = [point[0]-self.origin[0], point[1]-self.origin[1]]
        # Length of vector: ||v||
        lenvector = np.linalg.norm(vector[0] - vector[1])
        # If length is zero there is no angle
        if lenvector == 0:
            return -math.pi, 0
        # Normalize vector: v/||v||
        normalized = [vector[0]/lenvector, vector[1]/lenvector]
        dotprod  = normalized[0]*refvec[0] + normalized[1]*refvec[1] # x1*x2 + y1*y2
        diffprod = refvec[1]*normalized[0] - refvec[0]*normalized[1] # x1*y2 - y1*x2
        angle = math.atan2(diffprod, dotprod)
        # Negative angles represent counter-clockwise angles so we need to 
        # subtract them from 2*pi (360 degrees)
        if angle < 0:
            return 2*math.pi+angle, lenvector
        # I return first the angle because that's the primary sorting criterium
        # but if two vectors have the same angle then the shorter distance 
        # should come first.
        return angle, lenvector

#   converts RLE data to polygon array
def polygonFromMask(maskedArr, image_filepath): # https://www.kaggle.com/code/linrds/convert-rle-to-polygons/notebook

    contours, _ = cv2.findContours(maskedArr, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

    #   merge contours found using method in https://stackoverflow.com/questions/44501723/how-to-merge-contours-in-opencv
    #   get a list of points of each contour
    list_of_pts = [] 
    for ctr in contours:
        list_of_pts += [pt[0] for pt in ctr]

    #   order points clockwise
    center_pt = np.array(list_of_pts).mean(axis = 0) # get origin
    clock_ang_dist = clockwise_angle_and_distance(center_pt) # set origin
    list_of_pts = sorted(list_of_pts, key=clock_ang_dist) # use to sort

    #   force a list of points into cv2 format and merge them with cv2.convexHull instead
    countour = np.array(list_of_pts).reshape((-1,1,2)).astype(np.int32)
    countour = cv2.convexHull(countour)
    #print(countour)

    #   [DEBUGGING] Display contour for debugging
    # image = cv2.imread(image_filepath)
    # image = cv2.polylines(image, [countour], isClosed = True, color = (255, 0, 0), thickness = 4)
    # while(1):      
    #     cv2.imshow('image', image)
    #     if cv2.waitKey(20) & 0xFF == 27:
    #         break

    return countour

#   converts polygon array to json
def write_polygons_to_yolo_format(json_file, save_dir):
    #   file base cases
    if not os.path.exists(PATH_TO_CONVERTED_COCO_DIR):
        print("coco directory does not exists: '{}'".format(PATH_TO_CONVERTED_COCO_DIR), file=sys.stderr)
        return
    labels_file_path = os.path.join(PATH_TO_CONVERTED_COCO_DIR, 'labels')
    if not os.path.exists(labels_file_path):
        os.makedirs(labels_file_path)

    #   clear any left over files
    files_to_delete = glob.glob(os.path.join(labels_file_path, '*'))
    for f in files_to_delete:
        os.remove(f)

    #   load annotations JSON, parse into polygons, then write polygons to file in YOLO format
    with open(json_file) as f:
        imgs_anns = json.load(f)
        for i, annotation in enumerate(imgs_anns["annotations"]):
            print('Processing annotation {}/{}'.format(i, len(imgs_anns["annotations"])))

            #   initialize variables
            image_id = annotation["image_id"]
            image_filepath = os.path.join(PATH_TO_CONVERTED_COCO_DIR, "images", "camera_{}.png".format(image_id))
            category_id = annotation["category_id"]
            rle = annotation["segmentation"]
            height = rle.get('size')[0]
            width = rle.get('size')[1]
            #compressed_rle = mask_util.frPyObjects(rle, rle.get('size')[0], rle.get('size')[1])
            #mymask = mask_util.decode(compressed_rle)
            binary_mask_ground_truth = pycocotools.mask.decode(rle)

            polygons = polygonFromMask(binary_mask_ground_truth, image_filepath)

            #   post process polygons so their scaled by decimal offset from the image rather then perfect pixel coordinates
            polygonsCOCO = []
            for i, coordinate in enumerate(polygons):
                x, y = coordinate[0]
                polygonsCOCO.append(x / width)
                polygonsCOCO.append(y / height)
            #print(polygonsCOCO)

            #   clear file then write to file
            file_name = 'camera_{}.txt'.format(image_id)
            file_path = os.path.join(labels_file_path, file_name)
            with open(file_path, 'a') as file:
                polygonStr = ' '.join(map(str, polygonsCOCO))
                file.write('{} {}\n'.format(category_id, polygonStr))

def main():
    # Opening JSON file
    json_path = os.path.join(PATH_TO_CONVERTED_COCO_DIR, 'instances.json')
    write_polygons_to_yolo_format(json_path, 'annoation_results.json')

if __name__ == "__main__":
    main()
StevenBorkman commented 1 year ago

Awesome. And thanks for the updated code. Hopefully others in the community will benefit from it.