Spaces:
Running
Running
| import numpy as np | |
| import pandas as pd | |
| import scipy.signal | |
| import tensorflow as tf | |
| pd.options.mode.chained_assignment = None | |
| import logging | |
| import os | |
| import threading | |
| import obspy | |
| from scipy.interpolate import interp1d | |
| tf.compat.v1.disable_eager_execution() | |
| # from tensorflow.python.ops.linalg_ops import norm | |
| # from tensorflow.python.util import nest | |
| class Config: | |
| seed = 100 | |
| n_class = 2 | |
| fs = 100 | |
| dt = 1.0 / fs | |
| freq_range = [0, fs / 2] | |
| time_range = [0, 30] | |
| nperseg = 30 | |
| nfft = 60 | |
| plot = False | |
| nt = 3000 | |
| X_shape = [31, 201, 2] | |
| Y_shape = [31, 201, n_class] | |
| signal_shape = [31, 201] | |
| noise_shape = signal_shape | |
| use_seed = False | |
| queue_size = 10 | |
| noise_mean = 2 | |
| noise_std = 1 | |
| # noise_low = 1 | |
| # noise_high = 5 | |
| use_buffer = True | |
| snr_threshold = 10 | |
| # %% | |
| # def normalize(data, window=3000): | |
| # """ | |
| # data: nsta, chn, nt | |
| # """ | |
| # shift = window//2 | |
| # nt = len(data) | |
| # ## std in slide windows | |
| # data_pad = np.pad(data, ((window//2, window//2)), mode="reflect") | |
| # t = np.arange(0, nt, shift, dtype="int") | |
| # # print(f"nt = {nt}, nt+window//2 = {nt+window//2}") | |
| # std = np.zeros(len(t)) | |
| # mean = np.zeros(len(t)) | |
| # for i in range(len(std)): | |
| # std[i] = np.std(data_pad[i*shift:i*shift+window]) | |
| # mean[i] = np.mean(data_pad[i*shift:i*shift+window]) | |
| # t = np.append(t, nt) | |
| # std = np.append(std, [np.std(data_pad[-window:])]) | |
| # mean = np.append(mean, [np.mean(data_pad[-window:])]) | |
| # # print(t) | |
| # ## normalize data with interplated std | |
| # t_interp = np.arange(nt, dtype="int") | |
| # std_interp = interp1d(t, std, kind="slinear")(t_interp) | |
| # mean_interp = interp1d(t, mean, kind="slinear")(t_interp) | |
| # data = (data - mean_interp)/(std_interp) | |
| # return data, std_interp | |
| # %% | |
| def normalize(data, window=200): | |
| """ | |
| data: nsta, chn, nt | |
| """ | |
| shift = window // 2 | |
| nt = data.shape[1] | |
| ## std in slide windows | |
| data_pad = np.pad(data, ((0, 0), (window // 2, window // 2), (0, 0)), mode="reflect") | |
| t = np.arange(0, nt, shift, dtype="int") | |
| # print(f"nt = {nt}, nt+window//2 = {nt+window//2}") | |
| std = np.zeros(len(t)) | |
| mean = np.zeros(len(t)) | |
| for i in range(len(std)): | |
| std[i] = np.std(data_pad[:, i * shift : i * shift + window, :]) | |
| mean[i] = np.mean(data_pad[:, i * shift : i * shift + window, :]) | |
| t = np.append(t, nt) | |
| std = np.append(std, [np.std(data_pad[:, -window:, :])]) | |
| mean = np.append(mean, [np.mean(data_pad[:, -window:, :])]) | |
| # print(t) | |
| ## normalize data with interplated std | |
| t_interp = np.arange(nt, dtype="int") | |
| std_interp = interp1d(t, std, kind="slinear")(t_interp) | |
| std_interp[std_interp == 0] = 1.0 | |
| mean_interp = interp1d(t, mean, kind="slinear")(t_interp) | |
| data = (data - mean_interp[np.newaxis, :, np.newaxis]) / std_interp[np.newaxis, :, np.newaxis] | |
| return data, std_interp | |
| def normalize_batch(data, window=200): | |
| """ | |
| data: nbn, nf, nt, 2 | |
| """ | |
| assert len(data.shape) == 4 | |
| shift = window // 2 | |
| nbt, nf, nt, nimg = data.shape | |
| ## std in slide windows | |
| data_pad = np.pad(data, ((0, 0), (0, 0), (window // 2, window // 2), (0, 0)), mode="reflect") | |
| t = np.arange(0, nt + shift - 1, shift, dtype="int") # 201 => 0, 100, 200 | |
| std = np.zeros([nbt, len(t)]) | |
| mean = np.zeros([nbt, len(t)]) | |
| for i in range(std.shape[1]): | |
| std[:, i] = np.std(data_pad[:, :, i * shift : i * shift + window, :], axis=(1, 2, 3)) | |
| mean[:, i] = np.mean(data_pad[:, :, i * shift : i * shift + window, :], axis=(1, 2, 3)) | |
| std[:, -1], mean[:, -1] = std[:, -2], mean[:, -2] | |
| std[:, 0], mean[:, 0] = std[:, 1], mean[:, 1] | |
| ## normalize data with interplated std | |
| t_interp = np.arange(nt, dtype="int") | |
| std_interp = interp1d(t, std, kind="slinear")(t_interp) ##nbt, nt | |
| std_interp[std_interp == 0] = 1.0 | |
| mean_interp = interp1d(t, mean, kind="slinear")(t_interp) | |
| data = (data - mean_interp[:, np.newaxis, :, np.newaxis]) / std_interp[:, np.newaxis, :, np.newaxis] | |
| if len(t) > 3: ##need to address this normalization issue in training | |
| data /= 2.0 | |
| return data | |
| # %% | |
| def py_func_decorator(output_types=None, output_shapes=None, name=None): | |
| def decorator(func): | |
| def call(*args, **kwargs): | |
| nonlocal output_shapes | |
| # flat_output_types = nest.flatten(output_types) | |
| flat_output_types = tf.nest.flatten(output_types) | |
| # flat_values = tf.py_func( | |
| flat_values = tf.numpy_function(func, inp=args, Tout=flat_output_types, name=name) | |
| if output_shapes is not None: | |
| for v, s in zip(flat_values, output_shapes): | |
| v.set_shape(s) | |
| # return nest.pack_sequence_as(output_types, flat_values) | |
| return tf.nest.pack_sequence_as(output_types, flat_values) | |
| return call | |
| return decorator | |
| def dataset_map(iterator, output_types, output_shapes=None, num_parallel_calls=None, name=None): | |
| dataset = tf.data.Dataset.range(len(iterator)) | |
| def index_to_entry(idx): | |
| return iterator[idx] | |
| return dataset.map(index_to_entry, num_parallel_calls=num_parallel_calls) | |
| class DataReader(object): | |
| def __init__( | |
| self, | |
| signal_dir=None, | |
| signal_list=None, | |
| noise_dir=None, | |
| noise_list=None, | |
| queue_size=None, | |
| coord=None, | |
| config=Config(), | |
| ): | |
| self.config = config | |
| signal_list = pd.read_csv(signal_list, header=0) | |
| noise_list = pd.read_csv(noise_list, header=0) | |
| self.signal = signal_list | |
| self.noise = noise_list | |
| self.n_signal = len(self.signal) | |
| self.signal_dir = signal_dir | |
| self.noise_dir = noise_dir | |
| self.X_shape = config.X_shape | |
| self.Y_shape = config.Y_shape | |
| self.n_class = config.n_class | |
| self.coord = coord | |
| self.threads = [] | |
| self.queue_size = queue_size | |
| self.add_queue() | |
| self.buffer_signal = {} | |
| self.buffer_noise = {} | |
| self.buffer_channels_signal = {} | |
| self.buffer_channels_noise = {} | |
| def add_queue(self): | |
| with tf.device('/cpu:0'): | |
| self.sample_placeholder = tf.compat.v1.placeholder(dtype=tf.float32, shape=None) | |
| self.target_placeholder = tf.compat.v1.placeholder(dtype=tf.float32, shape=None) | |
| self.queue = tf.queue.PaddingFIFOQueue( | |
| self.queue_size, ['float32', 'float32'], shapes=[self.config.X_shape, self.config.Y_shape] | |
| ) | |
| self.enqueue = self.queue.enqueue([self.sample_placeholder, self.target_placeholder]) | |
| return 0 | |
| def dequeue(self, num_elements): | |
| output = self.queue.dequeue_many(num_elements) | |
| return output | |
| def get_snr(self, data, itp, dit=300): | |
| tmp_std = np.std(data[itp - dit : itp]) | |
| if tmp_std > 0: | |
| return np.std(data[itp : itp + dit]) / tmp_std | |
| else: | |
| return 0 | |
| def add_event(self, sample, channels, j): | |
| while np.random.uniform(0, 1) < 0.2: | |
| shift = None | |
| if channels not in self.buffer_channels_signal: | |
| self.buffer_channels_signal[channels] = self.signal[self.signal['channels'] == channels] | |
| fname = os.path.join(self.signal_dir, self.buffer_channels_signal[channels].sample(n=1).iloc[0]['fname']) | |
| try: | |
| if fname not in self.buffer_signal: | |
| meta = np.load(fname) | |
| data_FT = [] | |
| snr = [] | |
| for i in range(3): | |
| tmp_data = meta['data'][:, i] | |
| tmp_itp = meta['itp'] | |
| snr.append(self.get_snr(tmp_data, tmp_itp)) | |
| tmp_data -= np.mean(tmp_data) | |
| f, t, tmp_FT = scipy.signal.stft( | |
| tmp_data, | |
| fs=self.config.fs, | |
| nperseg=self.config.nperseg, | |
| nfft=self.config.nfft, | |
| boundary='zeros', | |
| ) | |
| data_FT.append(tmp_FT) | |
| data_FT = np.stack(data_FT, axis=-1) | |
| self.buffer_signal[fname] = { | |
| 'data_FT': data_FT, | |
| 'itp': tmp_itp, | |
| 'channels': meta['channels'], | |
| 'snr': snr, | |
| } | |
| meta_signal = self.buffer_signal[fname] | |
| except: | |
| logging.error("Failed reading signal: {}".format(fname)) | |
| continue | |
| if meta_signal['snr'][j] > self.config.snr_threshold: | |
| tmp_signal = np.zeros([self.X_shape[0], self.X_shape[1]], dtype=np.complex_) | |
| shift = np.random.randint(-self.X_shape[1], 1, None, 'int') | |
| tmp_signal[:, -shift:] = meta_signal['data_FT'][:, self.X_shape[1] : 2 * self.X_shape[1] + shift, j] | |
| if np.isinf(tmp_signal).any() or np.isnan(tmp_signal).any() or (not np.any(tmp_signal)): | |
| continue | |
| tmp_signal = tmp_signal / np.std(tmp_signal) | |
| sample += tmp_signal / np.random.uniform(1, 5) | |
| return sample | |
| def thread_main(self, sess, n_threads=1, start=0): | |
| stop = False | |
| while not stop: | |
| index = list(range(start, self.n_signal, n_threads)) | |
| np.random.shuffle(index) | |
| for i in index: | |
| fname_signal = os.path.join(self.signal_dir, self.signal.iloc[i]['fname']) | |
| try: | |
| if fname_signal not in self.buffer_signal: | |
| meta = np.load(fname_signal) | |
| data_FT = [] | |
| snr = [] | |
| for j in range(3): | |
| tmp_data = meta['data'][..., j] | |
| tmp_itp = meta['itp'] | |
| snr.append(self.get_snr(tmp_data, tmp_itp)) | |
| tmp_data -= np.mean(tmp_data) | |
| f, t, tmp_FT = scipy.signal.stft( | |
| tmp_data, | |
| fs=self.config.fs, | |
| nperseg=self.config.nperseg, | |
| nfft=self.config.nfft, | |
| boundary='zeros', | |
| ) | |
| data_FT.append(tmp_FT) | |
| data_FT = np.stack(data_FT, axis=-1) | |
| self.buffer_signal[fname_signal] = { | |
| 'data_FT': data_FT, | |
| 'itp': tmp_itp, | |
| 'channels': meta['channels'], | |
| 'snr': snr, | |
| } | |
| meta_signal = self.buffer_signal[fname_signal] | |
| except: | |
| logging.error("Failed reading signal: {}".format(fname_signal)) | |
| continue | |
| channels = meta_signal['channels'].tolist() | |
| start_tp = meta_signal['itp'].tolist() | |
| if channels not in self.buffer_channels_noise: | |
| self.buffer_channels_noise[channels] = self.noise[self.noise['channels'] == channels] | |
| fname_noise = os.path.join( | |
| self.noise_dir, self.buffer_channels_noise[channels].sample(n=1).iloc[0]['fname'] | |
| ) | |
| try: | |
| if fname_noise not in self.buffer_noise: | |
| meta = np.load(fname_noise) | |
| data_FT = [] | |
| for i in range(3): | |
| tmp_data = meta['data'][: self.config.nt, i] | |
| tmp_data -= np.mean(tmp_data) | |
| f, t, tmp_FT = scipy.signal.stft( | |
| tmp_data, | |
| fs=self.config.fs, | |
| nperseg=self.config.nperseg, | |
| nfft=self.config.nfft, | |
| boundary='zeros', | |
| ) | |
| data_FT.append(tmp_FT) | |
| data_FT = np.stack(data_FT, axis=-1) | |
| self.buffer_noise[fname_noise] = {'data_FT': data_FT, 'channels': meta['channels']} | |
| meta_noise = self.buffer_noise[fname_noise] | |
| except: | |
| logging.error("Failed reading noise: {}".format(fname_noise)) | |
| continue | |
| if self.coord.should_stop(): | |
| stop = True | |
| break | |
| j = np.random.choice([0, 1, 2]) | |
| if meta_signal['snr'][j] <= self.config.snr_threshold: | |
| continue | |
| tmp_noise = meta_noise['data_FT'][..., j] | |
| if np.isinf(tmp_noise).any() or np.isnan(tmp_noise).any() or (not np.any(tmp_noise)): | |
| continue | |
| tmp_noise = tmp_noise / np.std(tmp_noise) | |
| tmp_signal = np.zeros([self.X_shape[0], self.X_shape[1]], dtype=np.complex_) | |
| if np.random.random() < 0.9: | |
| shift = np.random.randint(-self.X_shape[1], 1, None, 'int') | |
| tmp_signal[:, -shift:] = meta_signal['data_FT'][:, self.X_shape[1] : 2 * self.X_shape[1] + shift, j] | |
| if np.isinf(tmp_signal).any() or np.isnan(tmp_signal).any() or (not np.any(tmp_signal)): | |
| continue | |
| tmp_signal = tmp_signal / np.std(tmp_signal) | |
| tmp_signal = self.add_event(tmp_signal, channels, j) | |
| if np.random.random() < 0.2: | |
| tmp_signal = np.fliplr(tmp_signal) | |
| ratio = 0 | |
| while ratio <= 0: | |
| ratio = self.config.noise_mean + np.random.randn() * self.config.noise_std | |
| # ratio = np.random.uniform(self.config.noise_low, self.config.noise_high) | |
| tmp_noisy_signal = tmp_signal + ratio * tmp_noise | |
| noisy_signal = np.stack([tmp_noisy_signal.real, tmp_noisy_signal.imag], axis=-1) | |
| if np.isnan(noisy_signal).any() or np.isinf(noisy_signal).any(): | |
| continue | |
| noisy_signal = noisy_signal / np.std(noisy_signal) | |
| tmp_mask = np.abs(tmp_signal) / (np.abs(tmp_signal) + np.abs(ratio * tmp_noise) + 1e-4) | |
| tmp_mask[tmp_mask >= 1] = 1 | |
| tmp_mask[tmp_mask <= 0] = 0 | |
| mask = np.zeros([tmp_mask.shape[0], tmp_mask.shape[1], self.n_class]) | |
| mask[:, :, 0] = tmp_mask | |
| mask[:, :, 1] = 1 - tmp_mask | |
| sess.run(self.enqueue, feed_dict={self.sample_placeholder: noisy_signal, self.target_placeholder: mask}) | |
| def start_threads(self, sess, n_threads=8): | |
| for i in range(n_threads): | |
| thread = threading.Thread(target=self.thread_main, args=(sess, n_threads, i)) | |
| thread.daemon = True | |
| thread.start() | |
| self.threads.append(thread) | |
| return self.threads | |
| class DataReader_test(DataReader): | |
| def __init__( | |
| self, | |
| signal_dir=None, | |
| signal_list=None, | |
| noise_dir=None, | |
| noise_list=None, | |
| queue_size=None, | |
| coord=None, | |
| config=Config(), | |
| ): | |
| self.config = config | |
| signal_list = pd.read_csv(signal_list, header=0) | |
| noise_list = pd.read_csv(noise_list, header=0) | |
| self.signal = signal_list | |
| self.noise = noise_list | |
| self.n_signal = len(self.signal) | |
| self.signal_dir = signal_dir | |
| self.noise_dir = noise_dir | |
| self.X_shape = config.X_shape | |
| self.Y_shape = config.Y_shape | |
| self.n_class = config.n_class | |
| self.coord = coord | |
| self.threads = [] | |
| self.queue_size = queue_size | |
| self.add_queue() | |
| self.buffer_signal = {} | |
| self.buffer_noise = {} | |
| self.buffer_channels_signal = {} | |
| self.buffer_channels_noise = {} | |
| def add_queue(self): | |
| self.sample_placeholder = tf.compat.v1.placeholder(dtype=tf.float32, shape=None) | |
| self.target_placeholder = tf.compat.v1.placeholder(dtype=tf.float32, shape=None) | |
| self.ratio_placeholder = tf.compat.v1.placeholder(dtype=tf.float32, shape=None) | |
| self.signal_placeholder = tf.compat.v1.placeholder(dtype=tf.complex64, shape=None) | |
| self.noise_placeholder = tf.compat.v1.placeholder(dtype=tf.complex64, shape=None) | |
| self.fname_placeholder = tf.compat.v1.placeholder(dtype=tf.string, shape=None) | |
| self.queue = tf.queue.PaddingFIFOQueue( | |
| self.queue_size, | |
| ['float32', 'float32', 'float32', 'complex64', 'complex64', 'string'], | |
| shapes=[ | |
| self.config.X_shape, | |
| self.config.Y_shape, | |
| [], | |
| self.config.signal_shape, | |
| self.config.noise_shape, | |
| [], | |
| ], | |
| ) | |
| self.enqueue = self.queue.enqueue( | |
| [ | |
| self.sample_placeholder, | |
| self.target_placeholder, | |
| self.ratio_placeholder, | |
| self.signal_placeholder, | |
| self.noise_placeholder, | |
| self.fname_placeholder, | |
| ] | |
| ) | |
| return 0 | |
| def dequeue(self, num_elements): | |
| output = self.queue.dequeue_up_to(num_elements) | |
| return output | |
| def thread_main(self, sess, n_threads=1, start=0): | |
| index = list(range(start, self.n_signal, n_threads)) | |
| for i in index: | |
| np.random.seed(i) | |
| fname = self.signal.iloc[i]['fname'] | |
| fname_signal = os.path.join(self.signal_dir, fname) | |
| meta = np.load(fname_signal) | |
| data_FT = [] | |
| snr = [] | |
| for j in range(3): | |
| tmp_data = meta['data'][..., j] | |
| tmp_itp = meta['itp'] | |
| snr.append(self.get_snr(tmp_data, tmp_itp)) | |
| tmp_data -= np.mean(tmp_data) | |
| f, t, tmp_FT = scipy.signal.stft( | |
| tmp_data, fs=self.config.fs, nperseg=self.config.nperseg, nfft=self.config.nfft, boundary='zeros' | |
| ) | |
| data_FT.append(tmp_FT) | |
| data_FT = np.stack(data_FT, axis=-1) | |
| meta_signal = {'data_FT': data_FT, 'itp': tmp_itp, 'channels': meta['channels'], 'snr': snr} | |
| channels = meta['channels'].tolist() | |
| start_tp = meta['itp'].tolist() | |
| if channels not in self.buffer_channels_noise: | |
| self.buffer_channels_noise[channels] = self.noise[self.noise['channels'] == channels] | |
| fname_noise = os.path.join( | |
| self.noise_dir, self.buffer_channels_noise[channels].sample(n=1, random_state=i).iloc[0]['fname'] | |
| ) | |
| meta = np.load(fname_noise) | |
| data_FT = [] | |
| for i in range(3): | |
| tmp_data = meta['data'][: self.config.nt, i] | |
| tmp_data -= np.mean(tmp_data) | |
| f, t, tmp_FT = scipy.signal.stft( | |
| tmp_data, fs=self.config.fs, nperseg=self.config.nperseg, nfft=self.config.nfft, boundary='zeros' | |
| ) | |
| data_FT.append(tmp_FT) | |
| data_FT = np.stack(data_FT, axis=-1) | |
| meta_noise = {'data_FT': data_FT, 'channels': meta['channels']} | |
| if self.coord.should_stop(): | |
| stop = True | |
| break | |
| j = np.random.choice([0, 1, 2]) | |
| tmp_noise = meta_noise['data_FT'][..., j] | |
| if np.isinf(tmp_noise).any() or np.isnan(tmp_noise).any() or (not np.any(tmp_noise)): | |
| continue | |
| tmp_noise = tmp_noise / np.std(tmp_noise) | |
| tmp_signal = np.zeros([self.X_shape[0], self.X_shape[1]], dtype=np.complex_) | |
| if np.random.random() < 0.9: | |
| shift = np.random.randint(-self.X_shape[1], 1, None, 'int') | |
| tmp_signal[:, -shift:] = meta_signal['data_FT'][:, self.X_shape[1] : 2 * self.X_shape[1] + shift, j] | |
| if np.isinf(tmp_signal).any() or np.isnan(tmp_signal).any() or (not np.any(tmp_signal)): | |
| continue | |
| tmp_signal = tmp_signal / np.std(tmp_signal) | |
| # tmp_signal = self.add_event(tmp_signal, channels, j) | |
| # if np.random.random() < 0.2: | |
| # tmp_signal = np.fliplr(tmp_signal) | |
| ratio = 0 | |
| while ratio <= 0: | |
| ratio = self.config.noise_mean + np.random.randn() * self.config.noise_std | |
| tmp_noisy_signal = tmp_signal + ratio * tmp_noise | |
| noisy_signal = np.stack([tmp_noisy_signal.real, tmp_noisy_signal.imag], axis=-1) | |
| if np.isnan(noisy_signal).any() or np.isinf(noisy_signal).any(): | |
| continue | |
| std_noisy_signal = np.std(noisy_signal) | |
| noisy_signal = noisy_signal / std_noisy_signal | |
| tmp_mask = np.abs(tmp_signal) / (np.abs(tmp_signal) + np.abs(ratio * tmp_noise) + 1e-4) | |
| tmp_mask[tmp_mask >= 1] = 1 | |
| tmp_mask[tmp_mask <= 0] = 0 | |
| mask = np.zeros([tmp_mask.shape[0], tmp_mask.shape[1], self.n_class]) | |
| mask[:, :, 0] = tmp_mask | |
| mask[:, :, 1] = 1 - tmp_mask | |
| sess.run( | |
| self.enqueue, | |
| feed_dict={ | |
| self.sample_placeholder: noisy_signal, | |
| self.target_placeholder: mask, | |
| self.ratio_placeholder: std_noisy_signal, | |
| self.signal_placeholder: tmp_signal, | |
| self.noise_placeholder: ratio * tmp_noise, | |
| self.fname_placeholder: fname, | |
| }, | |
| ) | |
| class DataReader_pred_queue(DataReader): | |
| def __init__(self, signal_dir, signal_list, queue_size, coord, config=Config()): | |
| self.config = config | |
| signal_list = pd.read_csv(signal_list) | |
| self.signal = signal_list | |
| self.n_signal = len(self.signal) | |
| self.n_class = config.n_class | |
| self.X_shape = config.X_shape | |
| self.Y_shape = config.Y_shape | |
| self.signal_dir = signal_dir | |
| self.coord = coord | |
| self.threads = [] | |
| self.queue_size = queue_size | |
| self.add_placeholder() | |
| def add_placeholder(self): | |
| self.sample_placeholder = tf.compat.v1.placeholder(dtype=tf.float32, shape=None) | |
| self.ratio_placeholder = tf.compat.v1.placeholder(dtype=tf.float32, shape=None) | |
| self.fname_placeholder = tf.compat.v1.placeholder(dtype=tf.string, shape=None) | |
| self.queue = tf.queue.PaddingFIFOQueue( | |
| self.queue_size, ['float32', 'float32', 'string'], shapes=[self.config.X_shape, [], []] | |
| ) | |
| self.enqueue = self.queue.enqueue([self.sample_placeholder, self.ratio_placeholder, self.fname_placeholder]) | |
| def dequeue(self, num_elements): | |
| output = self.queue.dequeue_up_to(num_elements) | |
| return output | |
| def thread_main(self, sess, n_threads=1, start=0): | |
| index = list(range(start, self.n_signal, n_threads)) | |
| shift = 0 | |
| for i in index: | |
| fname = self.signal.iloc[i]['fname'] | |
| data_signal = np.load(os.path.join(self.signal_dir, fname)) | |
| f, t, tmp_signal = scipy.signal.stft( | |
| scipy.signal.detrend(np.squeeze(data_signal['data'][shift : self.config.nt + shift])), | |
| fs=self.config.fs, | |
| nperseg=self.config.nperseg, | |
| nfft=self.config.nfft, | |
| boundary='zeros', | |
| ) | |
| noisy_signal = np.stack([tmp_signal.real, tmp_signal.imag], axis=-1) | |
| if np.isnan(noisy_signal).any() or np.isinf(noisy_signal).any() or (not np.any(noisy_signal)): | |
| continue | |
| std_noisy_signal = np.std(noisy_signal) | |
| if std_noisy_signal == 0: | |
| continue | |
| noisy_signal = noisy_signal / std_noisy_signal | |
| sess.run( | |
| self.enqueue, | |
| feed_dict={ | |
| self.sample_placeholder: noisy_signal, | |
| self.ratio_placeholder: std_noisy_signal, | |
| self.fname_placeholder: fname, | |
| }, | |
| ) | |
| class DataReader_pred: | |
| def __init__(self, signal_dir, signal_list, format="numpy", sampling_rate=100, config=Config()): | |
| self.buffer = {} | |
| self.config = config | |
| self.format = format | |
| self.dtype = "float32" | |
| try: | |
| signal_list = pd.read_csv(signal_list, sep="\t")["fname"] | |
| except: | |
| signal_list = pd.read_csv(signal_list)["fname"] | |
| self.signal_list = signal_list | |
| self.n_signal = len(self.signal_list) | |
| self.signal_dir = signal_dir | |
| self.sampling_rate = sampling_rate | |
| self.n_class = config.n_class | |
| FT_shape = self.get_data_shape() | |
| self.X_shape = [*FT_shape, 2] | |
| def get_data_shape(self): | |
| # fname = self.signal_list.iloc[0]['fname'] | |
| # data = np.load(os.path.join(self.signal_dir, fname), allow_pickle=True)["data"] | |
| # data = np.squeeze(data) | |
| base_name = self.signal_list[0] | |
| if self.format == "numpy": | |
| meta = self.read_numpy(os.path.join(self.signal_dir, base_name)) | |
| elif self.format == "mseed": | |
| meta = self.read_mseed(os.path.join(self.signal_dir, base_name)) | |
| elif self.format == "hdf5": | |
| meta = self.read_hdf5(base_name) | |
| data = meta["data"] | |
| data = np.transpose(data, [2, 1, 0]) | |
| if self.sampling_rate != 100: | |
| t = np.linspace(0, 1, data.shape[-1]) | |
| t_interp = np.linspace(0, 1, np.int(np.around(data.shape[-1] * 100.0 / self.sampling_rate))) | |
| data = interp1d(t, data, kind="slinear")(t_interp) | |
| f, t, tmp_signal = scipy.signal.stft( | |
| data, fs=self.config.fs, nperseg=self.config.nperseg, nfft=self.config.nfft, boundary='zeros' | |
| ) | |
| logging.info(f"Input data shape: {tmp_signal.shape} measured on file {base_name}") | |
| return tmp_signal.shape | |
| def __len__(self): | |
| return self.n_signal | |
| def read_numpy(self, fname): | |
| # try: | |
| if fname not in self.buffer: | |
| npz = np.load(fname) | |
| meta = {} | |
| if len(npz['data'].shape) == 1: | |
| meta["data"] = npz['data'][:, np.newaxis, np.newaxis] | |
| elif len(npz['data'].shape) == 2: | |
| meta["data"] = npz['data'][:, np.newaxis, :] | |
| else: | |
| meta["data"] = npz['data'] | |
| if "p_idx" in npz.files: | |
| if len(npz["p_idx"].shape) == 0: | |
| meta["itp"] = [[npz["p_idx"]]] | |
| else: | |
| meta["itp"] = npz["p_idx"] | |
| if "s_idx" in npz.files: | |
| if len(npz["s_idx"].shape) == 0: | |
| meta["its"] = [[npz["s_idx"]]] | |
| else: | |
| meta["its"] = npz["s_idx"] | |
| if "t0" in npz.files: | |
| meta["t0"] = npz["t0"] | |
| self.buffer[fname] = meta | |
| else: | |
| meta = self.buffer[fname] | |
| return meta | |
| # except: | |
| # logging.error("Failed reading {}".format(fname)) | |
| # return None | |
| def read_hdf5(self, fname): | |
| data = self.h5_data[fname][()] | |
| attrs = self.h5_data[fname].attrs | |
| meta = {} | |
| if len(data.shape) == 2: | |
| meta["data"] = data[:, np.newaxis, :] | |
| else: | |
| meta["data"] = data | |
| if "p_idx" in attrs: | |
| if len(attrs["p_idx"].shape) == 0: | |
| meta["itp"] = [[attrs["p_idx"]]] | |
| else: | |
| meta["itp"] = attrs["p_idx"] | |
| if "s_idx" in attrs: | |
| if len(attrs["s_idx"].shape) == 0: | |
| meta["its"] = [[attrs["s_idx"]]] | |
| else: | |
| meta["its"] = attrs["s_idx"] | |
| if "t0" in attrs: | |
| meta["t0"] = attrs["t0"] | |
| return meta | |
| def read_s3(self, format, fname, bucket, key, secret, s3_url, use_ssl): | |
| with self.s3fs.open(bucket + "/" + fname, 'rb') as fp: | |
| if format == "numpy": | |
| meta = self.read_numpy(fp) | |
| elif format == "mseed": | |
| meta = self.read_mseed(fp) | |
| else: | |
| raise (f"Format {format} not supported") | |
| return meta | |
| def read_mseed(self, fname): | |
| mseed = obspy.read(fname) | |
| mseed = mseed.detrend("spline", order=2, dspline=5 * mseed[0].stats.sampling_rate) | |
| mseed = mseed.merge(fill_value=0) | |
| starttime = min([st.stats.starttime for st in mseed]) | |
| endtime = max([st.stats.endtime for st in mseed]) | |
| mseed = mseed.trim(starttime, endtime, pad=True, fill_value=0) | |
| if mseed[0].stats.sampling_rate != self.sampling_rate: | |
| logging.warning(f"Sampling rate {mseed[0].stats.sampling_rate} != {self.sampling_rate} Hz") | |
| order = ['3', '2', '1', 'E', 'N', 'Z'] | |
| order = {key: i for i, key in enumerate(order)} | |
| comp2idx = {"3": 0, "2": 1, "1": 2, "E": 0, "N": 1, "Z": 2} | |
| t0 = starttime.strftime("%Y-%m-%dT%H:%M:%S.%f")[:-3] | |
| nt = len(mseed[0].data) | |
| data = np.zeros([nt, 3], dtype=self.dtype) | |
| ids = [x.get_id() for x in mseed] | |
| if len(ids) == 3: | |
| for j, id in enumerate(sorted(ids, key=lambda x: order[x[-1]])): | |
| data[:, j] = mseed.select(id=id)[0].data.astype(self.dtype) | |
| else: | |
| if len(ids) > 3: | |
| logging.warning(f"More than 3 channels {ids}!") | |
| for jj, id in enumerate(ids): | |
| j = comp2idx[id[-1]] | |
| data[:, j] = mseed.select(id=id)[0].data.astype(self.dtype) | |
| data = data[:, np.newaxis, :] | |
| meta = {"data": data, "t0": t0} | |
| return meta | |
| def __getitem__(self, i): | |
| # fname = self.signal.iloc[i]['fname'] | |
| # data = np.load(os.path.join(self.signal_dir, fname), allow_pickle=True)["data"] | |
| # data = np.squeeze(data) | |
| base_name = self.signal_list[i] | |
| if self.format == "numpy": | |
| meta = self.read_numpy(os.path.join(self.signal_dir, base_name)) | |
| elif self.format == "mseed": | |
| meta = self.read_mseed(os.path.join(self.signal_dir, base_name)) | |
| elif self.format == "hdf5": | |
| meta = self.read_hdf5(base_name) | |
| data = meta["data"] # nt, 1, nch | |
| data = np.transpose(data, [2, 1, 0]) # nch, 1, nt | |
| if np.mod(data.shape[-1], 3000) == 1: # 3001=>3000 | |
| data = data[..., :-1] | |
| if "t0" in meta: | |
| t0 = meta["t0"] | |
| else: | |
| t0 = "1970-01-01T00:00:00.000" | |
| if self.sampling_rate != 100: | |
| logging.warning(f"Resample from {self.sampling_rate} to 100!") | |
| t = np.linspace(0, 1, data.shape[-1]) | |
| t_interp = np.linspace(0, 1, np.int(np.around(data.shape[-1] * 100.0 / self.sampling_rate))) | |
| data = interp1d(t, data, kind="slinear")(t_interp) | |
| # sos = scipy.signal.butter(4, 0.1, 'high', fs=100, output='sos') ## for stability of long sequence | |
| # data = scipy.signal.sosfilt(sos, data) | |
| f, t, tmp_signal = scipy.signal.stft( | |
| data, fs=self.config.fs, nperseg=self.config.nperseg, nfft=self.config.nfft, boundary='zeros' | |
| ) # nch, 1, nf, nt | |
| noisy_signal = np.stack([tmp_signal.real, tmp_signal.imag], axis=-1) # nch, 1, nf, nt, 2 | |
| noisy_signal[np.isnan(noisy_signal)] = 0 | |
| noisy_signal[np.isinf(noisy_signal)] = 0 | |
| # noisy_signal, std_noisy_signal = normalize(noisy_signal) | |
| # return noisy_signal.astype(self.dtype), std_noisy_signal.astype(self.dtype), fname | |
| return noisy_signal.astype(self.dtype), base_name, t0 | |
| def dataset(self, batch_size, num_parallel_calls=4): | |
| dataset = dataset_map( | |
| self, | |
| output_types=(self.dtype, "string", "string"), | |
| output_shapes=(self.X_shape, None, None), | |
| num_parallel_calls=num_parallel_calls, | |
| ) | |
| dataset = tf.compat.v1.data.make_one_shot_iterator( | |
| dataset.batch(batch_size).prefetch(batch_size * 3) | |
| ).get_next() | |
| return dataset | |
| if __name__ == "__main__": | |
| # %% | |
| data_reader = DataReader_pred(signal_dir="./Dataset/yixiao/", signal_list="./Dataset/yixiao.csv") | |
| noisy_signal, std_noisy_signal, fname = data_reader[0] | |
| print(noisy_signal.shape, std_noisy_signal.shape, fname) | |
| batch = data_reader.dataset(10) | |
| init = tf.compat.v1.initialize_all_variables() | |
| sess = tf.compat.v1.Session() | |
| sess.run(init) | |
| print(sess.run(batch)) | |