泳池项目实验记录（三）

发表于 | 分类于 | 阅读次数

本周完成了泳池内人头与人体的匹配算法，即：匈牙利匹配算法 + IOU

代码如下：

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
def to_matching(outputs):
    head_list = []
    body_list = []
    head_indices = []
    body_indices = []
    person_list = []
    if len(outputs) > 0:
        for index, output in enumerate(outputs):
            id = int(output[4])
            cls = int(output[5])
            xmin, ymin, xmax, ymax = int(output[0]), int(output[1]), int(output[2]), int(output[3])
            # 坐标按顺时针
            bbox = [[xmin, ymin],
                    [xmax, ymin],
                    [xmax, ymax],
                    [xmin, ymax]]

            # ---datasets: head  body_overwater  body_underwater
            # if c == 0:
            #     head_list.append({"index": index, "id": id, "class": cls, "bbox": bbox})
            #     head_indices.append(index)
            # elif c == 1 or c == 2:
            #     body_list.append({"index": index, "id": id, "class": cls, "bbox": bbox})
            #     body_indices.append(index)

            # ---datasets: person  overwater underwater
            if cls == 0:
                body_list.append({"index": index, "id": id, "class": cls, "bbox": bbox})
            elif cls == 1 or cls == 2:
                head_list.append({"index": index, "id": id, "class": cls, "bbox": bbox})

        iou_matrix = np.zeros((len(head_list), len(body_list)))

        for i, head_points in enumerate(head_list):
            for j, body_points in enumerate(body_list):
                iou_matrix[i, j] = - intersection(head_points["bbox"], body_points["bbox"])
                # iou_matrix[i, j] = 1 - IOU(head_points, body_points)

        matched_row_indices, matched_col_indices = linear_sum_assignment(iou_matrix)
        iou_matrix = - iou_matrix

        idx = iou_matrix[matched_row_indices, matched_col_indices] > 0.95
        matched_row_indices = matched_row_indices[idx]
        matched_col_indices = matched_col_indices[idx]

        # 添加匹配上的 head 和 body
        for row_ind, col_ind in zip(matched_row_indices, matched_col_indices):
            head_indices.append(head_list[row_ind]["index"])
            body_indices.append(body_list[col_ind]["index"])
            head_list[row_ind]["bbox"] = to_center_with_corners(head_list[row_ind]["bbox"])
            body_list[col_ind]["bbox"] = to_center_with_corners(body_list[col_ind]["bbox"])
            person_list.append({"type": "person", "id": body_list[col_ind]["id"], "class": head_list[row_ind]["class"],
                                "bbox": [head_list[row_ind]["bbox"], body_list[col_ind]["bbox"]]})

        # 添加未匹配上的 head 或 body
        for index, output in enumerate(outputs):
            if index not in head_indices and index not in body_indices:
                id = int(output[4])
                cls = int(output[5])
                xmin, ymin, xmax, ymax = int(output[0]), int(output[1]), int(output[2]), int(output[3])
                bbox = to_center_with_corner([xmin, ymin, xmax, ymax])

                # ---datasets: head  body_overwater  body_underwater
                # if cls == 0:
                #     person_list.append({"type": "head", "id": id, "bbox": bbox})
                # else:
                #     person_list.append({"type": "body", "id": id, "bbox": bbox})

                # ---datasets: person  overwater underwater
                if cls == 0:
                    person_list.append({"type": "body", "id": id, "class": cls, "bbox": bbox})
                else:
                    person_list.append({"type": "head", "id": id, "class": cls, "bbox": bbox})

    return person_list
0%