Sync API from main repo

fast.py

@@ -37,8 +37,7 @@ app.state.model = None  # Initialize as None, load on first request
 def root():
     return dict(greeting="Hello")
 
-@app.post("/predict")
-async def predict(model_name: str, filepath_csv: UploadFile = File(...)):
+def model_loader(model_name):
     # Load model if not already loaded
     model_path = MODEL_DIR / f"{model_name}"
     encoder_name = encoder_from_model(model_name)
@@ -46,20 +45,23 @@ async def predict(model_name: str, filepath_csv: UploadFile = File(...)):
 
     # if model in model_path, load it, otherwise download it from HF
     if model_name not in model_cache:
-        # print("model_name", model_name)
-        # print("model_path", model_path)
         try:
             if not model_path.exists():
                 # Convert downloaded paths to Path objects
                 model_path = Path(hf_hub_download(repo_id=HF_REPO_ID, filename=f"{model_name}", cache_dir=CACHE_DIR))
                 encoder_path = Path(hf_hub_download(repo_id=HF_REPO_ID, filename=f"{encoder_name}", cache_dir=CACHE_DIR))
-            # print("model_path", model_path)
             model_cache[model_name] = load_model_by_type(model_path)  # Ensure string path for loading
             encoder_cache[model_name] = encoder_path
         except Exception as e:
             print(f"Error loading model: {str(e)}")  # Add debug print
             raise HTTPException(status_code=404, detail=f"Model {model_name} not found: {str(e)}")
-    model = app.state.model = model_cache[model_name]
+    return model_cache[model_name]
+
+
+@app.post("/predict")
+async def predict(model_name: str, filepath_csv: UploadFile = File(...)):
+    
+    model = app.state.model = model_loader(model_name)
 
     # Read the uploaded CSV file
     file_content = await filepath_csv.read()
@@ -68,6 +70,15 @@ async def predict(model_name: str, filepath_csv: UploadFile = File(...)):
     
     # Decode prediction using absolute path
     
-    y_pred = label_decoding(value=y_pred[0], path=encoder_cache[model_name])
+    y_pred = label_decoding(values=y_pred, path=encoder_cache[model_name])
     
     return {"prediction": y_pred}
+
+
+# @app.post("/predict_multibeats")
+# async def predict_multibeats(model_name: str, filepath_csv: UploadFile = File(...)):
+#     # Read the uploaded CSV file
+#     file_content = await filepath_csv.read()
+#     X = pd.read_csv(StringIO(file_content.decode('utf-8')))
+#     y_pred = model.predict_with_pipeline(X)
+#     return {"prediction": y_pred}

preproc.py

@@ -1,12 +1,18 @@
 from tslearn.utils import to_time_series_dataset
 from tslearn.preprocessing import TimeSeriesScalerMeanVariance
 import pickle
+from wfdb import rdrecord, rdann, processing
+from sklearn import preprocessing
+from scipy.signal import resample
 
-def preproc_single(X):
+import numpy as np
+import pandas as pd
+
+def preproc(X):
     # to be called in inference/api
     in_shape = X.shape
-    if X.shape != (1, 180):
-        print('File shape is not (1, 180) but ', in_shape)
+    if X.shape[1] != 180:
+        print('File shape is not (n, 180) but ', in_shape)
 
     X = to_time_series_dataset(X)
     X = X.reshape(in_shape[0], -1)
@@ -14,8 +20,74 @@ def preproc_single(X):
     X = scaler.fit_transform(X)
     return X.reshape(in_shape)
 
-def label_decoding(value, path):
+def apple_csv_to_data(filepath_csv):
+    # extract sampling rate 
+    with open(filepath_csv, 'r') as file:
+        for il,line in enumerate(file):
+            if line.startswith("Sample Rate"):
+                # Extract the sample rate
+                sample_rate = int(line.split(",")[1].split()[0])  # Split and get the numerical part
+                print(f"Sample Rate: {sample_rate}")
+                break
+            if il > 30:
+                print("Could not find sample rate in first 30 lines")
+                return None, None  
+    X = pd.read_csv(filepath_csv, skiprows=14, header=None)
+    return X, sample_rate
+
+def apple_trim_join(X, sample_rate=512, ns=2):
+    # There should be a less horrible way of doing this
+    # Ignore first two and last two seconds, that tend to be noisy --> 26 seconds ecg
+    X[1] = X[1].fillna(0)
+    X = X[0] + X[1] / (10 ** (X[1].astype(str).str.len() - 2)) # Ignoring the trailing ".0"
+    print(f"Ignoring first and last {ns} seconds")
+    X = X[ns*sample_rate:-ns*sample_rate].to_frame().T
+    X = X.iloc[0].to_numpy()
+    return X
+
+def apple_extract_beats(X, sample_rate=512):
+    X = apple_trim_join(X, sample_rate=sample_rate, ns=3)
+    # Scale and remove nans (should not happen anymore)
+    X = preprocessing.scale(X[~np.isnan(X)])
+
+    # I tried to hack the detection to make it learn peaks and
+    # not go with default, but it doesn't work
+    # I have tried:
+    # - Hardwiring n_calib_beats (not possible from user side)
+    #   to a lower number (5, 3).
+    # - Setting qrs_width to lower and higher values
+    # - Relax the correlation requirement to Rikers wavelet
+    # Maybe explore correlation with more robust wavelets
+    # wavelet = pywt.Wavelet('db4')
+    # (lib/python3.10/site-packages/wfdb/processing/qrs.py)
+ 
+    # Conf = processing.XQRS.Conf(qrs_width=0.1)
+    # qrs = processing.XQRS(sig = X,fs = sample_rate, conf=Conf)
+    # wfdb library doesn't allow to set n_calib_beats
+
+    qrs = processing.XQRS(sig = X,fs = sample_rate)
+    qrs.detect()
+    peaks = qrs.qrs_inds
+    print("Number of beats detected: ", len(peaks))
+    target_length = 180
+    beats = np.zeros((len(peaks), target_length))
+    
+    for i, peak in enumerate(peaks[1:-1]):
+        rr_interval = peaks[i + 1] - peaks[i]  # Distance to the next peak
+        window_size = int(rr_interval * 1.2)  # Extend by 20% to capture full P-QRS-T cycle
+        # Define the dynamic window around the R-peak
+        start = max(0, peak - window_size // 2)
+        end = min(len(X), peak + window_size // 2)
+        beat = resample(X[start:end], target_length)
+        beats[i] = beat
+    return beats
+
+def save_beats_csv(beats, filepath_csv):
+    pd.DataFrame(beats).to_csv(filepath_csv, index=False)
+
+def label_decoding(values, path):
     with open(path, "rb") as f:
         mapping = pickle.load(f)
     inverse_mapping = {v: k for k, v in mapping.items()}
-    return inverse_mapping[value]
+    # return inverse_mapping[values]
+    return [inverse_mapping[value] for value in values]

requirements.txt

@@ -5,6 +5,7 @@ huggingface-hub
 pandas==2.2.3
 numpy==1.26.4
 scikit-learn==1.2.2
+scipy
 tslearn
 tensorflow
 python-multipart

wrappers.py

@@ -1,5 +1,5 @@
 import numpy as np
-from preproc import preproc_single
+from preproc import preproc
 
 class BaseModelWrapper:
     def __init__(self, model):
@@ -7,7 +7,7 @@ class BaseModelWrapper:
 
     def preprocess(self, data):
         """Default preprocessing (can be overridden)."""
-        return preproc_single(data)
+        return preproc(data)
 
     def predict(self, data):
         """Call the model's prediction."""
@@ -28,7 +28,7 @@ class BaseModelWrapper:
 class LSTMWrapper(BaseModelWrapper):
     def preprocess(self, data):
         # LSTM requires additional dimension expansion
-        data = preproc_single(data)
+        data = preproc(data)
         return np.expand_dims(data, axis=1)  # Add time-step dimension
 
     def predict(self, data):