squareit/squareit.py at master · deckbsd/squareit · 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
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
import argparse
import wave

MIN = 1
MID = 2
MAX = 3
BYTES = 0

WAV_PCM_FORMATS = {'8': [1, 0, 128, 255], '16': [2, -32268, 0, 32767], '24': [3, -8388608, 0,  8388607], '32': [4, -2147483648, 0, 2147483647]}


def get_frame_value(frame):
    return int.from_bytes(frame, byteorder='little', signed=True)


def signal_analyze(wav_input_file, bits, channels):
    peak_list = [[] for declare in range(channels)]
    last_frame_value = [None] * channels
    current_frame_value = [0] * channels
    sign = [False] * channels
    last_position_from_middle = [None] * channels
    current_position_from_middle = [True] * channels
    total_frames = wav_input_file.getnframes()
    for i in range(total_frames):
        frame = wav_input_file.readframes(1)
        for channel in range(0, channels):
            index = channel * 2
            frame_part = frame[index:index + WAV_PCM_FORMATS[bits][BYTES]]
            current_frame_value[channel] = get_frame_value(frame_part)

            if current_frame_value[channel] >= WAV_PCM_FORMATS[bits][MID]:
                # High
                current_position_from_middle[channel] = True
            else:
                # Low
                current_position_from_middle[channel] = False

            if last_position_from_middle[channel] is None:
                last_position_from_middle[channel] = current_position_from_middle[channel]

            if current_position_from_middle[channel] != last_position_from_middle[channel] or i == (total_frames-1):
                (peak_list[channel]).append(last_frame_value[channel])
                last_frame_value[channel] = None
                last_position_from_middle[channel] = current_position_from_middle[channel]

            if current_frame_value[channel] >= 0:
                # +
                sign[channel] = True
            else:
                # -
                sign[channel] = False

            if sign[channel]:
                if last_frame_value[channel] is None or current_frame_value[channel] > last_frame_value[channel]:
                    last_frame_value[channel] = current_frame_value[channel]
            else:
                if last_frame_value[channel] is None or current_frame_value[channel] < last_frame_value[channel]:
                    last_frame_value[channel] = current_frame_value[channel]

    wav_input_file.rewind()
    return peak_list


if __name__ == "__main__":
    try:
        parser = argparse.ArgumentParser(description='transform sinusoidal signal to square.')
        parser.add_argument('--input', '-i', dest='input_file', required=True,
                            help='the wave file containing the input data. (required)')
        parser.add_argument('--output', '-o', dest='output_file', required=True,
                            help='the file where to put the transformation result. (required)')
        parser.add_argument('--bits', '-b', dest='bits', required=True,
                            help='The number of bits per sample')
        parser.add_argument('--amplification', '-a', dest='amp', default=False,
                            help='Do an amplification of the signal. (no by default)')
        parser.add_argument('--min', '-m', dest='min', default=None,
                            help='min')
        parser.add_argument('--mid', '-d', dest='mid', default=None,
                            help='mid')
        parser.add_argument('--max', '-x', dest='max', default=None,
                            help='max')
        args = parser.parse_args()

        input_file = args.input_file
        output_file = args.output_file
        selected_bits = args.bits
        to_amp = args.amp
        custom_min = False
        custom_max = False
        if args.min is not None:
            WAV_PCM_FORMATS[selected_bits][MIN] = int(args.min)
            custom_min = True
        if args.mid is not None:
            WAV_PCM_FORMATS[selected_bits][MID] = int(args.mid)
        if args.max is not None:
            WAV_PCM_FORMATS[selected_bits][MAX] = int(args.max)
            custom_max = True

        w = wave.open(input_file, 'r')
        wt = wave.open(output_file, "w")
        print("Copy parameters ...")

        output_params = w.getparams()
        channels = output_params.nchannels

        wt.setparams(output_params)
        if not to_amp:
            print("Analyzing signal...")
            peaks = signal_analyze(w, selected_bits, channels)
            print("Analyze finished [%d channels detected]" % channels)
            print("Copy and convert the signal...")
            peak_index = [0] * channels
            val_cmp = [0] * channels
            last_position_from_middle = [None] * channels
            current_position_from_middle = [True] * channels
            for i in range(w.getnframes()):
                frame = w.readframes(1)
                new_frame = bytearray()

                for channel in range(0, channels):
                    index = channel * 2
                    frame_cmp = frame[index:index + WAV_PCM_FORMATS[selected_bits][BYTES]]
                    frame_tmp = bytearray()
                    val_cmp[channel] = get_frame_value(frame_cmp)

                    if val_cmp[channel] >= WAV_PCM_FORMATS[selected_bits][MID]:
                        current_position_from_middle[channel] = True
                    else:
                        current_position_from_middle[channel] = False

                    if last_position_from_middle[channel] is None:
                        last_position_from_middle[channel] = current_position_from_middle[channel]

                    if current_position_from_middle[channel] != last_position_from_middle[channel]:
                        last_position_from_middle[channel] = current_position_from_middle[channel]
                        peak_index[channel] += 1

                    if val_cmp[channel] > WAV_PCM_FORMATS[selected_bits][MID]:
                        if custom_max :
                            val =  WAV_PCM_FORMATS[selected_bits][MAX]
                        else:
                            val = peaks[channel][peak_index[channel]]

                        frame_tmp = val.to_bytes(WAV_PCM_FORMATS[selected_bits][BYTES], 'little', signed=True)

                    if val_cmp[channel] < WAV_PCM_FORMATS[selected_bits][MID]:
                        if custom_min:
                            val =  WAV_PCM_FORMATS[selected_bits][MIN]
                        else:
                            val = peaks[channel][peak_index[channel]]

                        frame_tmp = val.to_bytes(WAV_PCM_FORMATS[selected_bits][BYTES], 'little', signed=True)

                    if val_cmp[channel] == WAV_PCM_FORMATS[selected_bits][MID]:
                        val = WAV_PCM_FORMATS[selected_bits][MID]
                        frame_tmp = val.to_bytes(WAV_PCM_FORMATS[selected_bits][BYTES], 'little', signed=True)

                    for byte in frame_tmp:
                        new_frame.append(byte)

                wt.writeframes(new_frame)
        else:
            print("Copy and convert the signal...")
            for i in range(w.getnframes()):
                frame = w.readframes(1)
                new_frame = bytearray()

                for channel in range(0, channels):
                    index = channel * 2
                    frame_cmp = frame[index:index + WAV_PCM_FORMATS[selected_bits][BYTES]]
                    frame_tmp = bytearray()
                    val_cmp = get_frame_value(frame_cmp)
                    if val_cmp > WAV_PCM_FORMATS[selected_bits][MID]:
                        val = WAV_PCM_FORMATS[selected_bits][MAX]
                        frame_tmp = val.to_bytes(WAV_PCM_FORMATS[selected_bits][BYTES], 'little', signed=True)

                    if val_cmp < WAV_PCM_FORMATS[selected_bits][MID]:
                        val = WAV_PCM_FORMATS[selected_bits][MIN]
                        frame_tmp = val.to_bytes(WAV_PCM_FORMATS[selected_bits][BYTES], 'little', signed=True)

                    if val_cmp == WAV_PCM_FORMATS[selected_bits][MID]:
                        val = WAV_PCM_FORMATS[selected_bits][MID]
                        frame_tmp = val.to_bytes(WAV_PCM_FORMATS[selected_bits][BYTES], 'little', signed=True)

                    for byte in frame_tmp:
                        new_frame.append(byte)

                wt.writeframes(new_frame)

        print("Transformation finished")
        w.close()
        wt.close()

    except Exception as ex:
        print(ex)