Fix start and end times for reading csv files (#51)

* Provide start and end times for reading CSV files if info json files don't contain these two timestamps. --------- Co-authored-by: “Aidan <“[email protected]”>
OxWearables · Jul 16, 2024 · e3e7684 · e3e7684
1 parent 949236d
commit e3e7684
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/src/asleep/get_sleep.py b/src/asleep/get_sleep.py
@@ -35,6 +35,9 @@
 NON_WEAR_THRESHOLD = 3  # H
 NON_WEAR_PREDICTION_FLAG = -1
 
+START_TIME_IDX = 0
+END_TIME_IDX = -1
+
 
 def load_model(model_path, force_download=False):
     """ Load trained model. Download if not exists. """
@@ -69,8 +72,8 @@ def get_parsed_data(raw_data_path, info_data_path, resample_hz, args):
         data = data.reset_index()
 
         # apply time shift
-        start_time = pd.to_datetime(info['StartTime'])
-        end_time = pd.to_datetime(info['EndTime'])
+        start_time = pd.to_datetime(info.get("StartTime", data["time"].iloc[START_TIME_IDX]))
+        end_time = pd.to_datetime(info.get("EndTime", data["time"].iloc[END_TIME_IDX]))
         info['StartTime'] = start_time + datetime.timedelta(hours=time_shift)
         info['EndTime'] = end_time + datetime.timedelta(hours=time_shift)
         info['StartTime'] = info['StartTime'].strftime('%Y-%m-%d %H:%M:%S')