init

Dockerfile

@@ -0,0 +1,26 @@
+FROM tensorflow/tensorflow
+
+# Create the environment:
+# COPY env.yml /app
+# RUN conda env create --name cs329s --file=env.yml
+# Make RUN commands use the new environment:
+# SHELL ["conda", "run", "-n", "cs329s", "/bin/bash", "-c"]
+
+RUN pip install tqdm obspy pandas 
+RUN pip install uvicorn fastapi kafka-python
+
+WORKDIR /opt
+
+# Copy files
+COPY deepdenoiser /opt/deepdenoiser
+# COPY model /opt/model
+RUN wget https://github.com/AI4EPS/models/releases/download/DeepDenoiser/model.tar && tar -xvf model.tar && rm model.tar
+
+# Expose API port
+EXPOSE 8000
+
+ENV PYTHONUNBUFFERED=1
+
+# Start API server
+#ENTRYPOINT ["conda", "run", "--no-capture-output", "-n", "cs329s", "uvicorn", "--app-dir", "phasenet", "app:app", "--reload", "--port", "8000", "--host", "0.0.0.0"]
+ENTRYPOINT ["uvicorn", "--app-dir", "deepdenoiser", "app:app", "--reload", "--port", "7860", "--host", "0.0.0.0"]

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2021 Weiqiang Zhu
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

deepdenoiser/app.py

@@ -0,0 +1,180 @@
+import os
+from collections import defaultdict, namedtuple
+from datetime import datetime, timedelta
+from json import dumps
+from typing import Any, AnyStr, Dict, List, NamedTuple, Union
+
+import numpy as np
+import requests
+import tensorflow as tf
+from fastapi import FastAPI
+from kafka import KafkaProducer
+from pydantic import BaseModel
+import scipy
+from scipy.interpolate import interp1d
+
+from model import UNet
+
+tf.compat.v1.disable_eager_execution()
+tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
+PROJECT_ROOT = os.path.realpath(os.path.join(os.path.dirname(__file__), ".."))
+JSONObject = Dict[AnyStr, Any]
+JSONArray = List[Any]
+JSONStructure = Union[JSONArray, JSONObject]
+
+app = FastAPI()
+X_SHAPE = [3000, 1, 3]
+SAMPLING_RATE = 100
+
+# load model
+model = UNet(mode="pred")
+sess_config = tf.compat.v1.ConfigProto()
+sess_config.gpu_options.allow_growth = True
+
+sess = tf.compat.v1.Session(config=sess_config)
+saver = tf.compat.v1.train.Saver(tf.compat.v1.global_variables())
+init = tf.compat.v1.global_variables_initializer()
+sess.run(init)
+latest_check_point = tf.train.latest_checkpoint(f"{PROJECT_ROOT}/model/190614-104802")
+print(f"restoring model {latest_check_point}")
+saver.restore(sess, latest_check_point)
+
+# Kafak producer
+use_kafka = False
+# BROKER_URL = 'localhost:9092'
+# BROKER_URL = 'my-kafka-headless:9092'
+
+try:
+    print("Connecting to k8s kafka")
+    BROKER_URL = "quakeflow-kafka-headless:9092"
+    producer = KafkaProducer(
+        bootstrap_servers=[BROKER_URL],
+        key_serializer=lambda x: dumps(x).encode("utf-8"),
+        value_serializer=lambda x: dumps(x).encode("utf-8"),
+    )
+    use_kafka = True
+    print("k8s kafka connection success!")
+except BaseException:
+    print("k8s Kafka connection error")
+    try:
+        print("Connecting to local kafka")
+        producer = KafkaProducer(
+            bootstrap_servers=["localhost:9092"],
+            key_serializer=lambda x: dumps(x).encode("utf-8"),
+            value_serializer=lambda x: dumps(x).encode("utf-8"),
+        )
+        use_kafka = True
+        print("local kafka connection success!")
+    except BaseException:
+        print("local Kafka connection error")
+
+
+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
+    # print(f"nt = {nt}, nt+window//2 = {nt+window//2}")
+    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)
+    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 get_prediction(meta):
+
+    FS = 100
+    NPERSEG = 30
+    NFFT = 60
+
+    vec = np.array(meta.vec)  # [batch, nt, chn]
+    nbt, nt, nch = vec.shape
+    vec = np.transpose(vec, [0, 2, 1])  # [batch, chn, nt]
+    vec = np.reshape(vec, [nbt * nch, nt])  ## [batch * chn, nt]
+
+    if np.mod(vec.shape[-1], 3000) == 1:  # 3001=>3000
+        vec = vec[..., :-1]
+
+    if meta.dt != 0.01:
+        t = np.linspace(0, 1, len(vec))
+        t_interp = np.linspace(0, 1, np.int(np.around(len(vec) * meta.dt * FS)))
+        vec = interp1d(t, vec, kind="slinear")(t_interp)
+
+    # sos = scipy.signal.butter(4, 0.1, 'high', fs=100, output='sos')  ## for stability of long sequence
+    # vec = scipy.signal.sosfilt(sos, vec)
+    f, t, tmp_signal = scipy.signal.stft(vec, fs=FS, nperseg=NPERSEG, nfft=NFFT, boundary='zeros')
+    noisy_signal = np.stack([tmp_signal.real, tmp_signal.imag], axis=-1)  # [batch * chn, nf, nt, 2]
+    noisy_signal[np.isnan(noisy_signal)] = 0
+    noisy_signal[np.isinf(noisy_signal)] = 0
+    X_input = normalize_batch(noisy_signal)
+
+    feed = {model.X: X_input, model.drop_rate: 0, model.is_training: False}
+    preds = sess.run(model.preds, feed_dict=feed)
+
+    _, denoised_signal = scipy.signal.istft(
+        (noisy_signal[..., 0] + noisy_signal[..., 1] * 1j) * preds[..., 0],
+        fs=FS,
+        nperseg=NPERSEG,
+        nfft=NFFT,
+        boundary='zeros',
+    )
+    # _, denoised_noise = scipy.signal.istft(
+    #     (noisy_signal[..., 0] + noisy_signal[..., 1] * 1j) * preds[..., 1],
+    #     fs=FS,
+    #     nperseg=NPERSEG,
+    #     nfft=NFFT,
+    #     boundary='zeros',
+    # )
+
+    denoised_signal = np.reshape(denoised_signal, [nbt, nch, nt])
+    denoised_signal = np.transpose(denoised_signal, [0, 2, 1])
+
+    result = meta.copy()
+    result.vec = denoised_signal.tolist()
+    return result
+
+
+class Data(BaseModel):
+    # id: Union[List[str], str]
+    # timestamp: Union[List[str], str]
+    # vec: Union[List[List[List[float]]], List[List[float]]]
+    id: List[str]
+    timestamp: List[str]
+    vec: List[List[List[float]]]
+    dt: float = 0.01
+
+
+@app.post("/predict")
+def predict(data: Data):
+
+    denoised = get_prediction(data)
+
+    return denoised
+
+
+@app.get("/healthz")
+def healthz():
+    return {"status": "ok"}

deepdenoiser/data_reader.py

@@ -0,0 +1,816 @@
+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))
+
+    @py_func_decorator(output_types, output_shapes, name=name)
+    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))

deepdenoiser/model.py

@@ -0,0 +1,495 @@
+import logging
+
+import numpy as np
+import tensorflow as tf
+
+from util import *
+
+tf.compat.v1.disable_eager_execution()
+
+
+class ModelConfig:
+
+    batch_size = 20
+    depths = 6
+    filters_root = 8
+    kernel_size = [3, 3]
+    pool_size = [2, 2]
+    dilation_rate = [1, 1]
+    class_weights = [1.0, 1.0, 1.0]
+    loss_type = "cross_entropy"
+    weight_decay = 0.0
+    optimizer = "adam"
+    momentum = 0.9
+    learning_rate = 0.01
+    decay_step = 1e9
+    decay_rate = 0.9
+    drop_rate = 0.0
+    summary = True
+
+    X_shape = [31, 201, 2]
+    n_channel = X_shape[-1]
+    Y_shape = [31, 201, 2]
+    n_class = Y_shape[-1]
+
+    def __init__(self, **kwargs):
+        for k, v in kwargs.items():
+            setattr(self, k, v)
+
+    def update_args(self, args):
+        for k, v in vars(args).items():
+            setattr(self, k, v)
+
+
+def crop_and_concat(net1, net2):
+    """
+    the size(net1) <= size(net2)
+    """
+    # net1_shape = net1.get_shape().as_list()
+    # net2_shape = net2.get_shape().as_list()
+    # # print(net1_shape)
+    # # print(net2_shape)
+    # # if net2_shape[1] >= net1_shape[1] and net2_shape[2] >= net1_shape[2]:
+    # offsets = [0, (net2_shape[1] - net1_shape[1]) // 2, (net2_shape[2] - net1_shape[2]) // 2, 0]
+    # size = [-1, net1_shape[1], net1_shape[2], -1]
+    # net2_resize = tf.slice(net2, offsets, size)
+    # return tf.concat([net1, net2_resize], 3)
+    # # else:
+    # #     offsets = [0, (net1_shape[1] - net2_shape[1]) // 2, (net1_shape[2] - net2_shape[2]) // 2, 0]
+    # #     size = [-1, net2_shape[1], net2_shape[2], -1]
+    # #     net1_resize = tf.slice(net1, offsets, size)
+    # #     return tf.concat([net1_resize, net2], 3)
+
+    ## dynamic shape
+    chn1 = net1.get_shape().as_list()[-1]
+    chn2 = net2.get_shape().as_list()[-1]
+    net1_shape = tf.shape(net1)
+    net2_shape = tf.shape(net2)
+    # print(net1_shape)
+    # print(net2_shape)
+    # if net2_shape[1] >= net1_shape[1] and net2_shape[2] >= net1_shape[2]:
+    offsets = [0, (net2_shape[1] - net1_shape[1]) // 2, (net2_shape[2] - net1_shape[2]) // 2, 0]
+    size = [-1, net1_shape[1], net1_shape[2], -1]
+    net2_resize = tf.slice(net2, offsets, size)
+
+    out = tf.concat([net1, net2_resize], 3)
+    out.set_shape([None, None, None, chn1 + chn2])
+    return out
+
+
+def crop_only(net1, net2):
+    """
+    the size(net1) <= size(net2)
+    """
+    net1_shape = net1.get_shape().as_list()
+    net2_shape = net2.get_shape().as_list()
+    # print(net1_shape)
+    # print(net2_shape)
+    # if net2_shape[1] >= net1_shape[1] and net2_shape[2] >= net1_shape[2]:
+    offsets = [0, (net2_shape[1] - net1_shape[1]) // 2, (net2_shape[2] - net1_shape[2]) // 2, 0]
+    size = [-1, net1_shape[1], net1_shape[2], -1]
+    net2_resize = tf.slice(net2, offsets, size)
+    # return tf.concat([net1, net2_resize], 3)
+    return net2_resize
+
+
+class UNet:
+    def __init__(self, config=ModelConfig(), input_batch=None, mode='train'):
+        self.depths = config.depths
+        self.filters_root = config.filters_root
+        self.kernel_size = config.kernel_size
+        self.dilation_rate = config.dilation_rate
+        self.pool_size = config.pool_size
+        self.X_shape = config.X_shape
+        self.Y_shape = config.Y_shape
+        self.n_channel = config.n_channel
+        self.n_class = config.n_class
+        self.class_weights = config.class_weights
+        self.batch_size = config.batch_size
+        self.loss_type = config.loss_type
+        self.weight_decay = config.weight_decay
+        self.optimizer = config.optimizer
+        self.decay_step = config.decay_step
+        self.decay_rate = config.decay_rate
+        self.momentum = config.momentum
+        self.learning_rate = config.learning_rate
+        self.global_step = tf.compat.v1.get_variable(name="global_step", initializer=0, dtype=tf.int32)
+        self.summary_train = []
+        self.summary_valid = []
+
+        self.build(input_batch, mode=mode)
+
+    def add_placeholders(self, input_batch=None, mode='train'):
+        if input_batch is None:
+            self.X = tf.compat.v1.placeholder(
+                dtype=tf.float32, shape=[None, None, None, self.X_shape[-1]], name='X'
+            )
+            self.Y = tf.compat.v1.placeholder(
+                dtype=tf.float32, shape=[None, None, None, self.n_class], name='y'
+            )
+        else:
+            self.X = input_batch[0]
+            if mode in ["train", "valid", "test"]:
+                self.Y = input_batch[1]
+            self.input_batch = input_batch
+
+        self.is_training = tf.compat.v1.placeholder(dtype=tf.bool, name="is_training")
+        # self.keep_prob = tf.placeholder(dtype=tf.float32, name="keep_prob")
+        self.drop_rate = tf.compat.v1.placeholder(dtype=tf.float32, name="drop_rate")
+        # self.learning_rate = tf.placeholder_with_default(tf.constant(0.01, dtype=tf.float32), shape=[], name="learning_rate")
+        # self.global_step = tf.placeholder_with_default(tf.constant(0, dtype=tf.int32), shape=[], name="global_step")
+
+    def add_prediction_op(self):
+        logging.info(
+            "Model: depths {depths}, filters {filters}, "
+            "filter size {kernel_size[0]}x{kernel_size[1]}, "
+            "pool size: {pool_size[0]}x{pool_size[1]}, "
+            "dilation rate: {dilation_rate[0]}x{dilation_rate[1]}".format(
+                depths=self.depths,
+                filters=self.filters_root,
+                kernel_size=self.kernel_size,
+                dilation_rate=self.dilation_rate,
+                pool_size=self.pool_size,
+            )
+        )
+
+        if self.weight_decay > 0:
+            weight_decay = tf.constant(self.weight_decay, dtype=tf.float32, name="weight_constant")
+            self.regularizer = tf.keras.regularizers.l2(l=0.5 * (weight_decay))
+        else:
+            self.regularizer = None
+
+        self.initializer = tf.compat.v1.keras.initializers.VarianceScaling(
+            scale=1.0, mode="fan_avg", distribution="uniform"
+        )
+
+        # down sample layers
+        convs = [None] * self.depths  # store output of each depth
+
+        with tf.compat.v1.variable_scope("Input"):
+            net = self.X
+            net = tf.compat.v1.layers.conv2d(
+                net,
+                filters=self.filters_root,
+                kernel_size=self.kernel_size,
+                activation=None,
+                use_bias=False,
+                padding='same',
+                dilation_rate=self.dilation_rate,
+                kernel_initializer=self.initializer,
+                kernel_regularizer=self.regularizer,
+                # bias_regularizer=self.regularizer,
+                name="input_conv",
+            )
+            net = tf.compat.v1.layers.batch_normalization(net, training=self.is_training, name="input_bn")
+            net = tf.nn.relu(net, name="input_relu")
+            # net = tf.nn.dropout(net, self.keep_prob)
+            net = tf.compat.v1.layers.dropout(net, rate=self.drop_rate, training=self.is_training, name="input_dropout")
+
+        for depth in range(0, self.depths):
+            with tf.compat.v1.variable_scope("DownConv_%d" % depth):
+                filters = int(2 ** (depth) * self.filters_root)
+
+                net = tf.compat.v1.layers.conv2d(
+                    net,
+                    filters=filters,
+                    kernel_size=self.kernel_size,
+                    activation=None,
+                    use_bias=False,
+                    padding='same',
+                    dilation_rate=self.dilation_rate,
+                    kernel_initializer=self.initializer,
+                    kernel_regularizer=self.regularizer,
+                    # bias_regularizer=self.regularizer,
+                    name="down_conv1_{}".format(depth + 1),
+                )
+                net = tf.compat.v1.layers.batch_normalization(
+                    net, training=self.is_training, name="down_bn1_{}".format(depth + 1)
+                )
+                net = tf.nn.relu(net, name="down_relu1_{}".format(depth + 1))
+                net = tf.compat.v1.layers.dropout(
+                    net, rate=self.drop_rate, training=self.is_training, name="down_dropout1_{}".format(depth + 1)
+                )
+
+                convs[depth] = net
+
+                if depth < self.depths - 1:
+                    net = tf.compat.v1.layers.conv2d(
+                        net,
+                        filters=filters,
+                        kernel_size=self.kernel_size,
+                        strides=self.pool_size,
+                        activation=None,
+                        use_bias=False,
+                        padding='same',
+                        # dilation_rate=self.dilation_rate,
+                        kernel_initializer=self.initializer,
+                        kernel_regularizer=self.regularizer,
+                        # bias_regularizer=self.regularizer,
+                        name="down_conv3_{}".format(depth + 1),
+                    )
+                    net = tf.compat.v1.layers.batch_normalization(
+                        net, training=self.is_training, name="down_bn3_{}".format(depth + 1)
+                    )
+                    net = tf.nn.relu(net, name="down_relu3_{}".format(depth + 1))
+                    net = tf.compat.v1.layers.dropout(
+                        net, rate=self.drop_rate, training=self.is_training, name="down_dropout3_{}".format(depth + 1)
+                    )
+
+        # up layers
+        for depth in range(self.depths - 2, -1, -1):
+            with tf.compat.v1.variable_scope("UpConv_%d" % depth):
+                filters = int(2 ** (depth) * self.filters_root)
+                net = tf.compat.v1.layers.conv2d_transpose(
+                    net,
+                    filters=filters,
+                    kernel_size=self.kernel_size,
+                    strides=self.pool_size,
+                    activation=None,
+                    use_bias=False,
+                    padding="same",
+                    kernel_initializer=self.initializer,
+                    kernel_regularizer=self.regularizer,
+                    # bias_regularizer=self.regularizer,
+                    name="up_conv0_{}".format(depth + 1),
+                )
+                net = tf.compat.v1.layers.batch_normalization(
+                    net, training=self.is_training, name="up_bn0_{}".format(depth + 1)
+                )
+                net = tf.nn.relu(net, name="up_relu0_{}".format(depth + 1))
+                net = tf.compat.v1.layers.dropout(
+                    net, rate=self.drop_rate, training=self.is_training, name="up_dropout0_{}".format(depth + 1)
+                )
+
+                # skip connection
+                net = crop_and_concat(convs[depth], net)
+                # net = crop_only(convs[depth], net)
+
+                net = tf.compat.v1.layers.conv2d(
+                    net,
+                    filters=filters,
+                    kernel_size=self.kernel_size,
+                    activation=None,
+                    use_bias=False,
+                    padding='same',
+                    dilation_rate=self.dilation_rate,
+                    kernel_initializer=self.initializer,
+                    kernel_regularizer=self.regularizer,
+                    # bias_regularizer=self.regularizer,
+                    name="up_conv1_{}".format(depth + 1),
+                )
+                net = tf.compat.v1.layers.batch_normalization(
+                    net, training=self.is_training, name="up_bn1_{}".format(depth + 1)
+                )
+                net = tf.nn.relu(net, name="up_relu1_{}".format(depth + 1))
+                net = tf.compat.v1.layers.dropout(
+                    net, rate=self.drop_rate, training=self.is_training, name="up_dropout1_{}".format(depth + 1)
+                )
+
+        # Output Map
+        with tf.compat.v1.variable_scope("Output"):
+            net = tf.compat.v1.layers.conv2d(
+                net,
+                filters=self.n_class,
+                kernel_size=(1, 1),
+                activation=None,
+                use_bias=True,
+                padding='same',
+                # dilation_rate=self.dilation_rate,
+                kernel_initializer=self.initializer,
+                kernel_regularizer=self.regularizer,
+                # bias_regularizer=self.regularizer,
+                name="output_conv",
+            )
+            # net = tf.nn.relu(net,
+            #                     name="output_relu")
+            # net = tf.layers.dropout(net,
+            #                         rate=self.drop_rate,
+            #                         training=self.is_training,
+            #                         name="output_dropout")
+            # net = tf.layers.batch_normalization(net,
+            #                                    training=self.is_training,
+            #                                    name="output_bn")
+            output = net
+
+        with tf.compat.v1.variable_scope("representation"):
+            self.representation = convs[-1]
+
+        with tf.compat.v1.variable_scope("logits"):
+            self.logits = output
+            tmp = tf.compat.v1.summary.histogram("logits", self.logits)
+            self.summary_train.append(tmp)
+
+        with tf.compat.v1.variable_scope("preds"):
+            self.preds = tf.nn.softmax(output)
+            tmp = tf.compat.v1.summary.histogram("preds", self.preds)
+            self.summary_train.append(tmp)
+
+    def add_loss_op(self):
+        if self.loss_type == "cross_entropy":
+            with tf.compat.v1.variable_scope("cross_entropy"):
+                flat_logits = tf.reshape(self.logits, [-1, self.n_class], name="logits")
+                flat_labels = tf.reshape(self.Y, [-1, self.n_class], name="labels")
+                if (np.array(self.class_weights) != 1).any():
+                    class_weights = tf.constant(np.array(self.class_weights, dtype=np.float32), name="class_weights")
+                    weight_map = tf.multiply(flat_labels, class_weights)
+                    weight_map = tf.reduce_sum(input_tensor=weight_map, axis=1)
+                    loss_map = tf.nn.softmax_cross_entropy_with_logits(logits=flat_logits, labels=flat_labels)
+                    #                     loss_map = tf.nn.sigmoid_cross_entropy_with_logits(logits=flat_logits,
+                    #                                                                       labels=flat_labels)
+                    weighted_loss = tf.multiply(loss_map, weight_map)
+                    loss = tf.reduce_mean(input_tensor=weighted_loss)
+                else:
+                    loss = tf.reduce_mean(
+                        input_tensor=tf.nn.softmax_cross_entropy_with_logits(logits=flat_logits, labels=flat_labels)
+                    )
+        #                     loss = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(logits=flat_logits,
+        #                                                                                   labels=flat_labels))
+        elif self.loss_type == "IOU":
+            with tf.compat.v1.variable_scope("IOU"):
+                eps = 1e-7
+                loss = 0
+                for i in range(1, self.n_class):
+                    intersection = eps + tf.reduce_sum(
+                        input_tensor=self.preds[:, :, :, i] * self.Y[:, :, :, i], axis=[1, 2]
+                    )
+                    union = (
+                        eps
+                        + tf.reduce_sum(input_tensor=self.preds[:, :, :, i], axis=[1, 2])
+                        + tf.reduce_sum(input_tensor=self.Y[:, :, :, i], axis=[1, 2])
+                    )
+                    loss += 1 - tf.reduce_mean(input_tensor=intersection / union)
+        elif self.loss_type == "mean_squared":
+            with tf.compat.v1.variable_scope("mean_squared"):
+                flat_logits = tf.reshape(self.logits, [-1, self.n_class], name="logits")
+                flat_labels = tf.reshape(self.Y, [-1, self.n_class], name="labels")
+                with tf.compat.v1.variable_scope("mean_squared"):
+                    loss = tf.compat.v1.losses.mean_squared_error(labels=flat_labels, predictions=flat_logits)
+        else:
+            raise ValueError("Unknown loss function: " % self.loss_type)
+
+        tmp = tf.compat.v1.summary.scalar("train_loss", loss)
+        self.summary_train.append(tmp)
+        tmp = tf.compat.v1.summary.scalar("valid_loss", loss)
+        self.summary_valid.append(tmp)
+
+        if self.weight_decay > 0:
+            with tf.compat.v1.name_scope('weight_loss'):
+                tmp = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.REGULARIZATION_LOSSES)
+                weight_loss = tf.add_n(tmp, name="weight_loss")
+            self.loss = loss + weight_loss
+        else:
+            self.loss = loss
+
+    def add_training_op(self):
+        if self.optimizer == "momentum":
+            self.learning_rate_node = tf.compat.v1.train.exponential_decay(
+                learning_rate=self.learning_rate,
+                global_step=self.global_step,
+                decay_steps=self.decay_step,
+                decay_rate=self.decay_rate,
+                staircase=True,
+            )
+            optimizer = tf.compat.v1.train.MomentumOptimizer(
+                learning_rate=self.learning_rate_node, momentum=self.momentum
+            )
+        elif self.optimizer == "adam":
+            self.learning_rate_node = tf.compat.v1.train.exponential_decay(
+                learning_rate=self.learning_rate,
+                global_step=self.global_step,
+                decay_steps=self.decay_step,
+                decay_rate=self.decay_rate,
+                staircase=True,
+            )
+
+            optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=self.learning_rate_node)
+        update_ops = tf.compat.v1.get_collection(tf.compat.v1.GraphKeys.UPDATE_OPS)
+        with tf.control_dependencies(update_ops):
+            self.train_op = optimizer.minimize(self.loss, global_step=self.global_step)
+        tmp = tf.compat.v1.summary.scalar("learning_rate", self.learning_rate_node)
+        self.summary_train.append(tmp)
+
+    def reset_learning_rate(self, sess, learning_rate, global_step):
+        self.learning_rate = learning_rate
+        assign_op = self.global_step.assign(global_step)
+        sess.run(assign_op)
+        if self.optimizer == "momentum":
+            self.learning_rate_node = tf.compat.v1.train.exponential_decay(
+                learning_rate=learning_rate,
+                global_step=self.global_step,
+                decay_steps=self.decay_step,
+                decay_rate=self.decay_rate,
+                staircase=True,
+            )
+            optimizer = tf.compat.v1.train.MomentumOptimizer(
+                learning_rate=self.learning_rate_node, momentum=self.momentum
+            )
+        elif self.optimizer == "adam":
+            self.learning_rate_node = tf.compat.v1.train.exponential_decay(
+                learning_rate=self.learning_rate,
+                global_step=self.global_step,
+                decay_steps=self.decay_step,
+                decay_rate=self.decay_rate,
+                staircase=True,
+            )
+
+            optimizer = tf.compat.v1.train.AdamOptimizer(learning_rate=self.learning_rate_node)
+
+    def train_on_batch(self, sess, X_batch, Y_batch, summary_writer, drop_rate=0.0):
+        feed = {self.drop_rate: drop_rate, self.is_training: True, self.X: X_batch, self.Y: Y_batch}
+        _, step_summary, step, loss = sess.run(
+            [self.train_op, self.summary_train, self.global_step, self.loss], feed_dict=feed
+        )
+        summary_writer.add_summary(step_summary, step)
+        return loss
+
+    def valid_on_batch(self, sess, X_batch, Y_batch, summary_writer, drop_rate=0.0):
+        feed = {self.drop_rate: drop_rate, self.is_training: False, self.X: X_batch, self.Y: Y_batch}
+        step_summary, step, loss, preds = sess.run(
+            [self.summary_valid, self.global_step, self.loss, self.preds], feed_dict=feed
+        )
+        summary_writer.add_summary(step_summary, step)
+        return loss, preds
+
+    def test_on_batch(self, sess, summary_writer):
+        feed = {self.drop_rate: 0, self.is_training: False}
+        (
+            step_summary,
+            step,
+            loss,
+            preds,
+            X_batch,
+            Y_batch,
+            ratio_batch,
+            signal_batch,
+            noise_batch,
+            fname_batch,
+        ) = sess.run(
+            [
+                self.summary_valid,
+                self.global_step,
+                self.loss,
+                self.preds,
+                self.X,
+                self.Y,
+                self.input_batch[2],
+                self.input_batch[3],
+                self.input_batch[4],
+                self.input_batch[5],
+            ],
+            feed_dict=feed,
+        )
+        summary_writer.add_summary(step_summary, step)
+
+        return loss, preds, X_batch, Y_batch, ratio_batch, signal_batch, noise_batch, fname_batch
+
+    def build(self, input_batch=None, mode='train'):
+        self.add_placeholders(input_batch, mode)
+        self.add_prediction_op()
+        if mode in ["train", "valid", "test"]:
+            self.add_loss_op()
+            self.add_training_op()
+            # self.add_metrics_op()
+            self.summary_train = tf.compat.v1.summary.merge(self.summary_train)
+            self.summary_valid = tf.compat.v1.summary.merge(self.summary_valid)
+        return 0

deepdenoiser/predict.py

@@ -0,0 +1,136 @@
+import argparse
+import logging
+import multiprocessing
+import os
+import time
+from functools import partial
+
+import numpy as np
+import tensorflow as tf
+from tqdm import tqdm
+
+from data_reader import DataReader_pred, normalize_batch
+from model import UNet
+from util import *
+
+tf.compat.v1.disable_eager_execution()
+tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
+
+
+def read_args():
+    """Returns args"""
+
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument("--format", default="numpy", type=str, help="Input data format: numpy or mseed")
+    parser.add_argument("--batch_size", default=20, type=int, help="Batch size")
+    parser.add_argument("--output_dir", default="output", help="Output directory (default: output)")
+    parser.add_argument("--model_dir", default=None, help="Checkpoint directory (default: None)")
+    parser.add_argument("--sampling_rate", default=100, type=int, help="sampling rate of pred data")
+    parser.add_argument("--data_dir", default="./Dataset/pred/", help="Input file directory")
+    parser.add_argument("--data_list", default="./Dataset/pred.csv", help="Input csv file")
+    parser.add_argument("--plot_figure", action="store_true", help="If plot figure")
+    parser.add_argument("--save_signal", action="store_true", help="If save denoised signal")
+    parser.add_argument("--save_noise", action="store_true", help="If save denoised noise")
+
+    args = parser.parse_args()
+    return args
+
+
+def pred_fn(args, data_reader, figure_dir=None, result_dir=None, log_dir=None):
+    current_time = time.strftime("%y%m%d-%H%M%S")
+    if log_dir is None:
+        log_dir = os.path.join(args.log_dir, "pred", current_time)
+    logging.info("Pred log: %s" % log_dir)
+    # logging.info("Dataset size: {}".format(data_reader.num_data))
+    if not os.path.exists(log_dir):
+        os.makedirs(log_dir)
+    if args.plot_figure:
+        figure_dir = os.path.join(log_dir, 'figures')
+        os.makedirs(figure_dir, exist_ok=True)
+    if args.save_signal or args.save_noise:
+        result_dir = os.path.join(log_dir, 'results')
+        os.makedirs(result_dir, exist_ok=True)
+
+    with tf.compat.v1.name_scope('Input_Batch'):
+        data_batch = data_reader.dataset(args.batch_size)
+
+    # model = UNet(input_batch=data_batch, mode='pred')
+    model = UNet(mode='pred')
+    sess_config = tf.compat.v1.ConfigProto()
+    sess_config.gpu_options.allow_growth = True
+    # sess_config.log_device_placement = False
+
+    with tf.compat.v1.Session(config=sess_config) as sess:
+
+        saver = tf.compat.v1.train.Saver(tf.compat.v1.global_variables())
+        init = tf.compat.v1.global_variables_initializer()
+        sess.run(init)
+
+        latest_check_point = tf.train.latest_checkpoint(args.model_dir)
+        logging.info(f"restoring models: {latest_check_point}")
+        saver.restore(sess, latest_check_point)
+
+        if args.plot_figure:
+            num_pool = multiprocessing.cpu_count()
+        else:
+            num_pool = 2
+        multiprocessing.set_start_method('spawn')
+        pool = multiprocessing.Pool(num_pool)
+        for _ in tqdm(range(0, data_reader.n_signal, args.batch_size), desc="Pred"):
+            X_batch, fname_batch, t0_batch = sess.run(data_batch)
+            nbt, nch, nst, nf, nt, nimg = X_batch.shape
+            X_batch_ = np.reshape(X_batch, [nbt * nch * nst, nf, nt, nimg])
+            X_batch_ = normalize_batch(X_batch_)
+            preds_batch = sess.run(
+                model.preds,
+                feed_dict={model.X: X_batch_, model.drop_rate: 0, model.is_training: False},
+            )
+            preds_batch = np.reshape(preds_batch, [nbt, nch, nst, nf, nt, preds_batch.shape[-1]])
+            # preds_batch, X_batch, ratio_batch, fname_batch = sess.run(
+            #     [model.preds, data_batch[0], data_batch[1], data_batch[2]],
+            #     feed_dict={model.drop_rate: 0, model.is_training: False},
+            # )
+
+            if args.save_signal or args.save_noise:
+                save_results(
+                    preds_batch,
+                    X_batch,
+                    fname=[x.decode() for x in fname_batch],
+                    t0=[x.decode() for x in t0_batch],
+                    save_signal=args.save_signal,
+                    save_noise=args.save_noise,
+                    result_dir=result_dir,
+                )
+
+            if args.plot_figure:
+                pool.starmap(
+                    partial(
+                        plot_figures,
+                        figure_dir=figure_dir,
+                    ),
+                    zip(preds_batch, X_batch, [x.decode() for x in fname_batch]),
+                )
+
+        pool.close()
+
+    return 0
+
+
+def main(args):
+
+    logging.basicConfig(format='%(asctime)s %(message)s', level=logging.INFO)
+
+    with tf.compat.v1.name_scope('create_inputs'):
+        data_reader = DataReader_pred(
+            format=args.format, signal_dir=args.data_dir, signal_list=args.data_list, sampling_rate=args.sampling_rate
+        )
+    logging.info("Dataset Size: {}".format(data_reader.n_signal))
+    pred_fn(args, data_reader, log_dir=args.output_dir)
+
+    return 0
+
+
+if __name__ == '__main__':
+    args = read_args()
+    main(args)

deepdenoiser/train.py

@@ -0,0 +1,557 @@
+#import warnings
+#warnings.filterwarnings('ignore', category=FutureWarning)
+import numpy as np
+import tensorflow as tf
+tf.compat.v1.disable_eager_execution()
+tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR)
+import argparse
+import os
+import time
+import logging
+from model import UNet
+from data_reader import *
+from util import *
+from tqdm import tqdm
+import multiprocessing
+from functools import partial
+
+
+def read_args():
+  """Returns args"""
+
+  parser = argparse.ArgumentParser()
+
+  parser.add_argument("--mode",
+                      default="train",
+                      help="train/valid/test/debug (default: train)")
+
+  parser.add_argument("--epochs",
+                      default=10,
+                      type=int,
+                      help="Number of epochs (default: 10)")
+
+  parser.add_argument("--batch_size",
+                      default=20,
+                      type=int,
+                      help="Batch size (default: 20)")
+
+  parser.add_argument("--learning_rate",
+                      default=0.001,
+                      type=float,
+                      help="learning rate (default: 0.001)")
+
+  parser.add_argument("--decay_step",
+                      default=-1,
+                      type=int,
+                      help="decay step (default: -1)")
+
+  parser.add_argument("--decay_rate",
+                      default=0.9,
+                      type=float,
+                      help="decay rate (default: 0.9)")
+
+  parser.add_argument("--momentum",
+                      default=0.9,
+                      type=float,
+                      help="momentum (default: 0.9)")
+
+  parser.add_argument("--filters_root",
+                      default=8,
+                      type=int,
+                      help="filters root (default: 8)")
+
+  parser.add_argument("--depth",
+                      default=6,
+                      type=int,
+                      help="depth (default: 6)")
+
+  parser.add_argument("--kernel_size",
+                      nargs="+",
+                      type=int,
+                      default=[3, 3],
+                      help="kernel size (default: [3, 3]")
+
+  parser.add_argument("--pool_size",
+                      nargs="+",
+                      type=int,
+                      default=[2, 2],
+                      help="pool size (default: [2, 2]")
+
+  parser.add_argument("--drop_rate",
+                      default=0,
+                      type=float,
+                      help="drop out rate (default: 0)")
+
+  parser.add_argument("--dilation_rate",
+                      nargs="+",
+                      type=int,
+                      default=[1, 1],
+                      help="dilation_rate (default: [1, 1]")
+
+  parser.add_argument("--loss_type",
+                      default="cross_entropy",
+                      help="loss type: cross_entropy, IOU, mean_squared (default: cross_entropy)")
+
+  parser.add_argument("--weight_decay",
+                      default=0,
+                      type=float,
+                      help="weight decay (default: 0)")
+
+  parser.add_argument("--optimizer",
+                      default="adam",
+                      help="optimizer: adam, momentum (default: adam)")
+
+  parser.add_argument("--summary",
+                      default=True,
+                      type=bool,
+                      help="summary (default: True)")
+
+  parser.add_argument("--class_weights",
+                      nargs="+",
+                      default=[1, 1],
+                      type=float,
+                      help="class weights (default: [1, 1]")
+
+  parser.add_argument("--log_dir",
+                      default="log",
+                      help="Tensorboard log directory (default: log)")
+
+  parser.add_argument("--model_dir",
+                      default=None,
+                      help="Checkpoint directory")
+
+  parser.add_argument("--num_plots",
+                      default=10,
+                      type=int,
+                      help="plotting trainning result (default: 10)")
+
+  parser.add_argument("--input_length",
+                      default=None,
+                      type=int,
+                      help="input length")
+  parser.add_argument("--sampling_rate",
+                      default=100,
+                      type=int,
+                      help="sampling rate of pred data in Hz (default: 100)")
+
+  parser.add_argument("--train_signal_dir",
+                      default="./Dataset/train/",
+                      help="Input file directory (default: ./Dataset/train/)")
+  parser.add_argument("--train_signal_list",
+                      default="./Dataset/train.csv",
+                      help="Input csv file (default: ./Dataset/train.csv)")
+  parser.add_argument("--train_noise_dir",
+                      default="./Dataset/train/",
+                      help="Input file directory (default: ./Dataset/train/)")
+  parser.add_argument("--train_noise_list",
+                      default="./Dataset/train.csv",
+                      help="Input csv file (default: ./Dataset/train.csv)")
+
+  parser.add_argument("--valid_signal_dir",
+                      default="./Dataset/",
+                      help="Input file directory (default: ./Dataset/)")
+  parser.add_argument("--valid_signal_list",
+                      default=None,
+                      help="Input csv file")
+  parser.add_argument("--valid_noise_dir",
+                      default="./Dataset/",
+                      help="Input file directory (default: ./Dataset/)")
+  parser.add_argument("--valid_noise_list",
+                      default=None,
+                      help="Input csv file")
+
+  parser.add_argument("--data_dir",
+                      default="./Dataset/pred/",
+                      help="Input file directory (default: ./Dataset/pred/)")
+  parser.add_argument("--data_list",
+                      default="./Dataset/pred.csv",
+                      help="Input csv file (default: ./Dataset/pred.csv)")
+
+  parser.add_argument("--output_dir",
+                      default=None,
+                      help="Output directory")
+
+  parser.add_argument("--fpred",
+                      default="preds.npz",
+                      help="ouput file name of test data")
+  parser.add_argument("--plot_figure",
+                      action="store_true",
+                      help="If plot figure for test")
+  parser.add_argument("--save_result",
+                      action="store_true",
+                      help="If save result for test")
+
+  args = parser.parse_args()
+  return args
+
+
+def set_config(args, data_reader):
+  config = Config()
+
+  config.X_shape = data_reader.X_shape
+  config.n_channel = config.X_shape[-1]
+  config.Y_shape = data_reader.Y_shape
+  config.n_class = config.Y_shape[-1]
+
+  config.depths = args.depth
+  config.filters_root = args.filters_root
+  config.kernel_size = args.kernel_size
+  config.pool_size = args.pool_size
+  config.dilation_rate = args.dilation_rate
+  config.batch_size = args.batch_size
+  config.class_weights = args.class_weights
+  config.loss_type = args.loss_type
+  config.weight_decay = args.weight_decay
+  config.optimizer = args.optimizer
+
+  config.learning_rate = args.learning_rate
+  if (args.decay_step == -1) and (args.mode == 'train'):
+    config.decay_step = data_reader.n_signal // args.batch_size
+  else:
+    config.decay_step = args.decay_step
+  config.decay_rate = args.decay_rate
+  config.momentum = args.momentum
+
+  config.summary = args.summary
+  config.drop_rate = args.drop_rate
+  config.class_weights = args.class_weights
+
+  return config
+
+
+def train_fn(args, data_reader, data_reader_valid=None):
+  current_time = time.strftime("%y%m%d-%H%M%S")
+  log_dir = os.path.join(args.log_dir, current_time)
+  logging.info("Training log: {}".format(log_dir))
+  if not os.path.exists(log_dir):
+    os.makedirs(log_dir)
+  figure_dir = os.path.join(log_dir, 'figures')
+  if not os.path.exists(figure_dir):
+    os.makedirs(figure_dir)
+
+  config = set_config(args, data_reader)
+  with open(os.path.join(log_dir, 'config.log'), 'w') as fp:
+    fp.write('\n'.join("%s: %s" % item for item in vars(config).items()))
+
+  with tf.compat.v1.name_scope('Input_Batch'):
+    batch = data_reader.dequeue(args.batch_size)
+    if data_reader_valid is not None:
+      batch_valid = data_reader_valid.dequeue(args.batch_size)
+
+  model = UNet(config)
+  sess_config = tf.compat.v1.ConfigProto()
+  sess_config.gpu_options.allow_growth = True
+  sess_config.log_device_placement = False
+
+  with tf.compat.v1.Session(config=sess_config) as sess:
+
+    summary_writer = tf.compat.v1.summary.FileWriter(log_dir, sess.graph)
+    saver = tf.compat.v1.train.Saver(tf.compat.v1.global_variables(), max_to_keep=5)
+    init = tf.compat.v1.global_variables_initializer()
+    sess.run(init)
+
+    if args.model_dir is not None:
+      logging.info("restoring models...")
+      latest_check_point = tf.train.latest_checkpoint(args.model_dir)
+      saver.restore(sess, latest_check_point)
+      model.reset_learning_rate(sess, learning_rate=0.01, global_step=0)
+
+
+    threads = data_reader.start_threads(sess, n_threads=multiprocessing.cpu_count())
+    if data_reader_valid is not None:
+      threads_valid = data_reader_valid.start_threads(sess, n_threads=multiprocessing.cpu_count())
+    flog = open(os.path.join(log_dir, 'loss.log'), 'w')
+
+    total_step = 0
+    mean_loss = 0
+    pool = multiprocessing.Pool(2)
+    for epoch in range(args.epochs):
+      progressbar = tqdm(range(0, data_reader.n_signal, args.batch_size), desc="{}: ".format(log_dir.split("/")[-1]))
+      for step in progressbar:
+        X_batch, Y_batch = sess.run(batch)
+        loss_batch = model.train_on_batch(sess, X_batch, Y_batch, summary_writer, args.drop_rate)
+        if epoch < 1:
+          mean_loss = loss_batch
+        else:
+          total_step += 1
+          mean_loss += (loss_batch-mean_loss)/total_step
+        progressbar.set_description("{}: epoch={}, loss={:.6f}, mean loss={:.6f}".format(log_dir.split("/")[-1], epoch, loss_batch, mean_loss))
+        flog.write("Epoch: {}, step: {}, loss: {}, mean loss: {}\n".format(epoch, step//args.batch_size, loss_batch, mean_loss))
+      saver.save(sess, os.path.join(log_dir, "model_{}.ckpt".format(epoch)))
+
+      ## valid
+      if data_reader_valid is not None:
+        mean_loss_valid = 0
+        total_step_valid = 0
+        progressbar = tqdm(range(0, data_reader_valid.n_signal, args.batch_size), desc="Valid: ")
+        for step in progressbar:
+          X_batch, Y_batch = sess.run(batch_valid)
+          loss_batch, preds_batch = model.valid_on_batch(sess, X_batch, Y_batch, summary_writer, args.drop_rate)
+          total_step_valid += 1
+          mean_loss_valid += (loss_batch-mean_loss_valid)/total_step_valid
+          progressbar.set_description("Valid: loss={:.6f}, mean loss={:.6f}".format(loss_batch, mean_loss_valid))
+          flog.write("Valid: {}, step: {}, loss: {}, mean loss: {}\n".format(epoch, step//args.batch_size, loss_batch, mean_loss_valid))
+
+        # plot_result(epoch, args.num_plots, figure_dir,  preds_batch, X_batch, Y_batch)
+        pool.map(partial(plot_result_thread, 
+                         epoch = epoch,
+                         preds = preds_batch,
+                         X = X_batch,
+                         Y = Y_batch,
+                         figure_dir = figure_dir),
+                range(args.num_plots))
+
+    flog.close()
+    pool.close()
+    data_reader.coord.request_stop()
+    if data_reader_valid is not None:
+      data_reader_valid.coord.request_stop()
+    try:
+      data_reader.coord.join(threads, stop_grace_period_secs=10, ignore_live_threads=True)
+      if data_reader_valid is not None:
+        data_reader_valid.coord.join(threads_valid, stop_grace_period_secs=10, ignore_live_threads=True)
+    except:
+      pass
+    sess.run(data_reader.queue.close(cancel_pending_enqueues=True))
+    if data_reader_valid is not None:
+      sess.run(data_reader_valid.queue.close(cancel_pending_enqueues=True))
+  return 0
+
+
+def test_fn(args, data_reader, figure_dir=None, result_dir=None):
+  current_time = time.strftime("%y%m%d-%H%M%S")
+  log_dir = os.path.join(args.log_dir, args.mode, current_time)
+  logging.info("{} log: {}".format(args.mode, log_dir))
+  if not os.path.exists(log_dir):
+    os.makedirs(log_dir)
+  if (args.plot_figure == True) and (figure_dir is None):
+    figure_dir = os.path.join(log_dir, 'figures')
+    if not os.path.exists(figure_dir):
+      os.makedirs(figure_dir)
+  if (args.save_result == True) and (result_dir is None):
+    result_dir = os.path.join(log_dir, 'results')
+    if not os.path.exists(result_dir):
+      os.makedirs(result_dir)
+
+  config = set_config(args, data_reader)
+  with open(os.path.join(log_dir, 'config.log'), 'w') as fp:
+    fp.write('\n'.join("%s: %s" % item for item in vars(config).items()))
+
+  with tf.compat.v1.name_scope('Input_Batch'):
+    batch = data_reader.dequeue(args.batch_size)
+
+  model = UNet(config, input_batch=batch, mode='test')
+  sess_config = tf.compat.v1.ConfigProto()
+  sess_config.gpu_options.allow_growth = True
+  sess_config.log_device_placement = False
+
+  with tf.compat.v1.Session(config=sess_config) as sess:
+
+    summary_writer = tf.compat.v1.summary.FileWriter(log_dir, sess.graph)
+    saver = tf.compat.v1.train.Saver(tf.compat.v1.global_variables(), max_to_keep=5)
+    init = tf.compat.v1.global_variables_initializer()
+    sess.run(init)
+
+    logging.info("restoring models...")
+    latest_check_point = tf.train.latest_checkpoint(args.model_dir)
+    saver.restore(sess, latest_check_point)
+
+    threads = data_reader.start_threads(sess, n_threads=multiprocessing.cpu_count())
+
+    flog = open(os.path.join(log_dir, 'loss.log'), 'w')
+    total_step = 0
+    mean_loss = 0
+    progressbar = tqdm(range(0, data_reader.n_signal, args.batch_size), desc=args.mode)
+    if args.plot_figure:                                                           
+      num_pool = multiprocessing.cpu_count()*2                                     
+    elif args.save_result:                                                         
+      num_pool = multiprocessing.cpu_count()                                       
+    else:                                                                          
+      num_pool = 2                                                                 
+    pool = multiprocessing.Pool(num_pool) 
+    for step in progressbar:
+
+      if step + args.batch_size >= data_reader.n_signal:
+        for t in threads:
+          t.join()
+        sess.run(data_reader.queue.close())
+      
+      loss_batch, preds_batch, X_batch, Y_batch, ratio_batch, \
+      signal_batch, noise_batch, fname_batch = model.test_on_batch(sess, summary_writer)
+      total_step += 1
+      mean_loss += (loss_batch-mean_loss)/total_step
+      progressbar.set_description("{}: loss={:.6f}, mean loss={:6f}".format(args.mode, loss_batch, mean_loss))
+      flog.write("step: {}, loss: {}\n".format(step, loss_batch))
+      flog.flush()
+
+      pool.map(partial(postprocessing_test, 
+                      preds=preds_batch, 
+                      X=X_batch*ratio_batch[:,np.newaxis,np.newaxis,np.newaxis],
+                      fname=fname_batch,
+                      figure_dir=figure_dir,
+                      result_dir=result_dir,
+                      signal_FT=signal_batch, 
+                      noise_FT=noise_batch), 
+                range(len(X_batch)))
+
+    flog.close()
+    pool.close()
+
+  return 0
+
+def pred_fn(args, data_reader, figure_dir=None, result_dir=None, log_dir=None):
+  current_time = time.strftime("%y%m%d-%H%M%S")
+  if log_dir is None:
+    log_dir = os.path.join(args.log_dir, "pred", current_time)
+  logging.info("Pred log: %s" % log_dir)
+  # logging.info("Dataset size: {}".format(data_reader.num_data))
+  if not os.path.exists(log_dir):
+    os.makedirs(log_dir)
+  if args.plot_figure:
+    figure_dir = os.path.join(log_dir, 'figures')
+    os.makedirs(figure_dir, exist_ok=True)
+  if args.save_result:
+    result_dir = os.path.join(log_dir, 'results')
+    os.makedirs(result_dir, exist_ok=True)
+
+  config = set_config(args, data_reader)
+  with open(os.path.join(log_dir, 'config.log'), 'w') as fp:
+    fp.write('\n'.join("%s: %s" % item for item in vars(config).items()))
+
+  with tf.compat.v1.name_scope('Input_Batch'):
+   data_batch = data_reader.dataset(args.batch_size)
+
+  # model = UNet(config, input_batch=batch, mode='pred')
+  model = UNet(config, mode='pred')
+  sess_config = tf.compat.v1.ConfigProto()
+  sess_config.gpu_options.allow_growth = True
+  #sess_config.log_device_placement = False
+
+  with tf.compat.v1.Session(config=sess_config) as sess:
+
+    saver = tf.compat.v1.train.Saver(tf.compat.v1.global_variables())
+    init = tf.compat.v1.global_variables_initializer()
+    sess.run(init)
+
+    logging.info("restoring models...")
+    latest_check_point = tf.train.latest_checkpoint(args.model_dir)
+    saver.restore(sess, latest_check_point)
+
+#    threads = data_reader.start_threads(sess, n_threads=multiprocessing.cpu_count())
+
+    if args.plot_figure:                                                           
+      num_pool = multiprocessing.cpu_count()                                   
+    elif args.save_result:                                                         
+      num_pool = multiprocessing.cpu_count()                                       
+    else:                                                                          
+      num_pool = 2                                                                 
+    multiprocessing.set_start_method('spawn')
+    pool = multiprocessing.Pool(num_pool)
+    for step in tqdm(range(0, data_reader.n_signal, args.batch_size), desc="Pred"):
+      #if step + args.batch_size >= data_reader.n_signal:
+      #  for t in threads:
+      #    t.join()
+      #  sess.run(data_reader.queue.close())
+      # X_batch = []
+      # ratio_batch = []
+      # fname_batch = []
+      # for i in range(step, min(step+args.batch_size, data_reader.n_signal)):
+      #   X, ratio, fname = data_reader[i]
+      #   if np.std(X) == 0:
+      #     continue
+      #   X_batch.append(X)
+      #   ratio_batch.append(ratio)
+      #   fname_batch.append(fname)
+      # X_batch = np.stack(X_batch, axis=0)
+      # ratio_batch = np.array(ratio_batch)
+      X_batch, ratio_batch, fname_batch = sess.run(data_batch)
+      preds_batch = sess.run(model.preds, feed_dict={model.X: X_batch,
+                                                     model.drop_rate: 0,
+                                                     model.is_training: False})
+      #preds_batch, X_batch, ratio_batch, fname_batch = sess.run([model.preds,
+      #                                                    batch[0],
+      #                                                    batch[1],
+      #                                                    batch[2]],
+      #                                                    feed_dict={model.drop_rate: 0,
+      #                                                               model.is_training: False})
+
+      pool.map(partial(postprocessing_pred,
+                       preds = preds_batch, 
+                       X = X_batch*ratio_batch[:,np.newaxis,:,np.newaxis],
+                       fname = [x.decode() for x in fname_batch],
+                       figure_dir = figure_dir, 
+                       result_dir = result_dir), 
+               range(len(X_batch)))
+
+      # for i in range(len(X_batch)):
+      #   postprocessing_thread(i,
+      #             preds = preds_batch, 
+      #             X = X_batch*ratio_batch[:,np.newaxis,np.newaxis,np.newaxis],
+      #             fname = fname_batch,
+      #             figure_dir = figure_dir, 
+      #             result_dir = result_dir)
+    
+    pool.close()
+
+  return 0
+
+def main(args):
+
+  logging.basicConfig(format='%(asctime)s %(message)s', level=logging.INFO)
+
+  coord = tf.train.Coordinator()
+
+  if args.mode == "train":
+    with tf.compat.v1.name_scope('create_inputs'):
+      data_reader = DataReader(
+          signal_dir=args.train_signal_dir,
+          signal_list=args.train_signal_list,
+          noise_dir=args.train_noise_dir,
+          noise_list=args.train_noise_list,
+          queue_size=args.batch_size*2,
+          coord=coord)
+      if (args.valid_signal_list is not None) and (args.valid_noise_list is not None):
+        data_reader_valid = DataReader(
+            signal_dir=args.valid_signal_dir,
+            signal_list=args.valid_signal_list,
+            noise_dir=args.valid_noise_dir,
+            noise_list=args.valid_noise_list,
+            queue_size=args.batch_size*2,
+            coord=coord)
+        logging.info("Dataset size: training %d, validation %d" %  (data_reader.n_signal, data_reader_valid.n_signal))
+      else:
+        data_reader_valid = None
+      logging.info("Dataset size: training %d, validation 0" %  (data_reader.n_signal))
+    train_fn(args, data_reader, data_reader_valid)
+  
+  elif args.mode == "valid" or args.mode == "test":
+    with tf.compat.v1.name_scope('create_inputs'):
+      data_reader = DataReader_test(
+          signal_dir=args.valid_signal_dir,
+          signal_list=args.valid_signal_list,
+          noise_dir=args.valid_noise_dir,
+          noise_list=args.valid_noise_list,
+          queue_size=args.batch_size*2,
+          coord=coord)
+    logging.info("Dataset Size: {}".format(data_reader.n_signal))
+    test_fn(args, data_reader)
+
+  elif args.mode == "pred":
+    with tf.compat.v1.name_scope('create_inputs'):
+      data_reader = DataReader_pred(
+          signal_dir=args.data_dir,
+          signal_list=args.data_list,
+          sampling_rate=args.sampling_rate)
+    logging.info("Dataset Size: {}".format(data_reader.n_signal))
+    pred_fn(args, data_reader, log_dir=args.output_dir)
+
+  else:
+    print("mode should be: train, valid, test, debug or pred")
+
+  coord.request_stop()
+  coord.join()
+  return 0
+
+if __name__ == '__main__':
+  args = read_args()
+  main(args)

deepdenoiser/util.py

@@ -0,0 +1,875 @@
+import os
+
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+import scipy
+from mpl_toolkits.axes_grid1.inset_locator import inset_axes, mark_inset
+from scipy import signal
+from tqdm import tqdm
+
+from data_reader import Config
+
+matplotlib.use('agg')
+
+
+def plot_result(epoch, num, figure_dir, preds, X, Y, mode="valid"):
+    config = Config()
+    for i in range(min(num, len(X))):
+
+        t, noisy_signal = scipy.signal.istft(
+            X[i, :, :, 0] + X[i, :, :, 1] * 1j, fs=config.fs, nperseg=config.nperseg, nfft=config.nfft, boundary='zeros'
+        )
+        t, ideal_denoised_signal = scipy.signal.istft(
+            (X[i, :, :, 0] + X[i, :, :, 1] * 1j) * Y[i, :, :, 0],
+            fs=config.fs,
+            nperseg=config.nperseg,
+            nfft=config.nfft,
+            boundary='zeros',
+        )
+        t, denoised_signal = scipy.signal.istft(
+            (X[i, :, :, 0] + X[i, :, :, 1] * 1j) * preds[i, :, :, 0],
+            fs=config.fs,
+            nperseg=config.nperseg,
+            nfft=config.nfft,
+            boundary='zeros',
+        )
+
+        plt.figure(i)
+        fig_size = plt.gcf().get_size_inches()
+        plt.gcf().set_size_inches(fig_size * [1.5, 1.5])
+        plt.subplot(4, 2, 1)
+        plt.pcolormesh(np.abs(X[i, :, :, 0] + X[i, :, :, 1] * 1j), vmin=0, vmax=2)
+        plt.title("Noisy signal")
+        plt.gca().set_xticklabels([])
+        plt.subplot(4, 2, 2)
+        plt.plot(t, noisy_signal, label='Noisy signal', linewidth=0.1)
+        signal_ylim = plt.gca().get_ylim()
+        plt.gca().set_xticklabels([])
+        plt.legend(loc='lower left')
+        plt.margins(x=0)
+
+        plt.subplot(4, 2, 3)
+        plt.pcolormesh(Y[i, :, :, 0], vmin=0, vmax=1)
+        plt.gca().set_xticklabels([])
+        plt.title("Y")
+        plt.subplot(4, 2, 4)
+        plt.pcolormesh(preds[i, :, :, 0], vmin=0, vmax=1)
+        plt.title("$\hat{Y}$")
+        plt.gca().set_xticklabels([])
+
+        plt.subplot(4, 2, 5)
+        plt.pcolormesh(np.abs(X[i, :, :, 0] + X[i, :, :, 1] * 1j) * Y[i, :, :, 0], vmin=0, vmax=2)
+        plt.title("Ideal denoised signal")
+        plt.gca().set_xticklabels([])
+        plt.subplot(4, 2, 6)
+        plt.pcolormesh(np.abs(X[i, :, :, 0] + X[i, :, :, 1] * 1j) * preds[i, :, :, 0], vmin=0, vmax=2)
+        plt.title("Denoised signal")
+        plt.gca().set_xticklabels([])
+
+        plt.subplot(4, 2, 7)
+        plt.plot(t, ideal_denoised_signal, label='Ideal denoised signal', linewidth=0.1)
+        plt.ylim(signal_ylim)
+        plt.xlabel("Time (s)")
+        plt.legend(loc='lower left')
+        plt.margins(x=0)
+        plt.subplot(4, 2, 8)
+        plt.plot(t, denoised_signal, label='Denoised signal', linewidth=0.1)
+        plt.ylim(signal_ylim)
+        plt.xlabel("Time (s)")
+        plt.legend(loc='lower left')
+        plt.margins(x=0)
+
+        plt.tight_layout()
+        plt.gcf().align_labels()
+        plt.savefig(os.path.join(figure_dir, "epoch{:03d}_{:03d}_{:}.png".format(epoch, i, mode)), bbox_inches='tight')
+        # plt.savefig(os.path.join(figure_dir, "epoch%03d_%03d.pdf" % (epoch, i)), bbox_inches='tight')
+        plt.close(i)
+    return 0
+
+
+def plot_result_thread(i, epoch, preds, X, Y, figure_dir, mode="valid"):
+    config = Config()
+    t, noisy_signal = scipy.signal.istft(
+        X[i, :, :, 0] + X[i, :, :, 1] * 1j, fs=config.fs, nperseg=config.nperseg, nfft=config.nfft, boundary='zeros'
+    )
+    t, ideal_denoised_signal = scipy.signal.istft(
+        (X[i, :, :, 0] + X[i, :, :, 1] * 1j) * Y[i, :, :, 0],
+        fs=config.fs,
+        nperseg=config.nperseg,
+        nfft=config.nfft,
+        boundary='zeros',
+    )
+    t, denoised_signal = scipy.signal.istft(
+        (X[i, :, :, 0] + X[i, :, :, 1] * 1j) * preds[i, :, :, 0],
+        fs=config.fs,
+        nperseg=config.nperseg,
+        nfft=config.nfft,
+        boundary='zeros',
+    )
+
+    plt.figure(i)
+    fig_size = plt.gcf().get_size_inches()
+    plt.gcf().set_size_inches(fig_size * [1.5, 1.5])
+    plt.subplot(4, 2, 1)
+    plt.pcolormesh(np.abs(X[i, :, :, 0] + X[i, :, :, 1] * 1j), vmin=0, vmax=2)
+    plt.title("Noisy signal")
+    plt.gca().set_xticklabels([])
+    plt.subplot(4, 2, 2)
+    plt.plot(t, noisy_signal, 'k', label='Noisy signal', linewidth=0.5)
+    signal_ylim = plt.gca().get_ylim()
+    plt.gca().set_xticklabels([])
+    plt.legend(loc='lower left')
+    plt.margins(x=0)
+
+    plt.subplot(4, 2, 3)
+    plt.pcolormesh(Y[i, :, :, 0], vmin=0, vmax=1)
+    plt.gca().set_xticklabels([])
+    plt.title("Y")
+    plt.subplot(4, 2, 4)
+    plt.pcolormesh(preds[i, :, :, 0], vmin=0, vmax=1)
+    plt.title("$\hat{Y}$")
+    plt.gca().set_xticklabels([])
+
+    plt.subplot(4, 2, 5)
+    plt.pcolormesh(np.abs(X[i, :, :, 0] + X[i, :, :, 1] * 1j) * Y[i, :, :, 0], vmin=0, vmax=2)
+    plt.title("Ideal denoised signal")
+    plt.gca().set_xticklabels([])
+    plt.subplot(4, 2, 6)
+    plt.pcolormesh(np.abs(X[i, :, :, 0] + X[i, :, :, 1] * 1j) * preds[i, :, :, 0], vmin=0, vmax=2)
+    plt.title("Denoised signal")
+    plt.gca().set_xticklabels([])
+
+    plt.subplot(4, 2, 7)
+    plt.plot(t, ideal_denoised_signal, 'k', label='Ideal denoised signal', linewidth=0.5)
+    plt.ylim(signal_ylim)
+    plt.xlabel("Time (s)")
+    plt.legend(loc='lower left')
+    plt.margins(x=0)
+    plt.subplot(4, 2, 8)
+    plt.plot(t, denoised_signal, 'k', label='Denoised signal', linewidth=0.5)
+    plt.ylim(signal_ylim)
+    plt.xlabel("Time (s)")
+    plt.legend(loc='lower left')
+    plt.margins(x=0)
+
+    plt.tight_layout()
+    plt.gcf().align_labels()
+    plt.savefig(os.path.join(figure_dir, "epoch{:03d}_{:03d}_{:}.png".format(epoch, i, mode)), bbox_inches='tight')
+    plt.close(i)
+    return 0
+
+
+def postprocessing_test(
+    i, preds, X, fname, figure_dir=None, result_dir=None, signal_FT=None, noise_FT=None, data_dir=None
+):
+    if (figure_dir is not None) or (result_dir is not None):
+        config = Config()
+        t1, noisy_signal = scipy.signal.istft(
+            X[i, :, :, 0] + X[i, :, :, 1] * 1j, fs=config.fs, nperseg=config.nperseg, nfft=config.nfft, boundary='zeros'
+        )
+        t1, denoised_signal = scipy.signal.istft(
+            (X[i, :, :, 0] + X[i, :, :, 1] * 1j) * preds[i, :, :, 0],
+            fs=config.fs,
+            nperseg=config.nperseg,
+            nfft=config.nfft,
+            boundary='zeros',
+        )
+        t1, denoised_noise = scipy.signal.istft(
+            (X[i, :, :, 0] + X[i, :, :, 1] * 1j) * (1 - preds[i, :, :, 0]),
+            fs=config.fs,
+            nperseg=config.nperseg,
+            nfft=config.nfft,
+            boundary='zeros',
+        )
+        t1, signal = scipy.signal.istft(
+            signal_FT[i, :, :], fs=config.fs, nperseg=config.nperseg, nfft=config.nfft, boundary='zeros'
+        )
+        t1, noise = scipy.signal.istft(
+            noise_FT[i, :, :], fs=config.fs, nperseg=config.nperseg, nfft=config.nfft, boundary='zeros'
+        )
+
+    if result_dir is not None:
+        try:
+            np.savez(
+                os.path.join(result_dir, fname[i].decode()),
+                preds=preds[i],
+                X=X[i],
+                signal_FT=signal_FT[i],
+                noise_FT=noise_FT[i],
+                noisy_signal=noisy_signal,
+                denoised_signal=denoised_signal,
+                denoised_noise=denoised_noise,
+                signal=signal,
+                noise=noise,
+            )
+        except FileNotFoundError:
+            os.makedirs(os.path.dirname(os.path.join(result_dir, fname[i].decode())), exist_ok=True)
+            np.savez(
+                os.path.join(result_dir, fname[i].decode()),
+                preds=preds[i],
+                X=X[i],
+                signal_FT=signal_FT[i],
+                noise_FT=noise_FT[i],
+                noisy_signal=noisy_signal,
+                denoised_signal=denoised_signal,
+                denoised_noise=denoised_noise,
+                signal=signal,
+                noise=noise,
+            )
+
+    if figure_dir is not None:
+        t_FT = np.linspace(config.time_range[0], config.time_range[1], X.shape[2])
+        f_FT = np.linspace(config.freq_range[0], config.freq_range[1], X.shape[1])
+
+        raw_data = None
+        if data_dir is not None:
+            raw_data = np.load(os.path.join(data_dir, fname[i].decode().split('/')[-1]))
+            itp = raw_data['itp']
+            its = raw_data['its']
+            ix1 = (750 - 50) / 100
+            ix2 = (750 + (its - itp) + 50) / 100
+            if ix2 - ix1 > 3:
+                ix2 = ix1 + 3
+
+        box = dict(boxstyle='round', facecolor='white', alpha=1)
+
+        text_loc = [0.05, 0.8]
+        plt.figure(i)
+        fig_size = plt.gcf().get_size_inches()
+        plt.gcf().set_size_inches(fig_size * [1, 2])
+        plt.subplot(511)
+        plt.pcolormesh(t_FT, f_FT, np.abs(signal_FT[i, :, :]), vmin=0, vmax=1)
+        plt.gca().set_xticklabels([])
+        plt.text(
+            text_loc[0],
+            text_loc[1],
+            '(i)',
+            horizontalalignment='center',
+            transform=plt.gca().transAxes,
+            fontsize="medium",
+            fontweight="bold",
+            bbox=box,
+        )
+        plt.subplot(512)
+        plt.pcolormesh(t_FT, f_FT, np.abs(noise_FT[i, :, :]), vmin=0, vmax=1)
+        plt.gca().set_xticklabels([])
+        plt.text(
+            text_loc[0],
+            text_loc[1],
+            '(ii)',
+            horizontalalignment='center',
+            transform=plt.gca().transAxes,
+            fontsize="medium",
+            fontweight="bold",
+            bbox=box,
+        )
+        plt.subplot(513)
+        plt.pcolormesh(t_FT, f_FT, np.abs(X[i, :, :, 0] + X[i, :, :, 1] * 1j), vmin=0, vmax=1)
+        plt.ylabel("Frequency (Hz)", fontsize='large')
+        plt.gca().set_xticklabels([])
+        plt.text(
+            text_loc[0],
+            text_loc[1],
+            '(iii)',
+            horizontalalignment='center',
+            transform=plt.gca().transAxes,
+            fontsize="medium",
+            fontweight="bold",
+            bbox=box,
+        )
+        plt.subplot(514)
+        plt.pcolormesh(t_FT, f_FT, np.abs(X[i, :, :, 0] + X[i, :, :, 1] * 1j) * preds[i, :, :, 0], vmin=0, vmax=1)
+        plt.gca().set_xticklabels([])
+        plt.text(
+            text_loc[0],
+            text_loc[1],
+            '(iv)',
+            horizontalalignment='center',
+            transform=plt.gca().transAxes,
+            fontsize="medium",
+            fontweight="bold",
+            bbox=box,
+        )
+        plt.subplot(515)
+        plt.pcolormesh(t_FT, f_FT, np.abs(X[i, :, :, 0] + X[i, :, :, 1] * 1j) * preds[i, :, :, 1], vmin=0, vmax=1)
+        plt.xlabel("Time (s)", fontsize='large')
+        plt.text(
+            text_loc[0],
+            text_loc[1],
+            '(v)',
+            horizontalalignment='center',
+            transform=plt.gca().transAxes,
+            fontsize="medium",
+            fontweight="bold",
+            bbox=box,
+        )
+
+        try:
+            plt.savefig(os.path.join(figure_dir, fname[i].decode().rstrip('.npz') + '_FT.png'), bbox_inches='tight')
+            # plt.savefig(os.path.join(figure_dir, fname[i].decode().rstrip('.npz')+'_FT.pdf'), bbox_inches='tight')
+        except FileNotFoundError:
+            os.makedirs(
+                os.path.dirname(os.path.join(figure_dir, fname[i].decode().rstrip('.npz') + '_FT.png')), exist_ok=True
+            )
+            plt.savefig(os.path.join(figure_dir, fname[i].decode().rstrip('.npz') + '_FT.png'), bbox_inches='tight')
+            # plt.savefig(os.path.join(figure_dir, fname[i].decode().rstrip('.npz')+'_FT.pdf'), bbox_inches='tight')
+        plt.close(i)
+
+        text_loc = [0.05, 0.8]
+        plt.figure(i)
+        fig_size = plt.gcf().get_size_inches()
+        plt.gcf().set_size_inches(fig_size * [1, 2])
+
+        ax3 = plt.subplot(513)
+        plt.plot(t1, noisy_signal, 'k', linewidth=0.5, label='Noisy signal')
+        plt.legend(loc='lower left', fontsize='medium')
+        plt.xlim([np.around(t1[0]), np.around(t1[-1])])
+        plt.ylim([-np.max(np.abs(noisy_signal)), np.max(np.abs(noisy_signal))])
+        signal_ylim = [-np.max(np.abs(noisy_signal[100:-100])), np.max(np.abs(noisy_signal[100:-100]))]
+        plt.ylim(signal_ylim)
+        plt.ylabel("Amplitude", fontsize='large')
+        plt.gca().set_xticklabels([])
+        plt.text(
+            text_loc[0],
+            text_loc[1],
+            '(iii)',
+            horizontalalignment='center',
+            transform=plt.gca().transAxes,
+            fontsize="medium",
+            fontweight="bold",
+            bbox=box,
+        )
+
+        ax1 = plt.subplot(511)
+        plt.plot(t1, signal, 'k', linewidth=0.5, label='Signal')
+        plt.legend(loc='lower left', fontsize='medium')
+        plt.xlim([np.around(t1[0]), np.around(t1[-1])])
+        plt.ylim(signal_ylim)
+        plt.gca().set_xticklabels([])
+        plt.text(
+            text_loc[0],
+            text_loc[1],
+            '(i)',
+            horizontalalignment='center',
+            transform=plt.gca().transAxes,
+            fontsize="medium",
+            fontweight="bold",
+            bbox=box,
+        )
+
+        plt.subplot(512)
+        plt.plot(t1, noise, 'k', linewidth=0.5, label='Noise')
+        plt.legend(loc='lower left', fontsize='medium')
+        plt.xlim([np.around(t1[0]), np.around(t1[-1])])
+        plt.ylim([-np.max(np.abs(noise)), np.max(np.abs(noise))])
+        noise_ylim = [-np.max(np.abs(noise[100:-100])), np.max(np.abs(noise[100:-100]))]
+        plt.ylim(noise_ylim)
+        plt.gca().set_xticklabels([])
+        plt.text(
+            text_loc[0],
+            text_loc[1],
+            '(ii)',
+            horizontalalignment='center',
+            transform=plt.gca().transAxes,
+            fontsize="medium",
+            fontweight="bold",
+            bbox=box,
+        )
+
+        ax4 = plt.subplot(514)
+        plt.plot(t1, denoised_signal, 'k', linewidth=0.5, label='Recovered signal')
+        plt.legend(loc='lower left', fontsize='medium')
+        plt.xlim([np.around(t1[0]), np.around(t1[-1])])
+        plt.ylim(signal_ylim)
+        plt.gca().set_xticklabels([])
+        plt.text(
+            text_loc[0],
+            text_loc[1],
+            '(iv)',
+            horizontalalignment='center',
+            transform=plt.gca().transAxes,
+            fontsize="medium",
+            fontweight="bold",
+            bbox=box,
+        )
+
+        plt.subplot(515)
+        plt.plot(t1, denoised_noise, 'k', linewidth=0.5, label='Recovered noise')
+        plt.legend(loc='lower left', fontsize='medium')
+        plt.xlim([np.around(t1[0]), np.around(t1[-1])])
+        plt.xlabel("Time (s)", fontsize='large')
+        plt.ylim(noise_ylim)
+        plt.text(
+            text_loc[0],
+            text_loc[1],
+            '(v)',
+            horizontalalignment='center',
+            transform=plt.gca().transAxes,
+            fontsize="medium",
+            fontweight="bold",
+            bbox=box,
+        )
+
+        if data_dir is not None:
+            axins = inset_axes(
+                ax1, width=2.0, height=1.0, loc='upper right', bbox_to_anchor=(1, 0.5), bbox_transform=ax1.transAxes
+            )
+            axins.plot(t1, signal, 'k', linewidth=0.5)
+            x1, x2 = ix1, ix2
+            y1 = -np.max(np.abs(signal[(t1 > ix1) & (t1 < ix2)]))
+            y2 = -y1
+            axins.set_xlim(x1, x2)
+            axins.set_ylim(y1, y2)
+            plt.xticks(visible=False)
+            plt.yticks(visible=False)
+            mark_inset(ax1, axins, loc1=1, loc2=3, fc="none", ec="0.5")
+
+            axins = inset_axes(
+                ax3, width=2.0, height=1.0, loc='upper right', bbox_to_anchor=(1, 0.3), bbox_transform=ax3.transAxes
+            )
+            axins.plot(t1, noisy_signal, 'k', linewidth=0.5)
+            x1, x2 = ix1, ix2
+            axins.set_xlim(x1, x2)
+            axins.set_ylim(y1, y2)
+            plt.xticks(visible=False)
+            plt.yticks(visible=False)
+            mark_inset(ax3, axins, loc1=1, loc2=3, fc="none", ec="0.5")
+
+            axins = inset_axes(
+                ax4, width=2.0, height=1.0, loc='upper right', bbox_to_anchor=(1, 0.5), bbox_transform=ax4.transAxes
+            )
+            axins.plot(t1, denoised_signal, 'k', linewidth=0.5)
+            x1, x2 = ix1, ix2
+            axins.set_xlim(x1, x2)
+            axins.set_ylim(y1, y2)
+            plt.xticks(visible=False)
+            plt.yticks(visible=False)
+            mark_inset(ax4, axins, loc1=1, loc2=3, fc="none", ec="0.5")
+
+        plt.savefig(os.path.join(figure_dir, fname[i].decode().rstrip('.npz') + '_wave.png'), bbox_inches='tight')
+        # plt.savefig(os.path.join(figure_dir, fname[i].decode().rstrip('.npz')+'_wave.pdf'), bbox_inches='tight')
+        plt.close(i)
+
+    return
+
+
+def postprocessing_pred(i, preds, X, fname, figure_dir=None, result_dir=None):
+
+    if (result_dir is not None) or (figure_dir is not None):
+        config = Config()
+
+        t1, noisy_signal = scipy.signal.istft(
+            (X[i, :, :, 0] + X[i, :, :, 1] * 1j),
+            fs=config.fs,
+            nperseg=config.nperseg,
+            nfft=config.nfft,
+            boundary='zeros',
+        )
+        t1, denoised_signal = scipy.signal.istft(
+            (X[i, :, :, 0] + X[i, :, :, 1] * 1j) * preds[i, :, :, 0],
+            fs=config.fs,
+            nperseg=config.nperseg,
+            nfft=config.nfft,
+            boundary='zeros',
+        )
+        t1, denoised_noise = scipy.signal.istft(
+            (X[i, :, :, 0] + X[i, :, :, 1] * 1j) * preds[i, :, :, 1],
+            fs=config.fs,
+            nperseg=config.nperseg,
+            nfft=config.nfft,
+            boundary='zeros',
+        )
+
+    if result_dir is not None:
+        try:
+            np.savez(
+                os.path.join(result_dir, fname[i]),
+                noisy_signal=noisy_signal,
+                denoised_signal=denoised_signal,
+                denoised_noise=denoised_noise,
+                t=t1,
+            )
+        except FileNotFoundError:
+            os.makedirs(os.path.dirname(os.path.join(result_dir, fname[i])))
+            np.savez(
+                os.path.join(result_dir, fname[i]),
+                noisy_signal=noisy_signal,
+                denoised_signal=denoised_signal,
+                denoised_noise=denoised_noise,
+                t=t1,
+            )
+
+    if figure_dir is not None:
+
+        t_FT = np.linspace(config.time_range[0], config.time_range[1], X.shape[2])
+        f_FT = np.linspace(config.freq_range[0], config.freq_range[1], X.shape[1])
+
+        box = dict(boxstyle='round', facecolor='white', alpha=1)
+        text_loc = [0.05, 0.77]
+
+        plt.figure(i)
+        fig_size = plt.gcf().get_size_inches()
+        plt.gcf().set_size_inches(fig_size * [1, 1.2])
+        vmax = np.std(np.abs(X[i, :, :, 0] + X[i, :, :, 1] * 1j)) * 1.8
+
+        plt.subplot(311)
+        plt.pcolormesh(
+            t_FT,
+            f_FT,
+            np.abs(X[i, :, :, 0] + X[i, :, :, 1] * 1j),
+            vmin=0,
+            vmax=vmax,
+            shading='auto',
+            label='Noisy signal',
+        )
+        plt.gca().set_xticklabels([])
+        plt.text(
+            text_loc[0],
+            text_loc[1],
+            '(i)',
+            horizontalalignment='center',
+            transform=plt.gca().transAxes,
+            fontsize="medium",
+            fontweight="bold",
+            bbox=box,
+        )
+        plt.subplot(312)
+        plt.pcolormesh(
+            t_FT,
+            f_FT,
+            np.abs(X[i, :, :, 0] + X[i, :, :, 1] * 1j) * preds[i, :, :, 0],
+            vmin=0,
+            vmax=vmax,
+            shading='auto',
+            label='Recovered signal',
+        )
+        plt.gca().set_xticklabels([])
+        plt.ylabel("Frequency (Hz)", fontsize='large')
+        plt.text(
+            text_loc[0],
+            text_loc[1],
+            '(ii)',
+            horizontalalignment='center',
+            transform=plt.gca().transAxes,
+            fontsize="medium",
+            fontweight="bold",
+            bbox=box,
+        )
+        plt.subplot(313)
+        plt.pcolormesh(
+            t_FT,
+            f_FT,
+            np.abs(X[i, :, :, 0] + X[i, :, :, 1] * 1j) * preds[i, :, :, 1],
+            vmin=0,
+            vmax=vmax,
+            shading='auto',
+            label='Recovered noise',
+        )
+        plt.xlabel("Time (s)", fontsize='large')
+        plt.text(
+            text_loc[0],
+            text_loc[1],
+            '(iii)',
+            horizontalalignment='center',
+            transform=plt.gca().transAxes,
+            fontsize="medium",
+            fontweight="bold",
+            bbox=box,
+        )
+
+        try:
+            plt.savefig(os.path.join(figure_dir, fname[i].rstrip('.npz') + '_FT.png'), bbox_inches='tight')
+            # plt.savefig(os.path.join(figure_dir, fname[i].split('/')[-1].rstrip('.npz')+'_FT.pdf'), bbox_inches='tight')
+        except FileNotFoundError:
+            os.makedirs(os.path.dirname(os.path.join(figure_dir, fname[i].rstrip('.npz') + '_FT.png')), exist_ok=True)
+            plt.savefig(os.path.join(figure_dir, fname[i].rstrip('.npz') + '_FT.png'), bbox_inches='tight')
+            # plt.savefig(os.path.join(figure_dir, fname[i].split('/')[-1].rstrip('.npz')+'_FT.pdf'), bbox_inches='tight')
+        plt.close(i)
+
+        plt.figure(i)
+        fig_size = plt.gcf().get_size_inches()
+        plt.gcf().set_size_inches(fig_size * [1, 1.2])
+
+        ax4 = plt.subplot(311)
+        plt.plot(t1, noisy_signal, 'k', label='Noisy signal', linewidth=0.5)
+        plt.xlim([np.around(t1[0]), np.around(t1[-1])])
+        signal_ylim = [-np.max(np.abs(noisy_signal[100:-100])), np.max(np.abs(noisy_signal[100:-100]))]
+        plt.ylim(signal_ylim)
+        plt.gca().set_xticklabels([])
+        plt.legend(loc='lower left', fontsize='medium')
+        plt.text(
+            text_loc[0],
+            text_loc[1],
+            '(i)',
+            horizontalalignment='center',
+            transform=plt.gca().transAxes,
+            fontsize="medium",
+            fontweight="bold",
+            bbox=box,
+        )
+
+        ax5 = plt.subplot(312)
+        plt.plot(t1, denoised_signal, 'k', label='Recovered signal', linewidth=0.5)
+        plt.xlim([np.around(t1[0]), np.around(t1[-1])])
+        plt.ylim(signal_ylim)
+        plt.gca().set_xticklabels([])
+        plt.ylabel("Amplitude", fontsize='large')
+        plt.legend(loc='lower left', fontsize='medium')
+        plt.text(
+            text_loc[0],
+            text_loc[1],
+            '(ii)',
+            horizontalalignment='center',
+            transform=plt.gca().transAxes,
+            fontsize="medium",
+            fontweight="bold",
+            bbox=box,
+        )
+
+        plt.subplot(313)
+        plt.plot(t1, denoised_noise, 'k', label='Recovered noise', linewidth=0.5)
+        plt.xlim([np.around(t1[0]), np.around(t1[-1])])
+        plt.ylim(signal_ylim)
+        plt.xlabel("Time (s)", fontsize='large')
+        plt.legend(loc='lower left', fontsize='medium')
+        plt.text(
+            text_loc[0],
+            text_loc[1],
+            '(iii)',
+            horizontalalignment='center',
+            transform=plt.gca().transAxes,
+            fontsize="medium",
+            fontweight="bold",
+            bbox=box,
+        )
+
+        plt.savefig(os.path.join(figure_dir, fname[i].rstrip('.npz') + '_wave.png'), bbox_inches='tight')
+        # plt.savefig(os.path.join(figure_dir, fname[i].rstrip('.npz')+'_wave.pdf'), bbox_inches='tight')
+        plt.close(i)
+
+    return
+
+
+def save_results(mask, X, fname, t0, save_signal=True, save_noise=True, result_dir="results"):
+
+    config = Config()
+
+    if save_signal:
+        _, denoised_signal = scipy.signal.istft(
+            (X[..., 0] + X[..., 1] * 1j) * mask[..., 0],
+            fs=config.fs,
+            nperseg=config.nperseg,
+            nfft=config.nfft,
+            boundary='zeros',
+        )  # nbt, nch, nst, nt
+        denoised_signal = np.transpose(denoised_signal, [0, 3, 2, 1])  # nbt, nt, nst, nch,
+    if save_noise:
+        _, denoised_noise = scipy.signal.istft(
+            (X[..., 0] + X[..., 1] * 1j) * mask[..., 1],
+            fs=config.fs,
+            nperseg=config.nperseg,
+            nfft=config.nfft,
+            boundary='zeros',
+        )
+        denoised_noise = np.transpose(denoised_noise, [0, 3, 2, 1])
+
+    if not os.path.exists(result_dir):
+        os.makedirs(result_dir)
+
+    for i in range(len(X)):
+        np.savez(
+            os.path.join(result_dir, fname[i]),
+            data=denoised_signal[i] if save_signal else None,
+            noise=denoised_noise[i] if save_noise else None,
+            t0=t0[i],
+        )
+
+
+def plot_figures(mask, X, fname, figure_dir="figures"):
+
+    config = Config()
+
+    # plot the last channel
+    mask = mask[-1, -1, ...]  # nch, nst, nf, nt, 2 => nf, nt, 2
+    X = X[-1, -1, ...]
+
+    t1, noisy_signal = scipy.signal.istft(
+        (X[..., 0] + X[..., 1] * 1j),
+        fs=config.fs,
+        nperseg=config.nperseg,
+        nfft=config.nfft,
+        boundary='zeros',
+    )
+    t1, denoised_signal = scipy.signal.istft(
+        (X[..., 0] + X[..., 1] * 1j) * mask[..., 0],
+        fs=config.fs,
+        nperseg=config.nperseg,
+        nfft=config.nfft,
+        boundary='zeros',
+    )
+    t1, denoised_noise = scipy.signal.istft(
+        (X[..., 0] + X[..., 1] * 1j) * mask[..., 1],
+        fs=config.fs,
+        nperseg=config.nperseg,
+        nfft=config.nfft,
+        boundary='zeros',
+    )
+
+    if not os.path.exists(figure_dir):
+        os.makedirs(figure_dir)
+
+    t_FT = np.linspace(config.time_range[0], config.time_range[1], X.shape[1])
+    f_FT = np.linspace(config.freq_range[0], config.freq_range[1], X.shape[0])
+
+    box = dict(boxstyle='round', facecolor='white', alpha=1)
+    text_loc = [0.05, 0.77]
+
+    plt.figure()
+    fig_size = plt.gcf().get_size_inches()
+    plt.gcf().set_size_inches(fig_size * [1, 1.2])
+    vmax = np.std(np.abs(X[:, :, 0] + X[:, :, 1] * 1j)) * 1.8
+
+    plt.subplot(311)
+    plt.pcolormesh(
+        t_FT,
+        f_FT,
+        np.abs(X[:, :, 0] + X[:, :, 1] * 1j),
+        vmin=0,
+        vmax=vmax,
+        shading='auto',
+        label='Noisy signal',
+    )
+    plt.gca().set_xticklabels([])
+    plt.text(
+        text_loc[0],
+        text_loc[1],
+        '(i)',
+        horizontalalignment='center',
+        transform=plt.gca().transAxes,
+        fontsize="medium",
+        fontweight="bold",
+        bbox=box,
+    )
+    plt.subplot(312)
+    plt.pcolormesh(
+        t_FT,
+        f_FT,
+        np.abs(X[:, :, 0] + X[:, :, 1] * 1j) * mask[:, :, 0],
+        vmin=0,
+        vmax=vmax,
+        shading='auto',
+        label='Recovered signal',
+    )
+    plt.gca().set_xticklabels([])
+    plt.ylabel("Frequency (Hz)", fontsize='large')
+    plt.text(
+        text_loc[0],
+        text_loc[1],
+        '(ii)',
+        horizontalalignment='center',
+        transform=plt.gca().transAxes,
+        fontsize="medium",
+        fontweight="bold",
+        bbox=box,
+    )
+    plt.subplot(313)
+    plt.pcolormesh(
+        t_FT,
+        f_FT,
+        np.abs(X[:, :, 0] + X[:, :, 1] * 1j) * mask[:, :, 1],
+        vmin=0,
+        vmax=vmax,
+        shading='auto',
+        label='Recovered noise',
+    )
+    plt.xlabel("Time (s)", fontsize='large')
+    plt.text(
+        text_loc[0],
+        text_loc[1],
+        '(iii)',
+        horizontalalignment='center',
+        transform=plt.gca().transAxes,
+        fontsize="medium",
+        fontweight="bold",
+        bbox=box,
+    )
+
+    try:
+        plt.savefig(os.path.join(figure_dir, fname.rstrip('.npz') + '_FT.png'), bbox_inches='tight')
+        # plt.savefig(os.path.join(figure_dir, fname[i].split('/')[-1].rstrip('.npz')+'_FT.pdf'), bbox_inches='tight')
+    except FileNotFoundError:
+        os.makedirs(os.path.dirname(os.path.join(figure_dir, fname.rstrip('.npz') + '_FT.png')), exist_ok=True)
+        plt.savefig(os.path.join(figure_dir, fname.rstrip('.npz') + '_FT.png'), bbox_inches='tight')
+        # plt.savefig(os.path.join(figure_dir, fname[i].split('/')[-1].rstrip('.npz')+'_FT.pdf'), bbox_inches='tight')
+    plt.close()
+
+    plt.figure()
+    fig_size = plt.gcf().get_size_inches()
+    plt.gcf().set_size_inches(fig_size * [1, 1.2])
+
+    ax4 = plt.subplot(311)
+    plt.plot(t1, noisy_signal, 'k', label='Noisy signal', linewidth=0.5)
+    plt.xlim([np.around(t1[0]), np.around(t1[-1])])
+    signal_ylim = [-np.max(np.abs(noisy_signal)), np.max(np.abs(noisy_signal))]
+    if signal_ylim[0] != signal_ylim[1]:
+        plt.ylim(signal_ylim)
+    plt.gca().set_xticklabels([])
+    plt.legend(loc='lower left', fontsize='medium')
+    plt.text(
+        text_loc[0],
+        text_loc[1],
+        '(i)',
+        horizontalalignment='center',
+        transform=plt.gca().transAxes,
+        fontsize="medium",
+        fontweight="bold",
+        bbox=box,
+    )
+
+    ax5 = plt.subplot(312)
+    plt.plot(t1, denoised_signal, 'k', label='Recovered signal', linewidth=0.5)
+    plt.xlim([np.around(t1[0]), np.around(t1[-1])])
+    if signal_ylim[0] != signal_ylim[1]:
+        plt.ylim(signal_ylim)
+    plt.gca().set_xticklabels([])
+    plt.ylabel("Amplitude", fontsize='large')
+    plt.legend(loc='lower left', fontsize='medium')
+    plt.text(
+        text_loc[0],
+        text_loc[1],
+        '(ii)',
+        horizontalalignment='center',
+        transform=plt.gca().transAxes,
+        fontsize="medium",
+        fontweight="bold",
+        bbox=box,
+    )
+
+    plt.subplot(313)
+    plt.plot(t1, denoised_noise, 'k', label='Recovered noise', linewidth=0.5)
+    plt.xlim([np.around(t1[0]), np.around(t1[-1])])
+    if signal_ylim[0] != signal_ylim[1]:
+        plt.ylim(signal_ylim)
+    plt.xlabel("Time (s)", fontsize='large')
+    plt.legend(loc='lower left', fontsize='medium')
+    plt.text(
+        text_loc[0],
+        text_loc[1],
+        '(iii)',
+        horizontalalignment='center',
+        transform=plt.gca().transAxes,
+        fontsize="medium",
+        fontweight="bold",
+        bbox=box,
+    )
+
+    plt.savefig(os.path.join(figure_dir, fname.rstrip('.npz') + '_wave.png'), bbox_inches='tight')
+    # plt.savefig(os.path.join(figure_dir, fname[i].rstrip('.npz')+'_wave.pdf'), bbox_inches='tight')
+    plt.close()
+
+    return
+
+
+if __name__ == "__main__":
+    pass

docs/README.md

@@ -0,0 +1,60 @@
+---
+title: DeepDenoiser
+emoji: 🌊
+colorFrom: purple
+colorTo: blue
+sdk: docker
+pinned: false
+---
+
+# DeepDenoiser: Seismic Signal Denoising and Decomposition Using Deep Neural Networks
+
+[![](https://github.com/AI4EPS/DeepDenoiser/workflows/documentation/badge.svg)](https://ai4eps.github.io/DeepDenoiser)
+## 1.  Install [miniconda](https://docs.conda.io/en/latest/miniconda.html) and requirements
+- Download DeepDenoiser repository
+```bash
+git clone https://github.com/wayneweiqiang/DeeoDenoiser.git
+cd DeepDenoiser
+```
+- Install to default environment
+```bash
+conda env update -f=env.yml -n base
+```
+- Install to "deepdenoiser" virtual envirionment
+```bash
+conda env create -f env.yml
+conda activate deepdenoiser
+```
+
+## 2. Pre-trained model
+Located in directory: **model/190614-104802**
+
+## 3. Related papers
+- Zhu, Weiqiang, S. Mostafa Mousavi, and Gregory C. Beroza. "Seismic Signal Denoising and Decomposition Using Deep Neural Networks." arXiv preprint arXiv:1811.02695 (2018).
+
+## 4. Interactive example
+See details in the [notebook](https://github.com/wayneweiqiang/DeepDenoiser/blob/master/docs/example_interactive.ipynb): [example_interactive.ipynb](example_interactive.ipynb)
+
+
+## 5. Batch prediction
+See details in the [notebook](https://github.com/wayneweiqiang/DeepDenoiser/blob/master/docs/example_batch_prediction.ipynb): [example_batch_prediction.ipynb](example_batch_prediction.ipynb)
+## 6. Train
+### Data format
+
+Required: two csv files for signal and noise, corresponding directories of the npz files.
+
+The csv file contains four columns: "fname", "itp", "channels"
+
+The npz file contains four variable: "data", "itp",  "channels"
+
+The shape of "data" variables has a shape of 9001 x 3
+
+The variables "itp" is the data points of first P arrival times.
+
+Note: In the demo data, for simplicity we use the waveform before itp as noise samples, so the train_noise_list is same as train_signal_list here.
+
+~~~bash
+python deepdenoiser/train.py --mode=train --train_signal_dir=./Dataset/train --train_signal_list=./Dataset/train.csv --train_noise_dir=./Dataset/train --train_noise_list=./Dataset/train.csv --batch_size=20
+~~~
+
+Please let us know of any bugs found in the code. Suggestions and collaborations are welcomed

env.yml

@@ -0,0 +1,19 @@
+name: deepdenoiser
+channels:
+  - defaults
+  - conda-forge
+dependencies:
+  - python=3.7
+  - numpy
+  - scipy
+  - matplotlib
+  - pandas
+  - scikit-learn
+  - tqdm
+  - obspy
+  - uvicorn
+  - fastapi
+  - kafka-python
+  - tensorflow
+
+

mkdocs.yml

@@ -0,0 +1,18 @@
+site_name: "DeepDenoiser"
+site_description: 'DeepDenoiser: Seismic Signal Denoising and Decomposition Using Deep Neural Networks'
+site_author: 'Weiqiang Zhu'
+docs_dir: docs/
+repo_name: 'wayneweiqiang/DeepDenoiser'
+repo_url: 'https://github.com/wayneweiqiang/DeepDenoiser'
+nav:
+    - Overview: README.md
+    - Interactive Example: example_interactive.ipynb
+    - Batch Prediction: example_batch_prediction.ipynb
+theme:
+  name: 'material'
+plugins:
+  - mkdocs-jupyter
+extra:
+  analytics:
+    provider: google
+    property: G-FMMP8CQRDZ

requirements.txt

@@ -0,0 +1,5 @@
+tensorflow
+matplotlib
+scipy
+pandas
+tqdm