RealImpact/preprocess_measurements.py at main · samuel-clarke/RealImpact · 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
import argparse
import glob
import os
import re

import numpy as np
import pandas as pd
from pydub.utils import db_to_float, ratio_to_db
import scipy.io.wavfile
from scipy.spatial.transform import Rotation
import trimesh

parser = argparse.ArgumentParser()
parser.add_argument('directory', type=str, help='foo help')
parser.add_argument('--no_transform', dest='transform', action='store_false', help='Do not use pretrained VGGish weights')
parser.set_defaults(transform=True)
parser.add_argument('--num_microphones', type=int, default=15, help='Random seed')
args = parser.parse_args()

data_dir = args.directory
MAX_INT16 = 32767
NUM_MICROPHONES = args.num_microphones
MIC_BAR_LENGTH = 1890 - 70

def normalize_gains(signal, gains):
    signal_fl = signal.astype(np.float32)
    max_gain = np.max(gains)
    for i in range(signal_fl.shape[0]):
        signal_fl[i, :] *= db_to_float(max_gain - gains[i])
    return signal_fl, max_gain

def normalize_gains_dB(signal, gains, dB):
    signal_fl = signal.astype(np.float32)
    max_gain = np.max(gains)
    for i in range(signal_fl.shape[0]):
        signal_fl[i, :] *= db_to_float(dB - gains[i])
    return signal_fl, max_gain

def window_hammer_signal(signal, threshold_ratio):
    # Find the max value in the hamnmer signal, then look
    # for where the signal is below a threshold on either side, i.e. when
    # there is no force being recorded. Those samples are then set to 0 as
    # they are assumed to be noise.
    max_val = np.max(signal, axis=1)
    max_val_idx = np.argmax(signal, axis=1)

    windowed_signal = np.zeros_like(signal)
    min_idx = np.zeros_like(max_val_idx)
    for i in range(signal.shape[0]):
        if (i % 100) == 0:
            print(i)
        curr_min_idx = max_val_idx[i] - np.argmax(signal[i, (max_val_idx[i])::-1] < (threshold_ratio * max_val[i]))
        max_idx = max_val_idx[i] + np.argmax(signal[i, max_val_idx[i]:] < (threshold_ratio * max_val[i]))
        windowed_signal[i, 0:(max_idx-curr_min_idx)] = signal[i, curr_min_idx:max_idx]
        min_idx[i] = curr_min_idx
    print(i)

    return windowed_signal, min_idx

def deconvolve_hammer(signal, hammer):
    windowed_hammers, window_idx = window_hammer_signal(hammer, 0.02)
    windowed_hammers = np.repeat(windowed_hammers, NUM_MICROPHONES, axis=0).astype(np.float32)
    window_idx = np.repeat(window_idx, NUM_MICROPHONES, axis=0)
    deconvolved_signal = np.zeros_like(signal).astype(np.float32)
    deconvolved_signal = deconvolved_signal[:, np.min(window_idx):]
    for i in range(signal.shape[0]):
        if (i % 100) == 0:
            print(i)
        end_ind = (signal.shape[1]-window_idx[i])
        deconvolved_signal[i, :end_ind] = signal[i, window_idx[i]:]
        deconvolved_signal[i, :end_ind] = np.real(np.fft.ifft(np.fft.fft(deconvolved_signal[i,:end_ind])
                                                       /np.fft.fft(windowed_hammers[i,:end_ind])))
    print(i)
    # Cut off end of deconvolved signal to remove ramp up at the end.
    deconvolved_signal = deconvolved_signal[:, :-int(fs*0.1)]
    return deconvolved_signal

def get_mic_world_space(angle, distance, ind):
    mic_z = -(MIC_BAR_LENGTH/2) + ind/14 * MIC_BAR_LENGTH
    mic_x = 230 + distance
    mic_y = -((45/2) + 20.95) * np.ones_like(angle)
    mic_points = np.vstack((mic_x, mic_y, mic_z)).transpose()
    rot = Rotation.from_euler('z', angle, degrees=True)
    pos_meters = rot.apply(mic_points) / 1000.0
    return pos_meters

vertex_dirs = glob.glob(os.path.join(data_dir, 'v_*'))
vertex_ids = []
hammers = []
hammer_gains = []
hammer_pads = []
audios = []
positions = []
hammer_gains = []
hammer_pads = []
mic_gains = []
mic_pads = []
for vd in vertex_dirs:
    print(vd)
    angle_offset = 0
    if 'degrees' in vd:
        result = re.search('.*_(.*)degrees', vd)
        print(result.group(1))
        angle_offset = int(result.group(1))
    df = pd.read_csv(os.path.join(vd, 'metadata.csv'))
    for i, row in df.iterrows():
        if row['valid'] == 1:
            angle = int(row['angle'])
            distance = int(row['distance'])
            rec_dir = os.path.join(vd, 'angle_%03i_distance_%04i'%(angle, distance))
            fs, data = scipy.io.wavfile.read(os.path.join(rec_dir, 'Force.wav'))
            hammers.append(data)
            vertex_ids.append(int(row['vertex_id']))
            hammer_gains.append(row['hammer_gain'])
            hammer_pads.append(row['hammer_pad'])
            for i in range(NUM_MICROPHONES):
                mic_gains.append(row['%i_gain'%(i+1)])
                mic_pads.append(row['%i_pad'%(i+1)])
                positions.append((angle+angle_offset, distance, i))
                fs, data = scipy.io.wavfile.read(os.path.join(rec_dir, 'Microphone_%i.wav'%(i+1)))
                audios.append(data)

audios = np.array(audios)
hammers = np.array(hammers)
positions = np.array(positions)
vertex_ids = np.array(vertex_ids)
mic_gains = np.array(mic_gains)
mic_pads = np.array(mic_pads, dtype=bool)
mic_adj_gains = mic_gains + 20 * np.logical_not(mic_pads)
hammer_gains = np.array(hammer_gains)
hammer_pads = np.array(hammer_pads)
hammer_adj_gains = hammer_gains + 20 * np.logical_not(hammer_pads)

# Normalize to 20dB adjusted, 0dB unadjusted gain
audios_norm, audios_gain = normalize_gains_dB(audios, mic_adj_gains, 20)
if np.max(np.abs(hammers)) > 100:
    hammers = hammers / MAX_INT16

hammers_norm, hammers_gain = normalize_gains_dB(hammers, hammer_adj_gains, 20)

hammers_newtons = 1.11 * (1/.02558) * hammers_norm

deconvolved = deconvolve_hammer(audios_norm, hammers_norm)
deconvolved_norm = deconvolved * 0.99 / np.max(np.abs(deconvolved))

listener_positions = get_mic_world_space(positions[:, 0], positions[:, 1], positions[:,2])

np.save(os.path.join(data_dir,'preprocessed/sounds.npy'), audios_norm / np.max(np.abs(audios_norm))/ 1.01)
np.save(os.path.join(data_dir, 'preprocessed/sounds_0db'), audios_norm)
np.save(os.path.join(data_dir,'preprocessed/hammers.npy'), hammers_norm)
np.save(os.path.join(data_dir, 'preprocessed/hammer_newtons.npy'), hammers_newtons)
np.save(os.path.join(data_dir,'preprocessed/deconvolved.npy'), deconvolved_norm)
np.save(os.path.join(data_dir,'preprocessed/deconvolved_0db.npy'), deconvolved)
np.save(os.path.join(data_dir,'preprocessed/vertexID.npy'), np.repeat(vertex_ids, 15, axis=0))
np.save(os.path.join(data_dir,'preprocessed/listenerXYZ.npy'), listener_positions)
if args.transform:
    mesh = trimesh.load(os.path.join(data_dir, 'preprocessed/transformed.obj'), process=False, maintain_order=True)
    vertex_positions = mesh.vertices[vertex_ids, :]
    vertex_positions = np.repeat(vertex_positions, 15, axis=0)
    np.save(os.path.join(data_dir,'preprocessed/vertexXYZ.npy'), vertex_positions)