Open glp-92 opened 3 weeks ago
import numpy as np
import cv2
def convert_yolo_label_to_mask(image_file_path, label_file_path):
original_image = cv2.imread(image_file_path, 0)
h, w = original_image.shape
mask = np.zeros((h, w), np.uint8)
with open(label_file_path, 'r') as f:
for line in map(lambda x: x.rsplit(), f.readlines()):
x_points = list(map(lambda x: int(float(x) * w), line[1::2]))
y_points = list(map(lambda y: int(float(y) * h), line[2::2]))
pts = np.array(list(zip(x_points, y_points)),np.int32).reshape((-1, 1, 2))
# cv2.polylines(mask, [pts], True, 255, 1) # Not filling
cv2.fillPoly(mask, [pts], 255)
return mask
im = convert_yolo_label_to_mask('1_040422.jpg', '1_040422.txt')
cv2.imwrite("mask.png", im)
import cv2
import numpy as np
def contours_join(contour_parent, contour_child):
"""
Join parent contour with child contour
Donut use case. Inside donut shouldn't detect anything
"""
def is_clockwise(contour):
value = 0
num = len(contour)
for i in range(len(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, dtype=np.int32)
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)
return merge_with_parent(contour_parent=contour_parent, contour=contour_child)
def group_child_contours_with_father(hierarchy):
"""
hierarchy: [next, previous, first child, parent]
returns:
{
father_key: {
"father": father_key,
"child": [child_keys]
}
}
"""
groups = {}
for i, h in enumerate(hierarchy.squeeze()):
father_index = h[3]
if father_index != -1:
if groups.get(father_index) is not None:
groups[father_index]["child"].append(i)
else:
groups[father_index] = {"father": father_index, "child": [i]}
else:
if groups.get(i) is not None:
groups[i]["father"] = i
else:
groups[i] = {"father": i, "child": []}
return groups
def convert_mask_to_yolo_seg_label(mask_path):
mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
height, width = mask.shape
label_str = ""
_, thresh = cv2.threshold(mask, 127, 255, 0)
contours, hierarchy = cv2.findContours(
thresh, cv2.RETR_CCOMP, cv2.CHAIN_APPROX_SIMPLE
) # instead of RETR_TREE
contour_groups = group_child_contours_with_father(hierarchy)
test_mask = np.zeros((height, width), dtype=np.uint8)
for contour_group in contour_groups.values():
label_str += "0" # One class
contour_parent = contours[contour_group["father"]]
for child in contour_group["child"]:
contour_parent = contours_join(
contour_parent=contour_parent, contour_child=contours[child]
)
contour_to_write = (
contour_parent.squeeze()
) # Coords in [[x, y]], [[x, y]] must remove 1 axis
for point in contour_to_write.tolist():
if not isinstance(point, list):
continue
label_str += f" {round(float(point[0]) / float(width), 6)}"
label_str += f" {round(float(point[1])/ float(height), 6)}"
label_str += "\n"
contour_parent = np.expand_dims(contour_parent, axis=0)
cv2.drawContours(test_mask, contour_parent, -1, 255, -1)
label_str = label_str.rstrip() # Remove last \n
return label_str, test_mask
label_str, mask_generated = convert_mask_to_yolo_seg_label(mask_path='image.png')
cv2.imwrite("mask1.png", mask_generated)
with open('mask.txt', "w") as f:
f.write(label_str)
test