glp-92 / Prueba_Tecnica_ML_DL

0 stars 0 forks source link

test #2

Open glp-92 opened 3 weeks ago

glp-92 commented 3 weeks ago

test

glp-92 commented 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)