Samsara/helper_functions.py at main · DataManagementLab/Samsara · 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
import time
import base64
import ujson

import numpy as np
import cv2 as cv

from pyflink.datastream.functions import MapFunction, ReduceFunction, FilterFunction

### MAP FUNCTIONS ######

class MapDecodeStream(MapFunction):

    def __init__(self, get_annotation):
        super().__init__()

        self.get_annotation = get_annotation

    def map(self, encoded_data):

        start = time.time()

        decoded_data = ujson.loads(encoded_data)

        # Decode the image
        decoded_image = base64.b64decode(decoded_data["frame"])
        np_data = np.fromstring(decoded_image,np.uint8)
        image = cv.imdecode(np_data, cv.IMREAD_UNCHANGED)

        end = time.time()

        return {
            "image" : image,
            "annotation" : self.get_annotation(decoded_data),
            "in" : time.time(),
            "operator_times" : [end-start]
        }

class MapResizeImage(MapFunction):

    def __init__(self, width, height):
        super().__init__()

        self.width = width
        self.height = height

    def map(self, data):
        start = time.time()

        image = cv.resize(data["image"], (self.width, self.height), interpolation=cv.INTER_AREA)

        end = time.time()
        data["operator_times"].append(end-start)

        return {
            "image" : image,
            "annotation" : data["annotation"],
            "in" : data["in"],
            "operator_times" : [end-start]
        }

class MapRecolorImage(MapFunction):

    def map(self, data):
        start = time.time()

        image = cv.cvtColor(data["image"], cv.COLOR_BGR2GRAY)

        end = time.time()
        data["operator_times"].append(end-start)

        return {
            "image" : image,
            "annotation" : data["annotation"],
            "in" : data["in"],
            "operator_times" : [end-start]
        }

class MapChangeFormat(MapFunction):
    def map(self, data):
        return {
            "result" : [data["result"]],
            "annotation" : [data["annotation"]],
            "in_first" : data["in"],
            "in_last" : data["in"],
            "llm_time" : [data["llm_time"]],
            "llm" : data["llm"],
            "operator_times" : data["operator_times"],
            "frames" : 1
        }

###############

### REDUCE FUNCTIONS ######

class ReduceWindowResults(ReduceFunction):

    def reduce(self, value1, value2):
        return {
            "result" : [*value1["result"], *value2["result"]],
            "annotation" : [*value1["annotation"], *value2["annotation"]],
            "in_first" : min(value1["in_first"], value2["in_first"]),
            "in_last" : max(value1["in_last"], value2["in_last"]),
            "llm_time" : [*value1["llm_time"], *value2["llm_time"]],
            "llm" : value1["llm"],
            "operator_times" : [value1["operator_times"][i] + value2["operator_times"][i] for i in range(len(value1["operator_times"]))],
            "frames" : value1["frames"] + value2["frames"]
        }

###############

### FILTER FUNCTIONS ######

class FilterNotEmpty(FilterFunction):
    def filter(self, data):
        return data["result"] != ""

class FilterNotSkipped(FilterFunction):
    def filter(self, data):
        return data["result"].upper() not in ["SKIPPED", "SKIP"]

###############