Use os.path.splitext instead of rstrip to remove .txt

Fixes #51
kootenpv · Feb 14, 2019 · dffc25b · dffc25b
1 parent a950343
commit dffc25b
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/whereami/get_data.py b/whereami/get_data.py
@@ -39,5 +39,5 @@ def get_train_data(folder=None):
  for line in f:
  data.append(json.loads(line))
  X.extend(data)
- y.extend([fname.rstrip(".txt")] * len(data))
+ y.extend([os.path.splitext(fname)[0]] * len(data))
  return X, y