cv-symbol-extraction/convertYoloToFormula.py at master · TNodeCode/cv-symbol-extraction · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
import yaml
import torch
from pathlib import Path
from torchmetrics import JaccardIndex
import numpy as np


def overlap(boxes1, boxes2):
    # top left corners of all combinations
    lt = torch.max(boxes1[:, None, :2], boxes2[:, :2])  # [N,M,2]

    # bottom right corners of all combinations
    rb = torch.min(boxes1[:, None, 2:], boxes2[:, 2:])  # [N,M,2]

    # width and hight of overlap area of boxes1 and boxes2.
    wh = (rb - lt).clamp(min=0)  # [N,M,2]
    inter = wh[:, :, 0] * wh[:, :, 1]  # [N,M]#

    boxes1Area = (boxes1[:, 2] - boxes1[:, 0]) * (boxes1[:, 3] - boxes1[:, 1])
    boxes2Area = (boxes2[:, 2] - boxes2[:, 0]) * (boxes2[:, 3] - boxes2[:, 1])

    iou = inter / (boxes1Area + boxes2Area - inter)

    return iou


def coordinatesConverter(bboxes, img_height, img_width):
    """
    Converts BBoxes from xywh (standardized) in xyxy
    """
    bboxes[:, 0] *= 2 * img_width
    bboxes[:, 1] *= 2 * img_height

    bboxes[:, 2] *= img_width
    bboxes[:, 3] *= img_height

    bboxes[:, 0] = (bboxes[:, 0] - bboxes[:, 2]) / 2
    bboxes[:, 1] = (bboxes[:, 1] - bboxes[:, 3]) / 2
    bboxes[:, 2] = bboxes[:, 2] + bboxes[:, 0]
    bboxes[:, 3] = bboxes[:, 3] + bboxes[:, 1]
    return bboxes


def convertYoloToSingleFormula(txt_file_name):
    # replace with Path to custom_data.yaml for encoding
    #with open("C:/Users/benwe/Downloads/custom_data.yaml", "r") as file:
    #    data = yaml.safe_load(file)["names"]
    #    labels_encode = {}
    #    for i in range(len(data)):
    #        labels_encode[data[i]] = i

    fomrula_code = 0

    yolo_output = Path("detections")
    if not yolo_output.exists():
        yolo_output.mkdir()
    yolo_labels_path = yolo_output / "labels"

    formulaLabels_path = yolo_output / "formulaLabels" / txt_file_name[:-4]
    if not formulaLabels_path.exists():
        formulaLabels_path.mkdir(parents=True)

    yolo_labels = [file for file in yolo_labels_path.glob("*.txt") if file.stem == txt_file_name[:-4]]
    output_counter = 0
    for file in yolo_labels:
        labels = open(file, "r").readlines()
        labels = list(map(lambda x: torch.Tensor(
            list(map(float, x.strip('\n').split(" ")))), labels))
        formulas_in_image = [box for box in labels if box[0] == fomrula_code]
        if len(formulas_in_image) == 2:
            first = formulas_in_image[0][1:]
            second = formulas_in_image[1][1:]
            # convert first and second to numpy arrays and to cpu
            first = first.cpu().numpy()
            second = second.cpu().numpy()
            intersection = np.logical_and(first, second)

            union = np.logical_or(first, second)

            iou_score = np.sum(intersection) / np.sum(union)

            # if iou_score > 0.95, then delete the box thats inside the other
            if iou_score > 0.95:
                if formulas_in_image[0][2] > formulas_in_image[1][2]:
                    formulas_in_image.pop(1)
                else:
                    formulas_in_image.pop(0)


        boxes2 = torch.stack(labels)[:, 1:]
        boxes2 = coordinatesConverter(boxes2, 640, 640)

        for formula in formulas_in_image:
            boxes1 = formula[1:].clone().reshape((1, 4))
            boxes1 = coordinatesConverter(boxes1, 640, 640)

            iou = overlap(boxes1, boxes2)
            _, boxes_in_formula_idx = torch.where(iou > 0)

            boxes_in_formula = torch.stack(
                labels)[boxes_in_formula_idx].tolist()
            boxes_in_formula = [[int(box[0]), *box[1:]]
                                for box in boxes_in_formula if box[0] != fomrula_code]

            label_txt = open(formulaLabels_path / f"{output_counter}.txt", 'a')
            label_txt.write(f"{formula[0]} {formula[1]} {formula[2]} {formula[3]} {formula[4]}\n")

            for line_ in boxes_in_formula:
                label_txt.write(
                    f"{line_[0]} {line_[1]} {line_[2]} {line_[3]} {line_[4]}")
                label_txt.write("\n")
            label_txt.close()
            output_counter += 1

            # # For visualization
            # print(boxes_in_formula)
            # for i in boxes_in_formula_idx:
            #     box = boxes2[i]
            #     plt.gca().add_patch(Rectangle((box[0], box[1]), box[2]-box[0], box[3]-box[1], linewidth=1, edgecolor='g', facecolor='none'))

            # plt.gca().add_patch(Rectangle((boxes1[0,0], boxes1[0,1]), boxes1[0,2]-boxes1[0,0], boxes1[0,3]-boxes1[0,1], linewidth=1, edgecolor='b', facecolor='none'))
            # img = torch.zeros((640, 640))
            # plt.imshow(img)
            # plt.show()