Implement proper labelling

Author: knoellle

crates/control/src/fall_state_detection.rs

@@ -62,6 +62,8 @@ impl FallStateDetection {
     }
 
     pub fn cycle(&mut self, context: CycleContext<impl PathsInterface>) -> Result<MainOutputs> {
+        // TODO: hadle primary state unstiff
+
         // let cycle_start = context.cycle_time.start_time;
         // let inertial_measurement_unit = context.sensor_data.inertial_measurement_unit;
         // let (roll, pitch, _) = context.robot_orientation.inner.euler_angles();

tools/machine-learning/fall_detection/data_loading/load_data.py

@@ -58,12 +58,22 @@ def unwrap(self, data: dict | str) -> str:
         raise ValueError(f"did not expect fall state {data}")
 
 
+class HasGroundContactUnwrapper(Unwrapper):
+    pass
+
+
+class PrimaryStateUnwrapper(Unwrapper):
+    pass
+
+
 OUTPUTS = {
     "Control.main_outputs.sensor_data": SensorDataUnwrapper(),
     "Control.main_outputs.robot_orientation": RobotOrientationUnwrapper(),
     "Control.main_outputs.center_of_mass": CenterOfMassUnwrapper(),
     "Control.main_outputs.zero_moment_point": ZeroMomentPointUnwrapper(),
     "Control.main_outputs.fall_state": FallStateUnwrapper(),
+    "Control.main_outputs.has_ground_contact": HasGroundContactUnwrapper(),
+    "Control.main_outputs.primary_state": PrimaryStateUnwrapper(),
 }
 
 

tools/machine-learning/fall_detection/dataset/dataset.py

@@ -28,6 +28,7 @@ def __init__(
 
         self.labeller = PseudoLabeller()
         self.features = features
+        print(dataframe.columns)
 
         number_of_nulls = (
             dataframe.select(features)
@@ -77,13 +78,15 @@ def to_windowed(
         )
         self.input_data = windows.select(self.features)
         # windowed_labels = windowed_dataframe.agg(pl.col("labels")).drop("index")
-        self.labels = pl.Series(
-            [
-                # todo: state prediction label index
-                row[-1]
-                # print(row)
-                for row in windows.get_column("labels")
-            ]
+        self.labels = pl.DataFrame(
+            {
+                "labels": [
+                    # todo: state prediction label index
+                    row[-1]
+                    # print(row)
+                    for row in windows.get_column("labels")
+                ]
+            }
         )
 
     def __len__(self) -> int:
@@ -92,6 +95,9 @@ def __len__(self) -> int:
     def n_features(self) -> int:
         return self.input_data.size(1)
 
+    def n_classes(self) -> int:
+        return len(self.labeller.label_type)
+
     def __getitem__(self, index: int) -> tuple[tf.Tensor, tf.Tensor]:
         mask = self.groups == index
         return self.input_data[mask], self.labels[mask]

tools/machine-learning/fall_detection/dataset/pseudo_labels.py

@@ -1,33 +1,46 @@
-from dataclasses import dataclass
+from enum import Enum
+
 import polars as pl
 
+from dataclasses import dataclass
+
 
 @dataclass
 class PseudoLabelParameters:
-    pitch_threshold: float = 1.0
+    falling_threshold: float = 0.6
+    fallen_threshold: float = 1.3
 
 
-LABELS = {
-    0: "Other",
-    1: "Fallen",
-}
+class Label(int, Enum):
+    Upright = 0
+    Falling = 1
+    Fallen = 2
 
 
 class PseudoLabeller:
+    label_type = Label
+
     def __init__(self, parameters: PseudoLabelParameters | None = None):
         self.parameters = parameters or PseudoLabelParameters()
 
     def generate_labels(self, data: pl.DataFrame) -> pl.Series:
+        has_ground_contact = pl.col("Control.main_outputs.has_ground_contact")
+        pitch = pl.col("Control.main_outputs.robot_orientation.pitch")
+        # print(data.select(primary_state))
         return (
             data.select(
-                self.map_to_schema(
-                    pl.col("Control.main_outputs.robot_orientation.pitch").abs()
-                    > self.parameters.pitch_threshold
+                pl.when(has_ground_contact)
+                .then(
+                    pl.when(pitch.abs() > self.parameters.falling_threshold)
+                    .then(pl.lit(Label.Falling))
+                    .otherwise(pl.lit(Label.Upright))
+                )
+                .otherwise(
+                    pl.when(pitch.abs() > self.parameters.fallen_threshold)
+                    .then(pl.lit(Label.Fallen))
+                    .otherwise(pl.lit(Label.Falling))
                 )
             )
             .to_series()
             .alias("labels")
         )
-
-    def map_to_schema(self, expression: pl.Expr) -> pl.Expr:
-        return pl.when(expression).then(pl.lit(1)).otherwise(pl.lit(0))

tools/machine-learning/fall_detection/scripts/inspect-data.py

@@ -2,20 +2,56 @@
 import plotly.express as px
 import plotly.io as pio
 from data_loading import load
+from dataset import FallenDataset
 
 
 def main():
     pio.renderers.default = "browser"
-    data = load("data.parquet")
-    # px.scatter(
-    #     data, x="time", y="Control.main_outputs.fall_state", color="robot_identifier"
-    # ).show()
+
+    df = load("data.parquet")
+    dataset = FallenDataset(
+        df,
+        group_keys=["robot_identifier", "match_identifier"],
+        features=[
+            pl.col("Control.main_outputs.robot_orientation.pitch"),
+            pl.col("Control.main_outputs.robot_orientation.roll"),
+            pl.col("Control.main_outputs.robot_orientation.yaw"),
+            pl.col("Control.main_outputs.has_ground_contact"),
+        ],
+    )
+    dataset.to_windowed(window_stride=1 / 83)
+
+    df = (
+        dataset.input_data.hstack(dataset.labels)
+        .with_row_index()
+        .hstack(
+            dataset.input_data.select(
+                pl.col(
+                    "Control.main_outputs.robot_orientation.pitch"
+                ).list.last()
+            ).rename({"Control.main_outputs.robot_orientation.pitch": "pitch"})
+        )
+    )
+    print(df)
+    # print(df[0, 0].shape)
+
     # px.scatter(
-    #     data.filter(pl.col("robot_identifier") == "10.1.24.33"),
+    #     data,
     #     x="time",
-    #     y="Control.main_outputs.robot_orientation.pitch",
-    #     color="Control.main_outputs.fall_state",
+    #     y="Control.main_outputs.fall_state",
+    #     color="robot_identifier",
     # ).show()
+    print(
+        df.select(
+            pl.col("Control.main_outputs.robot_orientation.pitch").list.last()
+        )
+    )
+    px.scatter(
+        df,  # .filter(pl.col("robot_identifier") == "10.1.24.33"),
+        x="index",
+        y="pitch",
+        color="labels",
+    ).show()
 
 
 if __name__ == "__main__":

tools/machine-learning/fall_detection/scripts/train_sequential.py

@@ -134,8 +134,6 @@ def split_data(input_data, labels):
 
 
 if __name__ == "__main__":
-    labels = ["Unknown", "Upright", "Falling", "Fallen"]
-
     df = load("data.parquet")
     dataset = FallenDataset(
         df,
@@ -149,14 +147,14 @@ def split_data(input_data, labels):
     )
     dataset.to_windowed()
 
-    model = build_model(len(labels))
+    model = build_model(dataset.n_classes())
 
     input_data = dataset.get_input_tensor()
     data_labels = dataset.get_labels_tensor()
     (x_train, y_train, x_test, y_test) = split_data(input_data, data_labels)
 
-    y_train = keras.utils.to_categorical(y_train, len(labels))
-    y_test = keras.utils.to_categorical(y_test, len(labels))
+    y_train = keras.utils.to_categorical(y_train, dataset.n_classes())
+    y_test = keras.utils.to_categorical(y_test, dataset.n_classes())
 
     train_model(model, x_train, y_train, x_test, y_test)