Update app.py

app.py

@@ -1,5 +1,4 @@
 import numpy as np
-
 from PIL import Image, ImageFilter
 
 import torch
@@ -13,6 +12,9 @@ from fastapi.responses import JSONResponse
 import uvicorn
 from fastapi.middleware.cors import CORSMiddleware
 
+from numpy.fft import rfft, irfft, rfftfreq
+import pandas as pd
+
 import datetime 
 import pytz
 
@@ -126,7 +128,7 @@ async def predict(file: UploadFile = File(...)):
     # image.save('input.jpg')
     
     image_rz = image.resize((256,256))
-    # image_rz.save('input_resize.jpg')
+    # image_rz.save('input1.jpg')
     
     print(f"image after fill bg : {image.size, type(image)}\n")
     
@@ -136,10 +138,114 @@ async def predict(file: UploadFile = File(...)):
     parkinson_predict = predict_parkinson(image)
     print(f"parkinson predict : {parkinson_predict}\n")
     
+    curr_time = datetime.datetime.now()
     print(f"end time : {datetime.datetime.now(pytz.timezone('Asia/Bangkok'))}")
     
     return JSONResponse([spiral_predict, parkinson_predict])
 
+SAMPLING_RATE  = 100
+DURATION = 10
+time_axis = np.linspace(1/SAMPLING_RATE, DURATION, SAMPLING_RATE*DURATION)
+print(f"len time axis : {len(time_axis)}")
+
+def spooled_tempfile_to_string(spooled_tempfile) -> str:
+    spooled_tempfile.seek(0)
+    
+    raw_content = spooled_tempfile.read()
+    
+    if isinstance(raw_content, bytes):
+        content = raw_content.decode('utf-8')
+    
+    return content
+
+def encode(raw_content) :
+    content = raw_content[raw_content.find('[') + 1 : raw_content.rfind(']')]
+
+    data_list = []
+
+    i = 0
+
+    while i < len(content) :
+        idx_open = content.find('[', i)
+        idx_close = content.find(']', idx_open)
+
+        if idx_open==-1 or idx_close==-1 : break
+
+        txt = content[idx_open+1 : idx_close]
+        row = [float(val) for val in txt.split(',')]
+
+        data_list.append(row)
+
+        i = idx_close + 1
+
+    df = pd.DataFrame(data_list, columns=['roll','pitch','yaw'], index=time_axis)
+
+    return df
+
+class Fourier :
+    def __init__(self, signal, sampling_rate, duration) :
+        self.signal = np.array(signal)
+        self.sampling_rate = sampling_rate
+        self.duration = duration
+
+        self.time = np.linspace(1/sampling_rate, duration, len(signal))
+
+        self.freq = self.get_frequencies()
+        self.raw_amplitudes = self.get_amplitudes()
+        self.amplitudes = self.norm_amplitudes()
+
+    def get_frequencies(self) :
+        return rfftfreq(len(self.signal), 1/self.sampling_rate)
+
+    def get_amplitudes(self) :
+        return rfft(self.signal)
+
+    def norm_amplitudes(self) :
+        return 2*np.abs(self.raw_amplitudes) / len(self.signal)
+
+def predict_tremor(signal, sampling_rate, duration) :
+    fourier = Fourier(signal, sampling_rate, duration)
+    amp = fourier.amplitudes
+
+    start_freq = 0.5
+    start_idx = np.argwhere(fourier.freq >= start_freq)[0][0]
+
+    idx_3hz = np.argwhere(fourier.freq >= 3)[0][0]
+    idx_6hz = np.argwhere(fourier.freq >= 6)[0][0]
+
+    max_amp_parkinson_range = np.max(amp[idx_3hz:idx_6hz+1])
+    if max_amp_parkinson_range > 1.5 : return 2
+    
+    max_amp_tremor_range = np.max(amp[start_idx:])
+    if max_amp_tremor_range > 1 : return 1
+
+    return 0
+
+@app.post("/predict_shake/")
+async def predict_shake(file: UploadFile = File(...)):
+    raw_content = spooled_tempfile_to_string(file.file)
+    
+    print(f"len raw_content : {len(raw_content)}")
+    print(f"raw_content: {raw_content[:10]}")
+    
+    df = encode(raw_content)
+    print(f"df shape : {df.shape}")
+
+    level_roll = predict_tremor(df['roll'], SAMPLING_RATE, DURATION)
+    level_pitch = predict_tremor(df['pitch'], SAMPLING_RATE, DURATION)
+    level_yaw = predict_tremor(df['yaw'], SAMPLING_RATE, DURATION)
+    
+    print(f"roll: {level_roll}, pitch: {level_pitch}, yaw: {level_yaw}")
+    
+    highest_level = np.max([level_roll, level_pitch, level_yaw])
+
+    level_names = ['low', 'mid', 'high']
+
+    curr_time = datetime.datetime.now()
+    print(f"curr time : {curr_time}")
+    
+    return level_names[highest_level]
+
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=8000)