Add "feature_names" to landmarks.npz; update corresponding code as well

noddingpigeon/video.py

@@ -17,13 +17,33 @@ class VideoSegment(Enum):
     LAST = auto()
 
 
+class Feature(Enum):
+    X_MIN = auto()
+    Y_MIN = auto()
+    WIDTH = auto()
+    HEIGHT = auto()
+    RIGHT_EYE_X = auto()
+    LEFT_EYE_X = auto()
+    NOSE_X = auto()
+    MOUTH_X = auto()
+    RIGHT_EAR_X = auto()
+    LEFT_EAR_X = auto()
+    RIGHT_EYE_Y = auto()
+    LEFT_EYE_Y = auto()
+    NOSE_Y = auto()
+    MOUTH_Y = auto()
+    RIGHT_EAR_Y = auto()
+    LEFT_EAR_Y = auto()
+
+
 def video_to_landmarks(
         video_path: Optional[Union[int, str]],
         max_num_frames: Optional[int] = Config.seq_length,
         video_segment: VideoSegment = VideoSegment.BEGINNING,
         end_padding: bool = True,
         drop_consecutive_duplicates: bool = False
 ) -> List[List[float]]:
+    assert Config.num_original_features == len(Feature)
     assert video_segment in VideoSegment
     video_path = video_path if video_path else 0  # For 0, webcam will be used.
     font_scale = 0.6

training/collect_landmarks.py

@@ -7,7 +7,7 @@
 
 import numpy as np
 
-from noddingpigeon.video import video_to_landmarks
+from noddingpigeon.video import Feature, video_to_landmarks
 from noddingpigeon.config import Config
 from noddingpigeon.model import load_landmarks
 
@@ -27,7 +27,9 @@ def _collect_landmarks_with_webcam(
     for label in labels:
         landmark_dict[label] = video_to_landmarks(None, max_num_frames)
         time.sleep(sleep_seconds)
-    np.savez_compressed(npz_path, **landmark_dict)
+
+    feature_names = [f.name for f in Feature]
+    np.savez_compressed(npz_path, feature_names=feature_names, **landmark_dict)
 
 
 if __name__ == "__main__":