-
Notifications
You must be signed in to change notification settings - Fork 2.1k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Brush Mask for semantic segmentation Export to COCO #5899
Comments
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 codeimport 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 visualizationdef 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() |
thx. this code worked for me. |
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 |
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. |
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
The text was updated successfully, but these errors were encountered: