robot-memory-visualization/flicker_remover.py at main · Mersen130/robot-memory-visualization · 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
from utilities import parse_box

class Flicker_Remover:
    def __init__(self) -> None:
        self.known_objs = set()
        self.last_frame_objs = set()
        self.last_frame_boxes = []
        self.second_last_frame_boxes = []
        self.threshold = 0.5

        self.flickered_objid_mapping = {}

    def update(self, boxes):
        # print(self.flickered_objid_mapping)
        self.apply_mapping(boxes)

        curr_objs, _, curr_obj2box = parse_box(boxes)

        for obj_id in curr_objs.difference(self.known_objs):  # objs enter
            self.known_objs.add(obj_id)

            similar_box = self.get_similar_boxes(boxes, curr_obj2box[obj_id])
            if similar_box:
                self.flickered_objid_mapping[obj_id] = similar_box
                boxes.remove(curr_obj2box[obj_id])
                continue

            similar_box = self.get_similar_boxes(self.last_frame_boxes, curr_obj2box[obj_id])
            if similar_box:
                self.flickered_objid_mapping[obj_id] = similar_box
                continue

            similar_box = self.get_similar_boxes(self.second_last_frame_boxes, curr_obj2box[obj_id])
            if similar_box:
                self.flickered_objid_mapping[obj_id] = similar_box

        # for obj_id in self.last_frame_objs.difference(curr_objs):  # objs left
        #     pass
        self.apply_mapping(boxes)

        self.last_frame_boxes, self.second_last_frame_boxes = boxes, self.last_frame_boxes
        return boxes

    def update_human(self, boxes_before_human, boxes_after_human):
        # compare boxes before and after human event
        objs_after, _, objs_after2box = parse_box(boxes_after_human)
        objs_before, _, objs_before2box = parse_box(boxes_before_human)

        for obj_id in objs_after.difference(objs_before):
            similar_box = self.get_similar_boxes(boxes_before_human, objs_after2box[obj_id])
            if similar_box:
                self.flickered_objid_mapping[obj_id] = similar_box
                continue

        self.apply_mapping(boxes_after_human)
        return boxes_after_human

    def apply_mapping(self, boxes):
        for i in range(len(boxes)):
            obj_id = int(boxes[i][4])
            if obj_id in self.flickered_objid_mapping:
                boxes[i][4] = self.flickered_objid_mapping[obj_id][4]

    def get_similar_boxes(self, boxes, target):
        """find the box that's similar to target in boxes, they are potentially same object"""
        t_x1, t_y1, t_x2, t_y2, t_obj_id, t_cls_id = \
            int(target[0]), int(target[1]), int(target[2]), int(target[3]), int(target[4]), int(target[5])
        b1 = (t_x1, t_y1, t_x2, t_y2)

        for box in boxes:
            x1, y1, x2, y2, obj_id, cls_id = \
                int(box[0]), int(box[1]), int(box[2]), int(box[3]), int(box[4]), int(box[5])
            b2 = (x1, y1, x2, y2)
            if b2 == b1:
                continue

            iou = self.bb_intersection_over_union(b1, b2)

            if iou >= self.threshold:
                return box


        return None

    def bb_intersection_over_union(self, boxA, boxB):
        # determine the (x, y)-coordinates of the intersection rectangle
        xA = max(boxA[0], boxB[0])
        yA = max(boxA[1], boxB[1])
        xB = min(boxA[2], boxB[2])
        yB = min(boxA[3], boxB[3])
        # compute the area of intersection rectangle
        interArea = max(0, xB - xA + 1) * max(0, yB - yA + 1)
        # compute the area of both the prediction and ground-truth
        # rectangles
        boxAArea = (boxA[2] - boxA[0] + 1) * (boxA[3] - boxA[1] + 1)
        boxBArea = (boxB[2] - boxB[0] + 1) * (boxB[3] - boxB[1] + 1)
        # compute the intersection over union by taking the intersection
        # area and dividing it by the sum of prediction + ground-truth
        # areas - the interesection area
        iou = interArea / float(boxAArea + boxBArea - interArea)
        # return the intersection over union value
        return iou

flicker_remover = Flicker_Remover()