Fix animals dataset labels

yupidevs · Nov 15, 2022 · 6c34693 · 6c34693
1 parent 09a40a4
commit 6c34693
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/recipes/animals.py b/recipes/animals.py
@@ -17,6 +17,8 @@
     "master/Starkey_OR_Main_Telemetry_1993-1996_Data.txt"
 )
 
+LABELS = {"E": "Elk", "D": "Deer", "C": "Cattle"}
+
 
 def build() -> Tuple[List[Trajectory], List[Any]]:
     raw_dir = _fetch_raw_data()
@@ -56,7 +58,7 @@ def _process_animal(animal_rows) -> Tuple[Trajectory, str]:
         lat.append(_lat)
         long.append(_long)
 
-    return Trajectory(x=long, y=lat, t=time), label
+    return Trajectory(x=long, y=lat, t=time), LABELS[label]
 
 
 def _yupify(raw_dir) -> Tuple[List[Trajectory], List[str]]: