update

deepdenoiser/app.py

@@ -6,15 +6,14 @@ from typing import Any, AnyStr, Dict, List, NamedTuple, Union
 
 import numpy as np
 import requests
+import scipy
 import tensorflow as tf
 from fastapi import FastAPI
 from kafka import KafkaProducer
+from model import UNet
 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__), ".."))
@@ -105,7 +104,6 @@ def normalize_batch(data, window=200):
 
 
 def get_prediction(meta):
-
     FS = 100
     NPERSEG = 30
     NFFT = 60
@@ -125,7 +123,7 @@ def get_prediction(meta):
 
     # 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')
+    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
@@ -139,7 +137,7 @@ def get_prediction(meta):
         fs=FS,
         nperseg=NPERSEG,
         nfft=NFFT,
-        boundary='zeros',
+        boundary="zeros",
     )
     # _, denoised_noise = scipy.signal.istft(
     #     (noisy_signal[..., 0] + noisy_signal[..., 1] * 1j) * preds[..., 1],
@@ -169,7 +167,6 @@ class Data(BaseModel):
 
 @app.post("/predict")
 def predict(data: Data):
-
     denoised = get_prediction(data)
 
     return denoised
@@ -177,4 +174,9 @@ def predict(data: Data):
 
 @app.get("/healthz")
 def healthz():
-    return {"status": "ok"}
+    return {"status": "ok"}
+
+
+@app.get("/")
+def read_root():
+    return {"Hello": "DeepDenoiser!"}