Initial commit

.gitignore

@@ -0,0 +1,20 @@
+__pycache__/
+*.pth
+*.npy
+*.json
+*.mp4
+*.mov
+*.avi
+*.mkv
+
+.venv/
+.ipynb_checkpoints/
+.DS_Store
+
+processed_data/
+env/
+venv/
+keypoints/
+outputs/
+logs/
+train_test_paths/

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) [2021] [AI4Bharat]
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,53 @@
+# INLCUDE - Isolated Indian Sign Language Recognition
+
+This repository contains code for training models on [INCLUDE](https://zenodo.org/record/4010759) dataset
+
+# Dependencies
+
+Install the dependencies through the following command
+
+```bash
+>> pip install -r requirements.txt
+```
+
+
+
+## Steps
+- Download the INCLUDE dataset
+- Run `generate_keypoints.py` to save keypoints from Mediapipe Hands and Blazepose for train, validation and test videos. 
+```bash
+>> python generate_keypoints.py --include_dir <path to downloaded dataset> --save_dir <path to save dir> --dataset <include/include50>
+```
+- Run `runner.py` to train a machine learning model on the dataset
+```bash
+>> python runner.py --dataset <include/include50> --use_augs --model transformer --data_dir <location to saved keypoints>
+```
+- Use the `--use_pretrained` flag to either perform only inference using pretrained model or resume training with the pretrained model. 
+```bash
+>> python runner.py --dataset <include/include50> --use_augs --model transformer --data_dir <location to saved keypoints> --use_pretrained <evaluate/resume_training>
+```
+- To get predictions for videos from a pretrained model, run the following command.
+```bash
+>> python evaluate.py --data_dir <dir with videos>
+```
+
+## Citation
+
+```
+@inproceedings{10.1145/3394171.3413528,
+author = {Sridhar, Advaith and Ganesan, Rohith Gandhi and Kumar, Pratyush and Khapra, Mitesh},
+title = {INCLUDE: A Large Scale Dataset for Indian Sign Language Recognition},
+year = {2020},
+isbn = {9781450379885},
+publisher = {Association for Computing Machinery},
+doi = {10.1145/3394171.3413528},
+numpages = {10},
+series = {MM '20}
+}
+```
+
+# INCLUDE
+# INCLUDE
+# INCLUDE
+# INCLUDE
+# INCLUDE

augment.py

@@ -0,0 +1,227 @@
+import pandas as pd
+import numpy as np
+
+
+class Augmentation:
+    def __init__(self, aug_func, p=1):
+        self.aug_func = aug_func
+        self.p = p
+
+    def __call__(self, df):
+        if np.random.rand() <= self.p:
+            return self.aug_func(df)
+        return df
+
+
+def OneOf(aug_a, aug_b):
+    if np.random.rand() < 0.5:
+        return aug_a
+    return aug_b
+
+
+def plus7rotation(df):
+    # +7 degree rotation
+    df_augmented = pd.DataFrame()
+    df_augmented["uid"] = df["uid"]
+    df_augmented["pose"] = ""
+    df_augmented["hand1"] = ""
+    df_augmented["hand2"] = ""
+    df_augmented["label"] = df["label"]
+
+    theta = 7 * (np.pi / 180)
+    c, s = np.cos(theta), np.sin(theta)
+    rotation_matrix = np.array([[c, -s], [s, c]])
+
+    for i in range(df.shape[0]):
+        for col in ["pose", "hand1", "hand2"]:
+            matrix = np.array(df.loc[i, col], dtype=np.float64)
+            matrix = np.matmul(matrix, rotation_matrix)
+            matrix = np.where(np.isnan(matrix), None, matrix).tolist()
+            df_augmented.at[i, col] = matrix
+
+    return df_augmented
+
+
+def minus7rotation(df):
+    # -7 degree rotation
+    df_augmented = pd.DataFrame()
+    df_augmented["uid"] = df["uid"]
+    df_augmented["pose"] = ""
+    df_augmented["hand1"] = ""
+    df_augmented["hand2"] = ""
+    df_augmented["label"] = df["label"]
+
+    theta = -7 * (np.pi / 180)
+    c, s = np.cos(theta), np.sin(theta)
+    rotation_matrix = np.array([[c, -s], [s, c]])
+
+    for i in range(df.shape[0]):
+        for col in ["pose", "hand1", "hand2"]:
+            matrix = np.array(df.loc[i, col], dtype=np.float64)
+            matrix = np.matmul(matrix, rotation_matrix)
+            matrix = np.where(np.isnan(matrix), None, matrix).tolist()
+            df_augmented.at[i, col] = matrix
+
+    return df_augmented
+
+
+def gaussSample(df):
+    # Random Gaussian sampling
+    df_augmented = df.copy()
+    dv = 0.05 * 10 ** -2
+    sv = 0.08 * 10 ** -2
+    lv = 0.08 * 10 ** -1
+    sigma = [
+        sv,
+        dv,
+        dv,
+        dv,
+        dv,
+        dv,
+        dv,
+        sv,
+        sv,
+        sv,
+        sv,
+        lv,
+        lv,
+        lv,
+        lv,
+        sv,
+        sv,
+        sv,
+        sv,
+        sv,
+        sv,
+        sv,
+        sv,
+        lv,
+        lv,
+    ]
+
+    ## Check if keypoints is range [0, 1]
+    x_width = 1920
+    y_height = 1080
+    for i in range(df.shape[0]):
+        if np.count_nonzero(df.loc[i, "pose"]) == 0:
+            break
+
+        pose = np.array(df.loc[i, "pose"], dtype=np.float64)
+        pose[:, 0] /= x_width
+        pose[:, 1] /= y_height
+        pose_variance = np.column_stack((sigma, sigma))
+        pose = np.random.normal(pose, pose_variance)
+        pose[:, 0] *= x_width
+        pose[:, 1] *= y_height
+        pose = np.where(np.isnan(pose), None, pose).tolist()
+
+        hand1 = np.array(df.loc[i, "hand1"], dtype=np.float64)
+        hand1[:, 0] /= x_width
+        hand1[:, 1] /= y_height
+        hand1 = np.random.normal(hand1, dv)
+        hand1[:, 0] *= x_width
+        hand1[:, 1] *= y_height
+        hand1 = np.where(np.isnan(hand1), None, hand1).tolist()
+
+        hand2 = np.array(df.loc[i, "hand2"], dtype=np.float64)
+        hand2[:, 0] /= x_width
+        hand2[:, 1] /= y_height
+        hand2 = np.random.normal(hand2, dv)
+        hand2[:, 0] *= x_width
+        hand2[:, 1] *= y_height
+        hand2 = np.where(np.isnan(hand2), None, hand2).tolist()
+
+        df_augmented.at[i, "pose"] = pose
+        df_augmented.at[i, "hand1"] = hand1
+        df_augmented.at[i, "hand2"] = hand2
+
+    return df_augmented
+
+
+def cutout(df):
+    # cutout
+    df_augmented = df.copy()
+
+    pad_idx = 0
+    for i in range(df.shape[0]):
+        if np.count_nonzero(df.loc[i, "pose"]) == 0:
+            pad_idx = i
+            break
+
+    for i in range(df.shape[0]):
+        if np.count_nonzero(df.loc[i, "pose"]) == 0:
+            break
+
+        if i < pad_idx:
+            pose = np.array(df.loc[i, "pose"])
+            hand1 = np.array(df.loc[i, "hand1"])
+            hand2 = np.array(df.loc[i, "hand2"])
+            pose_zero_idx = np.random.choice(25, 3, replace=False)
+            hand1_zero_idx = np.random.choice(21, 3, replace=False)
+            hand2_zero_idx = np.random.choice(21, 3, replace=False)
+
+            for i in pose_zero_idx:
+                pose[i] = [0, 0]
+            for i in hand1_zero_idx:
+                hand1[i] = [0, 0]
+            for i in hand2_zero_idx:
+                hand2[i] = [0, 0]
+
+            pose = pose.tolist()
+            hand1 = hand1.tolist()
+            hand2 = hand2.tolist()
+
+            df_augmented.at[i, "pose"] = pose
+            df_augmented.at[i, "hand1"] = hand1
+            df_augmented.at[i, "hand2"] = hand2
+
+    return df_augmented
+
+
+def downsample(df):
+    # downsample
+    frame_len = df.shape[0]
+    if frame_len < 15:
+        return df.copy()
+
+    df_augmented = df.copy()
+    drop_idx = np.random.choice(frame_len, 15)  # 154 frames , 15 frames
+    df_augmented = df_augmented.drop(index=drop_idx)
+    return df_augmented
+
+
+def upsample(df):
+    # upsample
+    def get_avg(df, idx, col):
+        aug_points = (
+            (
+                np.array(df.loc[idx - 1, col], dtype=np.float64)
+                + np.array(df.loc[idx, col], dtype=np.float64)
+            )
+            / 2
+        ).tolist()
+        return np.where(np.isnan(aug_points), None, aug_points).tolist()
+
+    frame_length = df.shape[0]
+    additional_frames = frame_length // 10
+    df_augmented = pd.DataFrame(
+        index=np.arange(frame_length + additional_frames),
+        columns=["uid", "pose", "hand1", "hand2", "label"],
+    )
+    df_augmented["uid"] = df.iloc[0].loc["uid"]
+
+    j = 0
+    for i in range(df_augmented.shape[0]):
+        if i % 10 != 0 or i == 0:
+            df_augmented.at[i, "pose"] = df.loc[j, "pose"]
+            df_augmented.at[i, "hand1"] = df.loc[j, "hand1"]
+            df_augmented.at[i, "hand2"] = df.loc[j, "hand2"]
+            j += 1
+            continue
+
+        df_augmented.at[i, "pose"] = get_avg(df, j, "pose")
+        df_augmented.at[i, "hand1"] = get_avg(df, j, "hand1")
+        df_augmented.at[i, "hand2"] = get_avg(df, j, "hand2")
+
+    df_augmented["label"] = df.iloc[0].loc["label"]
+    return df_augmented

augs_transformer.log

@@ -0,0 +1,67 @@
+Epoch: 1, train loss: 4.454522191776949, train acc: 0.011764705882352941, val loss: 4.167257332801819, val acc: 0.0325
+Epoch: 2, train loss: 4.139250940435073, train acc: 0.022058823529411766, val loss: 4.0199155569076535, val acc: 0.0375
+Epoch: 3, train loss: 4.015947541068582, train acc: 0.008823529411764706, val loss: 4.006150960922241, val acc: 0.025
+Epoch: 4, train loss: 3.9926619613871854, train acc: 0.023529411764705882, val loss: 3.9874421834945677, val acc: 0.0375
+Epoch: 5, train loss: 3.9719819349401138, train acc: 0.023529411764705882, val loss: 4.061394691467285, val acc: 0.0125
+Epoch: 6, train loss: 3.946228302226347, train acc: 0.011764705882352941, val loss: 4.082728552818298, val acc: 0.025
+Epoch: 7, train loss: 3.9648723742541145, train acc: 0.020588235294117647, val loss: 4.055822134017944, val acc: 0.025
+Epoch: 8, train loss: 3.946899579553043, train acc: 0.016176470588235296, val loss: 4.08174843788147, val acc: 0.0125
+Epoch: 9, train loss: 3.9401227221769446, train acc: 0.026470588235294117, val loss: 4.105770111083984, val acc: 0.0
+Epoch: 10, train loss: 3.9309966536129224, train acc: 0.025, val loss: 4.0282625436782835, val acc: 0.0
+Epoch: 11, train loss: 3.9365149301641127, train acc: 0.01911764705882353, val loss: 4.040050387382507, val acc: 0.025
+Epoch: 12, train loss: 3.9410057825200697, train acc: 0.010294117647058823, val loss: 4.0812424421310425, val acc: 0.025
+Epoch: 13, train loss: 3.926714330561021, train acc: 0.03235294117647059, val loss: 4.062881135940552, val acc: 0.0375
+Epoch: 14, train loss: 3.909536064372343, train acc: 0.030882352941176472, val loss: 4.074292922019959, val acc: 0.025
+Epoch: 15, train loss: 3.9082603931427, train acc: 0.016176470588235296, val loss: 4.133523559570312, val acc: 0.0125
+Epoch: 16, train loss: 3.9084412041832417, train acc: 0.029411764705882353, val loss: 4.1158226251602175, val acc: 0.0125
+Epoch: 17, train loss: 3.9092662194195915, train acc: 0.013235294117647059, val loss: 4.096873998641968, val acc: 0.0125
+Epoch: 1, train loss: 3.7796568870544434, train acc: 0.05416666666666667, val loss: 3.3048714796702066, val acc: 0.041666666666666664
+Epoch: 2, train loss: 2.96195813814799, train acc: 0.07083333333333333, val loss: 2.759005864461263, val acc: 0.08333333333333333
+Epoch: 3, train loss: 2.9263431151707966, train acc: 0.05, val loss: 2.9804892539978027, val acc: 0.125
+Epoch: 4, train loss: 2.9008073091506956, train acc: 0.09583333333333334, val loss: 2.922832171122233, val acc: 0.125
+Epoch: 5, train loss: 2.833145562807719, train acc: 0.07916666666666666, val loss: 2.975987990697225, val acc: 0.125
+Epoch: 6, train loss: 2.8469337860743207, train acc: 0.05, val loss: 2.910189708073934, val acc: 0.0
+Epoch: 7, train loss: 2.795868976910909, train acc: 0.07083333333333333, val loss: 2.9394245942433677, val acc: 0.041666666666666664
+Epoch: 8, train loss: 2.8105882883071898, train acc: 0.06666666666666667, val loss: 2.85969344774882, val acc: 0.041666666666666664
+Epoch: 9, train loss: 2.773549795150757, train acc: 0.075, val loss: 2.7630320390065513, val acc: 0.08333333333333333
+Epoch: 10, train loss: 2.673665165901184, train acc: 0.11666666666666667, val loss: 2.9171629746754966, val acc: 0.041666666666666664
+Epoch: 11, train loss: 2.7379541556040445, train acc: 0.09166666666666666, val loss: 2.762899319330851, val acc: 0.125
+Epoch: 12, train loss: 2.7289541959762573, train acc: 0.08333333333333333, val loss: 2.7536725203196206, val acc: 0.08333333333333333
+Epoch: 13, train loss: 2.6597306728363037, train acc: 0.1, val loss: 2.668971618016561, val acc: 0.125
+Epoch: 14, train loss: 2.6211634238560992, train acc: 0.13333333333333333, val loss: 2.6979883511861167, val acc: 0.16666666666666666
+Epoch: 15, train loss: 2.519361448287964, train acc: 0.12083333333333333, val loss: 2.602913777033488, val acc: 0.08333333333333333
+Epoch: 16, train loss: 2.257799812157949, train acc: 0.20416666666666666, val loss: 2.008861700693766, val acc: 0.375
+Epoch: 17, train loss: 1.9873774607976278, train acc: 0.2875, val loss: 1.886928955713908, val acc: 0.4166666666666667
+Epoch: 18, train loss: 1.8596484184265136, train acc: 0.32083333333333336, val loss: 1.5878275632858276, val acc: 0.4166666666666667
+Epoch: 19, train loss: 1.7037100315093994, train acc: 0.37916666666666665, val loss: 1.804202914237976, val acc: 0.25
+Epoch: 20, train loss: 1.6452774405479431, train acc: 0.4083333333333333, val loss: 1.3881817261377971, val acc: 0.375
+Epoch: 21, train loss: 1.4983825981616974, train acc: 0.4708333333333333, val loss: 1.309925119082133, val acc: 0.5
+Epoch: 22, train loss: 1.4139545102914175, train acc: 0.42083333333333334, val loss: 1.2405953407287598, val acc: 0.5
+Epoch: 23, train loss: 1.3079413563013076, train acc: 0.5416666666666666, val loss: 1.0623323718706768, val acc: 0.7083333333333334
+Epoch: 24, train loss: 1.0987703889608382, train acc: 0.6458333333333334, val loss: 0.971619725227356, val acc: 0.7083333333333334
+Epoch: 25, train loss: 1.0566187125941118, train acc: 0.6375, val loss: 1.001390238602956, val acc: 0.6666666666666666
+Epoch: 26, train loss: 0.9702375590801239, train acc: 0.6625, val loss: 1.0340282917022705, val acc: 0.6666666666666666
+Epoch: 27, train loss: 0.8936778622368972, train acc: 0.6791666666666667, val loss: 0.8146964112917582, val acc: 0.7916666666666666
+Epoch: 28, train loss: 0.7876667633652688, train acc: 0.7291666666666666, val loss: 0.805329442024231, val acc: 0.7083333333333334
+Epoch: 29, train loss: 0.7577408134937287, train acc: 0.7458333333333333, val loss: 0.6200626691182455, val acc: 0.8333333333333334
+Epoch: 30, train loss: 0.7406811500589053, train acc: 0.7625, val loss: 0.756562332312266, val acc: 0.7083333333333334
+Epoch: 31, train loss: 0.6476593340436617, train acc: 0.8, val loss: 0.41417930523554486, val acc: 0.9166666666666666
+Epoch: 32, train loss: 0.6368213320771853, train acc: 0.7833333333333333, val loss: 0.7533869743347168, val acc: 0.7083333333333334
+Epoch: 33, train loss: 0.8467342297236125, train acc: 0.7, val loss: 0.46440934141476947, val acc: 0.875
+Epoch: 34, train loss: 0.6197561246653398, train acc: 0.7666666666666667, val loss: 0.3979839086532593, val acc: 0.8333333333333334
+Epoch: 35, train loss: 0.6288725058237712, train acc: 0.7625, val loss: 0.35973605265220004, val acc: 0.8333333333333334
+Epoch: 36, train loss: 0.6005093584458033, train acc: 0.7958333333333333, val loss: 0.4219923714796702, val acc: 0.7916666666666666
+Epoch: 37, train loss: 0.5010966459910074, train acc: 0.8375, val loss: 0.3479422728220622, val acc: 0.875
+Epoch: 38, train loss: 0.5946198473374049, train acc: 0.7791666666666667, val loss: 0.3922339826822281, val acc: 0.875
+Epoch: 39, train loss: 0.4414273346463839, train acc: 0.8666666666666667, val loss: 0.25101877748966217, val acc: 0.9583333333333334
+Epoch: 40, train loss: 0.42124902034799255, train acc: 0.8541666666666666, val loss: 0.27305885155995685, val acc: 0.9166666666666666
+Epoch: 41, train loss: 0.39148166303833326, train acc: 0.8666666666666667, val loss: 0.3132087141275406, val acc: 0.875
+Epoch: 42, train loss: 0.44249625156323114, train acc: 0.8416666666666667, val loss: 0.3087519307931264, val acc: 0.875
+Epoch: 43, train loss: 0.399545910085241, train acc: 0.8833333333333333, val loss: 0.28369801739851636, val acc: 0.9166666666666666
+Epoch: 44, train loss: 0.40615768507122996, train acc: 0.8583333333333333, val loss: 0.3653286099433899, val acc: 0.8333333333333334
+Epoch: 45, train loss: 0.45086419557531676, train acc: 0.825, val loss: 0.3478381074965, val acc: 0.9166666666666666
+Epoch: 46, train loss: 0.36654019008080163, train acc: 0.8583333333333333, val loss: 0.30720309416453045, val acc: 0.9166666666666666
+Epoch: 47, train loss: 0.32245828956365585, train acc: 0.9083333333333333, val loss: 0.17537023747960725, val acc: 0.875
+Epoch: 48, train loss: 0.45214990166326363, train acc: 0.8166666666666667, val loss: 0.19804753363132477, val acc: 0.9583333333333334
+Epoch: 49, train loss: 0.34067513768871627, train acc: 0.8666666666666667, val loss: 0.2660779853661855, val acc: 0.9166666666666666
+Epoch: 50, train loss: 0.31909968542555966, train acc: 0.875, val loss: 0.2961270858844121, val acc: 0.875

cnn_runner.py

@@ -0,0 +1,222 @@
+import os
+import json
+
+import numpy as np
+import torch
+from tqdm.auto import tqdm
+
+from models import CNN
+from configs import CnnConfig
+from utils import load_json
+import cv2
+import glob
+
+
+def replace_nan(x, y):
+    x = 0.0 if np.isnan(x) else x
+    y = 0.0 if np.isnan(y) else y
+    return x, y
+
+
+def draw_hands(
+    image,
+    hand_x,
+    hand_y,
+    connections,
+    connection_color,
+    thickness,
+    point_color,
+    frame_length,
+    frame_width,
+):
+    for connection in connections:
+        x0 = hand_x[connection[0]] * frame_width
+        y0 = hand_y[connection[0]] * frame_length
+        x1 = hand_x[connection[1]] * frame_width
+        y1 = hand_y[connection[1]] * frame_length
+
+        x0, y0 = replace_nan(x0, y0)
+        x1, y1 = replace_nan(x1, y1)
+
+        cv2.line(
+            image,
+            (int(x0), int(y0)),
+            (int(x1), int(y1)),
+            connection_color,
+            thickness,
+        )
+
+    for x, y in zip(hand_x, hand_y):
+        x, y = replace_nan(x, y)
+        cv2.circle(image, (int(x), int(y)), thickness, point_color, thickness)
+
+    return image
+
+
+def draw_pose(
+    image,
+    pose_x,
+    pose_y,
+    links,
+    connection_color,
+    thickness,
+    point_color,
+    frame_length,
+    frame_width,
+):
+    for link in links:
+        x0 = pose_x[link[0]] * frame_width
+        y0 = pose_y[link[0]] * frame_length
+        x1 = pose_x[link[1]] * frame_width
+        y1 = pose_y[link[1]] * frame_length
+
+        x0, y0 = replace_nan(x0, y0)
+        x1, y1 = replace_nan(x1, y1)
+
+        cv2.line(
+            image,
+            (int(x0), int(y0)),
+            (int(x1), int(y1)),
+            connection_color,
+            thickness,
+        )
+        cv2.circle(image, (int(x0), int(y0)), thickness, point_color, thickness)
+        cv2.circle(image, (int(x1), int(y1)), thickness, point_color, thickness)
+
+    return image
+
+
+def cnn_feat(file_path, save_dir):
+    video_record = load_json(file_path)
+
+    FRAME_LENGTH = 1080
+    FRAME_WIDTH = 1920
+    POINT_COLOR = (255, 0, 0)
+    CONNECTION_COLOR = (0, 255, 0)
+    THICKNESS = 2
+
+    connections = [
+        (0, 1),
+        (1, 2),
+        (2, 3),
+        (3, 4),
+        (5, 6),
+        (6, 7),
+        (7, 8),
+        (9, 10),
+        (10, 11),
+        (11, 12),
+        (13, 14),
+        (14, 15),
+        (15, 16),
+        (17, 18),
+        (18, 19),
+        (19, 20),
+        (0, 5),
+        (5, 9),
+        (9, 13),
+        (13, 17),
+        (0, 17),
+    ]
+
+    links = [
+        (11, 12),
+        (11, 23),
+        (12, 24),
+        (23, 24),
+        (11, 13),
+        (13, 15),
+        (12, 14),
+        (14, 16),
+        (15, 21),
+        (15, 17),
+        (17, 19),
+        (19, 15),
+        (22, 16),
+        (16, 18),
+        (18, 20),
+        (16, 20),
+    ]
+
+    model = CNN(CnnConfig)
+    features = np.empty((0, CnnConfig.output_dim))
+
+    assert video_record["n_frames"] > 0, "Number of frames should be greater than zero"
+    for i in range(video_record["n_frames"]):
+        image = np.zeros((FRAME_LENGTH, FRAME_WIDTH, 3), np.uint8)
+        pose_x = video_record["pose_x"][i]
+        pose_y = video_record["pose_y"][i]
+        hand1_x = video_record["hand1_x"][i]
+        hand1_y = video_record["hand1_y"][i]
+        hand2_x = video_record["hand2_x"][i]
+        hand2_y = video_record["hand2_y"][i]
+
+        if hand1_x[0] != 0:
+            image = draw_hands(
+                image,
+                hand1_x,
+                hand1_y,
+                connections,
+                CONNECTION_COLOR,
+                THICKNESS,
+                POINT_COLOR,
+                FRAME_LENGTH,
+                FRAME_WIDTH,
+            )
+
+        if hand2_x[0] != 0:
+            image = draw_hands(
+                image,
+                hand2_x,
+                hand2_y,
+                connections,
+                CONNECTION_COLOR,
+                THICKNESS,
+                POINT_COLOR,
+                FRAME_LENGTH,
+                FRAME_WIDTH,
+            )
+
+        image = draw_pose(
+            image,
+            pose_x,
+            pose_y,
+            links,
+            CONNECTION_COLOR,
+            THICKNESS,
+            POINT_COLOR,
+            FRAME_LENGTH,
+            FRAME_WIDTH,
+        )
+        image = image.astype(np.float32) / 255
+        image = cv2.resize(image, (224, 224))
+        feat = model(torch.FloatTensor(image).permute(2, 0, 1).unsqueeze(0))
+        features = np.vstack([features, feat.numpy()])
+
+    save_path = os.path.join(save_dir, video_record["uid"] + ".npy")
+    np.save(save_path, features)
+
+
+def runner(args, mode):
+    json_files = sorted(
+        glob.glob(
+            os.path.join(args.data_dir, f"{args.dataset}_{mode}_keypoints", "*.json")
+        )
+    )
+
+    save_dir = os.path.join(args.data_dir, f"{args.dataset}_{mode}_features")
+    if not os.path.exists(save_dir):
+        os.mkdir(save_dir)
+
+    saved_files = glob.glob(os.path.join(save_dir, "*.npy"))
+    if len(saved_files) != len(json_files):
+        for file_path in tqdm(json_files, desc=f"Saving CNN features - {mode} files"):
+            cnn_feat(file_path, save_dir)
+    else:
+        print(mode, "CNN features already exist!")
+
+
+def save_cnn_features(args):
+    runner(args, mode="train")
+    runner(args, mode="val")
+    runner(args, mode="test")

commands.txt

@@ -0,0 +1,52 @@
+#Then recreate it on any machine with:
+python -m venv venv
+source venv/bin/activate
+pip install -r requirements.txt
+
+#What to run now (recommended path)
+
+#1. Sanity-check that videos are readable (this should now show 10/10):
+venv/bin/python generate_keypoints.py \
+--include_dir /home/ravijaanthony/Documents/dev/IIT/FYP/Code/DATA \
+ --save_dir ./processed_data -\
+ -dataset include50 \
+ --scan \
+  --preflight \
+ --preflight_n 10
+
+#2. Generate keypoints using Holistic (pose + hands + face) and save eyebrow landmarks:
+rm -rf processed_data/include50_*_keypoints
+venv/bin/python generate_keypoints.py \
+    --include_dir /home/ravijaanthony/Documents/dev/IIT/FYP/Code/DATA \
+    --save_dir ./processed_data \
+    --dataset include50 \
+    --scan \
+    --use_holistic \
+    --face_mode full \
+    --splits all \
+    --split_seed 0
+
+#3. Train (same as before, but make sure you use the venv python):
+venv/bin/python runner.py \
+--dataset include50 \
+--use_augs \
+--model transformer \
+--data_dir ./processed_data \
+--batch_size 8
+
+
+# For INCLUDE dataset (263 classes)
+venv/bin/python runner.py \
+  --dataset include \
+  --model transformer \
+  --transformer_size large \
+  --data_dir ./processed_data \
+  --use_pretrained evaluate
+
+# For INCLUDE50 dataset (50 classes)
+venv/bin/python runner.py \
+  --dataset include50 \
+  --model transformer \
+  --transformer_size small \
+  --data_dir ./processed_data \
+  --use_pretrained evaluate

configs.py

@@ -0,0 +1,63 @@
+from dataclasses import dataclass, field
+import transformers
+import timm
+
+
+@dataclass
+class LstmConfig:
+    input_size: int = 134
+    hidden_size: int = 256
+    num_layers: int = 5
+    batch_first: bool = True
+    bidirectional: bool = True
+    dropout: float = 0.2
+
+
+@dataclass
+class XgbConfig:
+    booster: str = "gbtree"
+    silent: int = 0
+    max_depth: int = 2
+    subsample: float = 0.9923301318585108
+    colsample_bytree: float = 0.7747027267489391
+    reg_lambda: int = 3
+    objective: str = "multi:softprob"
+    eval_metric: str = "mlogloss"
+    tree_method: str = "gpu_hist"  ## change it to `hist` if gpu not available
+
+
+@dataclass
+class TransformerConfig:
+    size: str
+    input_size: int = 134
+    max_position_embeddings: int = field(default=256, repr=False)
+    layer_norm_eps: float = field(default=1e-12, repr=False)
+    hidden_dropout_prob: float = field(default=0.1, repr=False)
+    hidden_size: int = field(default=512, repr=False)
+    num_attention_heads: int = field(default=8, repr=False)
+    num_hidden_layers: int = field(default=4, repr=False)
+    model_config: transformers.BertConfig = field(init=False)
+
+    def __post_init__(self):
+        assert self.size in ["small", "large"]
+        if self.size == "small":
+            self.hidden_size = 256
+            self.num_attention_heads = 4
+            self.num_hidden_layers = 2
+
+        self.model_config = transformers.BertConfig(
+            hidden_size=self.hidden_size,
+            num_attention_heads=self.num_attention_heads,
+            num_hidden_layers=self.num_hidden_layers,
+            max_position_embeddings=self.max_position_embeddings,
+        )
+
+
+@dataclass
+class CnnConfig:
+    model: str = "mobilenetv2_100"
+    output_dim: int = 1280
+
+    def __post_init__(self):
+        available_models = timm.list_models(pretrained=True)
+        assert self.model in available_models

dataset.py

@@ -0,0 +1,171 @@
+import json
+import glob
+import os
+
+import torch
+from torch.utils import data
+import numpy as np
+import pandas as pd
+from augment import (
+    Augmentation,
+    OneOf,
+    plus7rotation,
+    minus7rotation,
+    gaussSample,
+    cutout,
+    upsample,
+    downsample,
+)
+
+
+class KeypointsDataset(data.Dataset):
+    def __init__(
+        self,
+        keypoints_dir,
+        use_augs,
+        label_map,
+        mode="train",
+        max_frame_len=200,
+        frame_length=1080,
+        frame_width=1920,
+    ):
+        self.files = sorted(glob.glob(os.path.join(keypoints_dir, "*.json")))
+        self.mode = mode
+        self.use_augs = use_augs
+        self.label_map = label_map
+        self.max_frame_len = max_frame_len
+        self.frame_length = frame_length
+        self.frame_width = frame_width
+        
+         # Filter out files with labels not in label_map
+        valid_files = []
+        for file_path in self.files:
+            row = pd.read_json(file_path, typ="series")
+            label = "".join([i for i in row.label if i.isalpha()]).lower()
+            if label in self.label_map:
+                valid_files.append(file_path)
+        self.files = valid_files
+
+        self.augs = [
+            Augmentation(OneOf(plus7rotation, minus7rotation), p=0.4),
+            Augmentation(gaussSample, p=0.4),
+            Augmentation(cutout, p=0.4),
+            Augmentation(OneOf(upsample, downsample), p=0.4),
+        ]
+
+    def augment(self, df):
+        for aug in self.augs:
+            df = aug(df)
+        return df
+
+    def interpolate(self, arr):
+
+        arr_x = arr[:, :, 0]
+        arr_x = pd.DataFrame(arr_x)
+        arr_x = arr_x.interpolate(method="linear", limit_direction="both").to_numpy()
+
+        arr_y = arr[:, :, 1]
+        arr_y = pd.DataFrame(arr_y)
+        arr_y = arr_y.interpolate(method="linear", limit_direction="both").to_numpy()
+
+        if np.count_nonzero(~np.isnan(arr_x)) == 0:
+            arr_x = np.zeros(arr_x.shape)
+        if np.count_nonzero(~np.isnan(arr_y)) == 0:
+            arr_y = np.zeros(arr_y.shape)
+
+        arr_x = arr_x * self.frame_width
+        arr_y = arr_y * self.frame_length
+
+        return np.stack([arr_x, arr_y], axis=-1)
+
+    def combine_xy(self, x, y):
+        x, y = np.array(x), np.array(y)
+        _, length = x.shape
+        x = x.reshape((-1, length, 1))
+        y = y.reshape((-1, length, 1))
+        return np.concatenate((x, y), -1).astype(np.float32)
+
+    def __getitem__(self, idx):
+        file_path = self.files[idx]
+        row = pd.read_json(file_path, typ="series")
+        label = row.label
+        label = "".join([i for i in label if i.isalpha()]).lower()
+
+        pose = self.combine_xy(row.pose_x, row.pose_y)
+        h1 = self.combine_xy(row.hand1_x, row.hand1_y)
+        h2 = self.combine_xy(row.hand2_x, row.hand2_y)
+
+        pose = self.interpolate(pose)
+        h1 = self.interpolate(h1)
+        h2 = self.interpolate(h2)
+
+        df = pd.DataFrame.from_dict(
+            {
+                "uid": row.uid,
+                "pose": pose.tolist(),
+                "hand1": h1.tolist(),
+                "hand2": h2.tolist(),
+                "label": label,
+            }
+        )
+        if self.mode == "train" and self.use_augs:
+            df = self.augment(df)
+
+        pose = (
+            np.array(list(map(np.array, df.pose.values)))
+            .reshape(-1, 50)
+            .astype(np.float32)
+        )
+        h1 = (
+            np.array(list(map(np.array, df.hand1.values)))
+            .reshape(-1, 42)
+            .astype(np.float32)
+        )
+        h2 = (
+            np.array(list(map(np.array, df.hand2.values)))
+            .reshape(-1, 42)
+            .astype(np.float32)
+        )
+        final_data = np.concatenate((pose, h1, h2), -1)
+        final_data = np.pad(
+            final_data,
+            ((0, self.max_frame_len - final_data.shape[0]), (0, 0)),
+            "constant",
+        )
+        return {
+            "uid": row.uid,
+            "data": torch.FloatTensor(final_data),
+            "label": self.label_map[label],
+            "lablel_string": label,
+        }
+
+    def __len__(self):
+        return len(self.files)
+
+
+class FeaturesDatset(data.Dataset):
+    def __init__(self, features_dir, label_map, mode="train", max_frame_len=200):
+        self.features_dir = features_dir
+        self.file_paths = sorted(glob.glob(os.path.join(features_dir, "*.npy")))
+        self.label_map = label_map
+        self.mode = mode
+        self.max_frame_len = max_frame_len
+
+    def __getitem__(self, i):
+        file_path = self.file_paths[i]
+        data = np.load(file_path)
+        data = np.pad(
+            data,
+            ((0, self.max_frame_len - data.shape[0]), (0, 0)),
+            "constant",
+        )
+        label = os.path.basename(file_path).split("_")[0]
+        return {
+            "uid": os.path.basename(file_path).split(".")[0],
+            "data": torch.FloatTensor(data),
+            "label": self.label_map[label],
+            "lablel_string": label,
+        }
+
+    def __len__(self):
+        return len(self.file_paths)

evaluate.py

@@ -0,0 +1,169 @@
+import os
+import glob
+import argparse
+import json
+import numpy as np
+import pandas as pd
+from tqdm import tqdm
+import torch
+import torch.nn as nn
+from torch.utils import data
+from generate_keypoints import process_video
+from models import Transformer
+from configs import TransformerConfig
+from utils import load_json, load_label_map
+import shutil
+
+parser = argparse.ArgumentParser(description="Evaluate function")
+parser.add_argument("--data_dir", required=True, help="data directory")
+args = parser.parse_args()
+
+
+class KeypointsDataset(data.Dataset):
+    def __init__(
+        self,
+        keypoints_dir,
+        max_frame_len=200,
+        frame_length=1080,
+        frame_width=1920,
+    ):
+        self.files = sorted(glob.glob(os.path.join(keypoints_dir, "*.json")))
+        self.max_frame_len = max_frame_len
+        self.frame_length = frame_length
+        self.frame_width = frame_width
+
+    def interpolate(self, arr):
+
+        arr_x = arr[:, :, 0]
+        arr_x = pd.DataFrame(arr_x)
+        arr_x = arr_x.interpolate(method="linear", limit_direction="both").to_numpy()
+
+        arr_y = arr[:, :, 1]
+        arr_y = pd.DataFrame(arr_y)
+        arr_y = arr_y.interpolate(method="linear", limit_direction="both").to_numpy()
+
+        if np.count_nonzero(~np.isnan(arr_x)) == 0:
+            arr_x = np.zeros(arr_x.shape)
+        if np.count_nonzero(~np.isnan(arr_y)) == 0:
+            arr_y = np.zeros(arr_y.shape)
+
+        arr_x = arr_x * self.frame_width
+        arr_y = arr_y * self.frame_length
+
+        return np.stack([arr_x, arr_y], axis=-1)
+
+    def combine_xy(self, x, y):
+        x, y = np.array(x), np.array(y)
+        _, length = x.shape
+        x = x.reshape((-1, length, 1))
+        y = y.reshape((-1, length, 1))
+        return np.concatenate((x, y), -1).astype(np.float32)
+
+    def __getitem__(self, idx):
+        file_path = self.files[idx]
+        row = pd.read_json(file_path, typ="series")
+        label = row.label
+        label = "".join([i for i in label if i.isalpha()]).lower()
+
+        pose = self.combine_xy(row.pose_x, row.pose_y)
+        h1 = self.combine_xy(row.hand1_x, row.hand1_y)
+        h2 = self.combine_xy(row.hand2_x, row.hand2_y)
+
+        pose = self.interpolate(pose)
+        h1 = self.interpolate(h1)
+        h2 = self.interpolate(h2)
+
+        df = pd.DataFrame.from_dict(
+            {
+                "uid": row.uid,
+                "pose": pose.tolist(),
+                "hand1": h1.tolist(),
+                "hand2": h2.tolist(),
+                "label": label,
+            }
+        )
+
+        pose = (
+            np.array(list(map(np.array, df.pose.values)))
+            .reshape(-1, 50)
+            .astype(np.float32)
+        )
+        h1 = (
+            np.array(list(map(np.array, df.hand1.values)))
+            .reshape(-1, 42)
+            .astype(np.float32)
+        )
+        h2 = (
+            np.array(list(map(np.array, df.hand2.values)))
+            .reshape(-1, 42)
+            .astype(np.float32)
+        )
+        final_data = np.concatenate((pose, h1, h2), -1)
+        final_data = np.pad(
+            final_data,
+            ((0, self.max_frame_len - final_data.shape[0]), (0, 0)),
+            "constant",
+        )
+        return {
+            "uid": row.uid,
+            "data": torch.FloatTensor(final_data),
+        }
+
+    def __len__(self):
+        return len(self.files)
+
+
+@torch.no_grad()
+def inference(dataloader, model, device, label_map):
+    model.eval()
+    predictions = []
+
+    for batch in tqdm(dataloader, desc="Eval"):
+        input_data = batch["data"].to(device)
+        output = model(input_data).detach().cpu()
+        output = torch.argmax(torch.softmax(output, dim=-1), dim=-1).numpy()
+        predictions.append({"uid": batch["uid"][0], "predicted_label": label_map[output[0]]})
+
+    return predictions
+
+
+video_paths = glob.glob(os.path.join(args.data_dir, "*"))
+save_dir = "keypoints_dir"
+if os.path.isdir(save_dir):
+    shutil.rmtree(save_dir)
+os.mkdir(save_dir)
+for path in tqdm(video_paths, desc="Processing Videos"):
+    process_video(path, save_dir)
+
+label_map = load_label_map("include")
+dataset = KeypointsDataset(
+    keypoints_dir=save_dir,
+    max_frame_len=169,
+)
+
+dataloader = data.DataLoader(
+    dataset,
+    batch_size=1,
+    shuffle=False,
+    num_workers=4,
+    pin_memory=True,
+)
+label_map = dict(zip(label_map.values(), label_map.keys()))
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+config = TransformerConfig(size="large", max_position_embeddings=256)
+model = Transformer(config=config, n_classes=263)
+model = model.to(device)
+
+pretrained_model_name = "include_no_cnn_transformer_large.pth"
+pretrained_model_links = load_json("pretrained_links.json")
+if not os.path.isfile(pretrained_model_name):
+    link = pretrained_model_links[pretrained_model_name]
+    torch.hub.download_url_to_file(link, pretrained_model_name, progress=True)
+
+ckpt = torch.load(pretrained_model_name)
+model.load_state_dict(ckpt["model"])
+print("### Model loaded ###")
+
+preds = inference(dataloader, model, device, label_map)
+print(json.dumps(preds, indent=2))

generate_keypoints.py

@@ -0,0 +1,693 @@
+import os
+import json
+import glob
+import multiprocessing
+import argparse
+import os.path
+import cv2
+import mediapipe as mp
+from tqdm.auto import tqdm
+from joblib import Parallel, delayed
+import numpy as np
+import gc
+import warnings
+
+
+# MediaPipe FaceMesh landmark indices (subset) for eyebrows.
+# These are commonly used indices for left/right eyebrows.
+_EYEBROW_IDXS = [
+    # left eyebrow
+    70,
+    63,
+    105,
+    66,
+    107,
+    55,
+    65,
+    52,
+    53,
+    46,
+    # right eyebrow
+    336,
+    296,
+    334,
+    293,
+    300,
+    276,
+    283,
+    282,
+    295,
+    285,
+]
+
+
+def _label_from_path(path: str) -> str:
+    # Assumes dataset layout: <include_dir>/<label>/<video_file>
+    label = path.split("/")[-2]
+    return "".join([i for i in label if i.isalpha()]).lower()
+
+
+def _uid_from_path(path: str, label: str) -> str:
+    uid = os.path.splitext(os.path.basename(path))[0]
+    return "_".join([label, uid])
+
+
+def _try_open_video(path: str) -> cv2.VideoCapture:
+    cap = cv2.VideoCapture(path)
+    if cap.isOpened():
+        return cap
+
+    # Try forcing FFmpeg backend if available.
+    if hasattr(cv2, "CAP_FFMPEG"):
+        cap2 = cv2.VideoCapture(path, cv2.CAP_FFMPEG)
+        if cap2.isOpened():
+            return cap2
+
+    return cap
+
+
+def _pose25_from_mp_pose(pose_landmarks):
+    if pose_landmarks is None:
+        return [np.nan] * 25, [np.nan] * 25
+
+    lm = pose_landmarks.landmark
+
+    def get_xy(i: int):
+        return lm[i].x, lm[i].y
+
+    def avg_xy(a: int, b: int):
+        ax, ay = get_xy(a)
+        bx, by = get_xy(b)
+        return (ax + bx) / 2.0, (ay + by) / 2.0
+
+    # Map MediaPipe Pose (33) -> OpenPose BODY_25-like layout used by this repo (25).
+    # 0 Nose
+    nose = get_xy(0)
+    # 1 Neck (approx: midpoint of shoulders)
+    neck = avg_xy(11, 12)
+
+    rshoulder = get_xy(12)
+    relbow = get_xy(14)
+    rwrist = get_xy(16)
+
+    lshoulder = get_xy(11)
+    lelbow = get_xy(13)
+    lwrist = get_xy(15)
+
+    midhip = avg_xy(23, 24)
+
+    rhip = get_xy(24)
+    rknee = get_xy(26)
+    rankle = get_xy(28)
+
+    lhip = get_xy(23)
+    lknee = get_xy(25)
+    lankle = get_xy(27)
+
+    # Eyes / ears
+    reye = get_xy(5)
+    leye = get_xy(2)
+    rear = get_xy(8)
+    lear = get_xy(7)
+
+    # Feet (MediaPipe doesn't provide small-toe separately; we duplicate foot_index)
+    lbigtoe = get_xy(31)
+    lsmalltoe = get_xy(31)
+    lheel = get_xy(29)
+
+    rbigtoe = get_xy(32)
+    rsmalltoe = get_xy(32)
+    rheel = get_xy(30)
+
+    points = [
+        nose,
+        neck,
+        rshoulder,
+        relbow,
+        rwrist,
+        lshoulder,
+        lelbow,
+        lwrist,
+        midhip,
+        rhip,
+        rknee,
+        rankle,
+        lhip,
+        lknee,
+        lankle,
+        reye,
+        leye,
+        rear,
+        lear,
+        lbigtoe,
+        lsmalltoe,
+        lheel,
+        rbigtoe,
+        rsmalltoe,
+        rheel,
+    ]
+
+    xs = [p[0] for p in points]
+    ys = [p[1] for p in points]
+    return xs, ys
+
+
+def _hand21_from_landmarks(hand_landmarks):
+    if hand_landmarks is None:
+        return [np.nan] * 21, [np.nan] * 21
+    xs, ys = [], []
+    for landmark in hand_landmarks.landmark:
+        xs.append(landmark.x)
+        ys.append(landmark.y)
+    return xs, ys
+
+
+def _face_from_landmarks(face_landmarks, face_mode: str):
+    if face_mode == "none":
+        return None, None
+
+    if face_mode == "full":
+        idxs = None
+        n = 468
+    elif face_mode == "eyebrows":
+        idxs = _EYEBROW_IDXS
+        n = len(_EYEBROW_IDXS)
+    else:
+        raise ValueError(f"Unsupported face_mode: {face_mode}")
+
+    if face_landmarks is None:
+        return [np.nan] * n, [np.nan] * n
+
+    lm = face_landmarks.landmark
+    if idxs is None:
+        xs = [p.x for p in lm]
+        ys = [p.y for p in lm]
+        return xs, ys
+
+    xs, ys = [], []
+    for i in idxs:
+        xs.append(lm[i].x)
+        ys.append(lm[i].y)
+    return xs, ys
+
+
+def process_video(
+    path: str,
+    save_dir: str,
+    *,
+    use_holistic: bool = False,
+    face_mode: str = "none",
+    write_placeholders: bool = False,
+):
+    if not os.path.isfile(path):
+        warnings.warn(path + " file not found")
+        return
+
+    label = _label_from_path(path)
+    uid = _uid_from_path(path, label)
+
+    cap = _try_open_video(path)
+    if not cap.isOpened():
+        warnings.warn(f"OpenCV could not open video: {path}")
+        return
+
+    pose_points_x, pose_points_y = [], []
+    hand1_points_x, hand1_points_y = [], []
+    hand2_points_x, hand2_points_y = [], []
+    face_points_x, face_points_y = [], []
+
+    n_frames = 0
+
+    if use_holistic:
+        holistic = mp.solutions.holistic.Holistic(
+            static_image_mode=False,
+            model_complexity=1,
+            smooth_landmarks=True,
+            enable_segmentation=False,
+            refine_face_landmarks=True,
+            min_detection_confidence=0.5,
+            min_tracking_confidence=0.5,
+        )
+        try:
+            while cap.isOpened():
+                ret, image = cap.read()
+                if not ret:
+                    break
+
+                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+                results = holistic.process(image)
+
+                pose_x, pose_y = _pose25_from_mp_pose(results.pose_landmarks)
+                # Use holistic's left/right hands (subject perspective)
+                h1x, h1y = _hand21_from_landmarks(results.left_hand_landmarks)
+                h2x, h2y = _hand21_from_landmarks(results.right_hand_landmarks)
+
+                fx, fy = _face_from_landmarks(results.face_landmarks, face_mode)
+
+                pose_points_x.append(pose_x)
+                pose_points_y.append(pose_y)
+                hand1_points_x.append(h1x)
+                hand1_points_y.append(h1y)
+                hand2_points_x.append(h2x)
+                hand2_points_y.append(h2y)
+
+                if face_mode != "none":
+                    face_points_x.append(fx)
+                    face_points_y.append(fy)
+
+                n_frames += 1
+        finally:
+            holistic.close()
+    else:
+        # Legacy path: Pose + Hands separately (kept for backward compatibility).
+        hands = mp.solutions.hands.Hands(
+            min_detection_confidence=0.5, min_tracking_confidence=0.5
+        )
+        pose = mp.solutions.pose.Pose(
+            static_image_mode=False,
+            model_complexity=1,
+            enable_segmentation=False,
+            min_detection_confidence=0.5,
+            min_tracking_confidence=0.5,
+        )
+        try:
+            while cap.isOpened():
+                ret, image = cap.read()
+                if not ret:
+                    break
+                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+
+                hand_results = hands.process(image)
+                pose_results = pose.process(image)
+
+                # Hands
+                hand1 = (
+                    hand_results.multi_hand_landmarks[0]
+                    if hand_results.multi_hand_landmarks
+                    else None
+                )
+                hand2 = (
+                    hand_results.multi_hand_landmarks[1]
+                    if hand_results.multi_hand_landmarks
+                    and len(hand_results.multi_hand_landmarks) > 1
+                    else None
+                )
+                h1x, h1y = _hand21_from_landmarks(hand1)
+                h2x, h2y = _hand21_from_landmarks(hand2)
+
+                # Pose (mapped to 25)
+                pose_x, pose_y = _pose25_from_mp_pose(pose_results.pose_landmarks)
+
+                pose_points_x.append(pose_x)
+                pose_points_y.append(pose_y)
+                hand1_points_x.append(h1x)
+                hand1_points_y.append(h1y)
+                hand2_points_x.append(h2x)
+                hand2_points_y.append(h2y)
+
+                n_frames += 1
+        finally:
+            hands.close()
+            pose.close()
+
+    cap.release()
+
+    if n_frames == 0 and not write_placeholders:
+        warnings.warn(f"No frames decoded (n_frames=0). Skipping: {path}")
+        return
+
+    # If no frames decoded but placeholders requested, mimic previous behavior.
+    if n_frames == 0 and write_placeholders:
+        pose_points_x = [[np.nan] * 25]
+        pose_points_y = [[np.nan] * 25]
+        hand1_points_x = [[np.nan] * 21]
+        hand1_points_y = [[np.nan] * 21]
+        hand2_points_x = [[np.nan] * 21]
+        hand2_points_y = [[np.nan] * 21]
+        if face_mode != "none":
+            n_face = 468 if face_mode == "full" else len(_EYEBROW_IDXS)
+            face_points_x = [[np.nan] * n_face]
+            face_points_y = [[np.nan] * n_face]
+
+    save_data = {
+        "uid": uid,
+        "label": label,
+        "pose_x": pose_points_x,
+        "pose_y": pose_points_y,
+        "hand1_x": hand1_points_x,
+        "hand1_y": hand1_points_y,
+        "hand2_x": hand2_points_x,
+        "hand2_y": hand2_points_y,
+        "n_frames": n_frames,
+    }
+    if face_mode != "none":
+        save_data["face_x"] = face_points_x
+        save_data["face_y"] = face_points_y
+
+    with open(os.path.join(save_dir, f"{uid}.json"), "w") as f:
+        json.dump(save_data, f)
+
+    del save_data
+    gc.collect()
+
+
+def _build_label_dir_map(include_dir: str) -> dict[str, str]:
+    """Maps lowercased directory name -> actual directory name."""
+    mapping: dict[str, str] = {}
+    try:
+        for name in os.listdir(include_dir):
+            full = os.path.join(include_dir, name)
+            if os.path.isdir(full):
+                mapping[name.lower()] = name
+                cleaned = ''.join([c for c in name if c.isalpha()]).lower()
+                if cleaned and cleaned not in mapping:
+                    mapping[cleaned] = name
+    except FileNotFoundError:
+        return mapping
+    return mapping
+
+
+def _resolve_video_path(include_dir: str, rel_path: str, label_dir_map: dict[str, str]) -> str:
+    """Resolves train_test_paths entries against different dataset folder layouts.
+
+    Supports:
+    1) Official INCLUDE layout: <include_dir>/<Category>/<N. Label>/<file>
+    2) Flat label layout:      <include_dir>/<Label>/<file>
+    """
+    direct = os.path.join(include_dir, rel_path)
+    if os.path.isfile(direct):
+        return direct
+
+    # Try <include_dir>/<label>/<filename>
+    parts = [p for p in rel_path.split('/') if p]
+    filename = os.path.basename(rel_path)
+    label = None
+    if len(parts) >= 2:
+        label_part = parts[-2]
+        label = ''.join([c for c in label_part if c.isalpha()]).lower()
+
+    if label:
+        dir_name = label_dir_map.get(label)
+        if dir_name:
+            cand = os.path.join(include_dir, dir_name, filename)
+            if os.path.isfile(cand):
+                return cand
+
+        # Sometimes folders are title-cased (e.g., Court)
+        for variant in (label.capitalize(), label.title()):
+            dir_name = label_dir_map.get(variant.lower())
+            if dir_name:
+                cand = os.path.join(include_dir, dir_name, filename)
+                if os.path.isfile(cand):
+                    return cand
+
+    return direct
+
+
+def load_file(path, include_dir):
+    with open(path, "r") as fp:
+        rel_paths = [ln.strip() for ln in fp.read().splitlines() if ln.strip()]
+
+    label_dir_map = _build_label_dir_map(include_dir)
+    return [_resolve_video_path(include_dir, rel, label_dir_map) for rel in rel_paths]
+
+
+def scan_include_dir_videos(include_dir: str) -> list[str]:
+    """Scans include_dir for videos in a flat label-folder layout.
+
+    Expected layout:
+      <include_dir>/<Label>/*.MOV (or mp4/avi/mkv)
+    """
+    exts = ("*.MOV", "*.mov", "*.MP4", "*.mp4", "*.AVI", "*.avi", "*.MKV", "*.mkv")
+    videos: list[str] = []
+    for ext in exts:
+        videos.extend(glob.glob(os.path.join(include_dir, "*", ext)))
+    # Deduplicate + stable order
+    return sorted(set(videos))
+
+
+def split_paths(paths: list[str], seed: int = 0, train_ratio: float = 0.8, val_ratio: float = 0.1):
+    if not paths:
+        return [], [], []
+    rng = np.random.default_rng(seed)
+    idx = np.arange(len(paths))
+    rng.shuffle(idx)
+    paths = [paths[i] for i in idx]
+
+    n = len(paths)
+    n_train = int(train_ratio * n)
+    n_val = int(val_ratio * n)
+
+    train = paths[:n_train]
+    val = paths[n_train : n_train + n_val]
+    test = paths[n_train + n_val :]
+    return train, val, test
+
+
+def load_train_test_val_paths(include_dir: str, dataset: str):
+    train_paths = load_file(f"train_test_paths/{dataset}_train.txt", include_dir)
+    val_paths = load_file(f"train_test_paths/{dataset}_val.txt", include_dir)
+    test_paths = load_file(f"train_test_paths/{dataset}_test.txt", include_dir)
+    return train_paths, val_paths, test_paths
+
+
+def save_keypoints(
+    *,
+    dataset: str,
+    file_paths: list[str],
+    mode: str,
+    save_root: str,
+    n_jobs: int,
+    use_holistic: bool,
+    face_mode: str,
+    limit: int,
+    no_parallel: bool,
+    write_placeholders: bool,
+):
+    save_dir = os.path.join(save_root, f"{dataset}_{mode}_keypoints")
+    os.makedirs(save_dir, exist_ok=True)
+
+    if limit and limit > 0:
+        file_paths = file_paths[:limit]
+
+    existing = [p for p in file_paths if os.path.isfile(p)]
+    missing = len(file_paths) - len(existing)
+    if missing:
+        warnings.warn(f"{missing} {mode} videos missing on disk (skipping them)")
+
+    if no_parallel or n_jobs == 1:
+        for path in tqdm(existing, desc=f"processing {mode} videos"):
+            process_video(
+                path,
+                save_dir,
+                use_holistic=use_holistic,
+                face_mode=face_mode,
+                write_placeholders=write_placeholders,
+            )
+        return
+
+    Parallel(n_jobs=n_jobs, backend="multiprocessing")(
+        delayed(process_video)(
+            path,
+            save_dir,
+            use_holistic=use_holistic,
+            face_mode=face_mode,
+            write_placeholders=write_placeholders,
+        )
+        for path in tqdm(existing, desc=f"processing {mode} videos")
+    )
+
+
+def preflight_videos(video_paths: list[str], n: int = 10):
+    """Quickly checks existence + OpenCV decode for a small sample."""
+    sample = video_paths[: max(0, n)]
+    exists = 0
+    opened = 0
+    first_read_ok = 0
+
+    for path in sample:
+        if not os.path.isfile(path):
+            continue
+        exists += 1
+        cap = _try_open_video(path)
+        if not cap.isOpened():
+            cap.release()
+            continue
+        opened += 1
+        ok, _frame = cap.read()
+        if ok:
+            first_read_ok += 1
+        cap.release()
+
+    print('Preflight results:')
+    print(f'  sample size: {len(sample)}')
+    print(f'  exists on disk: {exists}')
+    print(f'  OpenCV opened: {opened}')
+    print(f'  first frame decoded: {first_read_ok}')
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Generate keypoints from MediaPipe")
+    parser.add_argument(
+        "--include_dir",
+        default="",
+        type=str,
+        required=True,
+        help="path to the location of INCLUDE/INCLUDE50 videos",
+    )
+    parser.add_argument(
+        "--save_dir",
+        default="",
+        type=str,
+        required=True,
+        help="location to output json files",
+    )
+    parser.add_argument(
+        "--dataset", default="include", type=str, help="options: include or include50"
+    )
+    parser.add_argument(
+        "--video",
+        default=None,
+        type=str,
+        help="process a single video path (bypasses train_test_paths)",
+    )
+    parser.add_argument(
+        "--use_holistic",
+        action="store_true",
+        help="use MediaPipe Holistic (pose+hands+face) instead of separate pose/hands",
+    )
+    parser.add_argument(
+        "--face_mode",
+        default="none",
+        choices=["none", "eyebrows", "full"],
+        help="face landmarks to save when using holistic",
+    )
+    parser.add_argument(
+        "--splits",
+        default="all",
+        choices=["train", "val", "test", "all"],
+        help="which splits to process",
+    )
+    parser.add_argument(
+        "--limit",
+        default=0,
+        type=int,
+        help="process only the first N videos per split (useful for testing)",
+    )
+    parser.add_argument(
+        "--jobs",
+        default=multiprocessing.cpu_count(),
+        type=int,
+        help="number of parallel workers",
+    )
+    parser.add_argument(
+        "--no_parallel",
+        action="store_true",
+        help="disable multiprocessing (easier debugging)",
+    )
+    parser.add_argument(
+        "--write_placeholders",
+        action="store_true",
+        help="write placeholder NaN JSONs even when videos can't be decoded",
+    )
+
+    parser.add_argument(
+        "--preflight",
+        action="store_true",
+        help="only check that videos exist and can be decoded, then exit",
+    )
+    parser.add_argument(
+        "--preflight_n",
+        default=10,
+        type=int,
+        help="number of videos to check in preflight",
+    )
+
+    parser.add_argument(
+        "--scan",
+        action="store_true",
+        help="ignore train_test_paths and scan include_dir/<label>/* for videos, then make a random train/val/test split",
+    )
+    parser.add_argument(
+        "--split_seed",
+        default=0,
+        type=int,
+        help="seed used when --scan creates train/val/test splits",
+    )
+
+    args = parser.parse_args()
+
+    if args.video is not None:
+        single_dir = os.path.join(args.save_dir, f"{args.dataset}_single_keypoints")
+        os.makedirs(single_dir, exist_ok=True)
+        process_video(
+            args.video,
+            single_dir,
+            use_holistic=args.use_holistic,
+            face_mode=args.face_mode,
+            write_placeholders=args.write_placeholders,
+        )
+        print(f"Saved single-video keypoints to: {single_dir}")
+        raise SystemExit(0)
+
+    if args.scan:
+        all_paths = scan_include_dir_videos(args.include_dir)
+        train_paths, val_paths, test_paths = split_paths(all_paths, seed=args.split_seed)
+    else:
+        train_paths, val_paths, test_paths = load_train_test_val_paths(
+            args.include_dir, args.dataset
+        )
+
+    if args.preflight:
+        paths = []
+        if args.splits in ("train", "all"):
+            paths.extend(train_paths)
+        if args.splits in ("val", "all"):
+            paths.extend(val_paths)
+        if args.splits in ("test", "all"):
+            paths.extend(test_paths)
+        preflight_videos(paths, n=args.preflight_n)
+        raise SystemExit(0)
+
+
+    if args.splits in ("val", "all"):
+        save_keypoints(
+            dataset=args.dataset,
+            file_paths=val_paths,
+            mode="val",
+            save_root=args.save_dir,
+            n_jobs=args.jobs,
+            use_holistic=args.use_holistic,
+            face_mode=args.face_mode,
+            limit=args.limit,
+            no_parallel=args.no_parallel,
+            write_placeholders=args.write_placeholders,
+        )
+
+    if args.splits in ("test", "all"):
+        save_keypoints(
+            dataset=args.dataset,
+            file_paths=test_paths,
+            mode="test",
+            save_root=args.save_dir,
+            n_jobs=args.jobs,
+            use_holistic=args.use_holistic,
+            face_mode=args.face_mode,
+            limit=args.limit,
+            no_parallel=args.no_parallel,
+            write_placeholders=args.write_placeholders,
+        )
+
+    if args.splits in ("train", "all"):
+        save_keypoints(
+            dataset=args.dataset,
+            file_paths=train_paths,
+            mode="train",
+            save_root=args.save_dir,
+            n_jobs=args.jobs,
+            use_holistic=args.use_holistic,
+            face_mode=args.face_mode,
+            limit=args.limit,
+            no_parallel=args.no_parallel,
+            write_placeholders=args.write_placeholders,
+        )

models/__init__.py

@@ -0,0 +1,4 @@
+from .lstm import LSTM
+from .transformer import Transformer
+from .xgboost import Xgboost
+from .cnn import CNN

models/cnn.py

@@ -0,0 +1,11 @@
+import torch.nn as nn
+import timm
+
+
+class CNN(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.model = timm.create_model(config.model, pretrained=True, num_classes=0)
+
+    def forward(self, x):
+        return self.model(x).detach()

models/lstm.py

@@ -0,0 +1,23 @@
+from dataclasses import asdict
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class LSTM(nn.Module):
+    def __init__(self, config, n_classes=50):
+        super().__init__()
+        config_dict = asdict(config)
+        self.lstm = nn.LSTM(**config_dict)
+        in_features = (
+            config.hidden_size * 2 if config.bidirectional else config.hidden_size
+        )
+        self.l1 = nn.Linear(in_features=in_features, out_features=n_classes)
+
+    def forward(self, x):
+        x, (_, _) = self.lstm(x)
+        x = torch.max(x, dim=1).values
+        x = F.dropout(x, p=0.3)
+        x = self.l1(x)
+        return x

models/transformer.py

@@ -0,0 +1,61 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import transformers
+
+
+class PositionEmbedding(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.position_embeddings = nn.Embedding(
+            config.max_position_embeddings, config.hidden_size
+        )
+        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.register_buffer(
+            "position_ids", torch.arange(config.max_position_embeddings).expand((1, -1))
+        )
+        self.position_embedding_type = getattr(
+            config, "position_embedding_type", "absolute"
+        )
+
+    def forward(self, x):
+        input_shape = x.size()
+        seq_length = input_shape[1]
+        position_ids = self.position_ids[:, :seq_length]
+
+        position_embeddings = self.position_embeddings(position_ids)
+        embeddings = x + position_embeddings
+        embeddings = self.LayerNorm(embeddings)
+        embeddings = self.dropout(embeddings)
+        return embeddings
+
+
+class Transformer(nn.Module):
+    def __init__(self, config, n_classes=50):
+        super().__init__()
+        self.l1 = nn.Linear(
+            in_features=config.input_size, out_features=config.hidden_size
+        )
+        self.embedding = PositionEmbedding(config)
+
+        setattr(config.model_config, "_attn_implementation", "eager")
+        
+        self.layers = nn.ModuleList(
+            [
+                transformers.BertLayer(config.model_config)
+                for _ in range(config.num_hidden_layers)
+            ]
+        )
+        self.l2 = nn.Linear(in_features=config.hidden_size, out_features=n_classes)
+
+    def forward(self, x):
+        x = self.l1(x)
+        x = self.embedding(x)
+        for layer in self.layers:
+            x = layer(x)[0]
+
+        x = torch.max(x, dim=1).values
+        x = F.dropout(x, p=0.2)
+        x = self.l2(x)
+        return x

models/xgboost.py

@@ -0,0 +1,27 @@
+from dataclasses import asdict
+import pickle
+
+import xgboost
+
+
+class Xgboost:
+    def __init__(self, config):
+        config = asdict(config)
+        self.model = xgboost.XGBClassifier(use_label_encoder=False, **config)
+
+    def __call__(self, x):
+        return self.model.predict(x)
+
+    def fit(self, x_train, y_train, x_val, y_val):
+        self.model.fit(
+            x_train,
+            y_train,
+            eval_set=[(x_train, y_train), (x_val, y_val)],
+            verbose=True,
+        )
+
+    def save(self, save_path):
+        pickle.dump(self.model, open(save_path, "wb"))
+
+    def load(self, load_path):
+        self.model = pickle.load(open(load_path, "rb"))

requirements.txt

@@ -0,0 +1,7 @@
+torch
+transformers
+timm
+mediapipe
+tqdm
+scikit-learn
+xgboost ## installs CPU version, check docs for GPU compatible version

runner.py

@@ -0,0 +1,88 @@
+import warnings
+import argparse
+
+import train_nn
+import train_xgb
+from cnn_runner import save_cnn_features
+
+parser = argparse.ArgumentParser(
+    description="INCLUDE trainer for xgboost, lstm and transformer"
+)
+parser.add_argument("--seed", default=0, type=int, help="seed value")
+parser.add_argument(
+    "--dataset", default="include", type=str, help="options: include or include50"
+)
+parser.add_argument(
+    "--use_augs",
+    action="store_true",
+    help="use augmented data",
+)
+parser.add_argument(
+    "--use_cnn",
+    action="store_true",
+    help="use mobilenet to convert keypoints to videos and generate embeddings from CNN",
+)
+parser.add_argument(
+    "--model",
+    default="lstm",
+    type=str,
+    help="options: lstm, transformer, xgboost",
+)
+parser.add_argument(
+    "--data_dir",
+    default="",
+    type=str,
+    required=True,
+    help="location to train, val and test json files",
+)
+parser.add_argument(
+    "--save_path",
+    default="./",
+    type=str,
+    help="location to save trained model",
+)
+parser.add_argument(
+    "--epochs", default=50, type=int, help="number of epochs to train the model"
+)
+parser.add_argument("--batch_size", default=128, type=int, help="batch size of data")
+parser.add_argument(
+    "--learning_rate",
+    default=1e-4,
+    type=float,
+    help="learning rate for training neural net",
+)
+parser.add_argument(
+    "--transformer_size", default="small", type=str, help="options: small, large"
+)
+parser.add_argument(
+    "--use_pretrained",
+    default=None,
+    help="use pretrained model. options: evaluate, resume_training",
+)
+args = parser.parse_args()
+
+
+if __name__ == "__main__":
+
+    if args.model == "xgboost":
+        if args.use_pretrained:
+            raise Exception("Pre-trained models are not available for XGBoost")
+        if args.use_cnn:
+            warnings.warn(
+                "use_cnn flag set to true for xgboost model. xgboost will not use cnn features"
+            )
+        train_xgb.fit(args)
+        train_xgb.evaluate(args)
+
+    else:
+        if args.use_cnn:
+            save_cnn_features(args)
+            if args.use_augs:
+                warnings.warn("cannot perform augmentation on cnn features")
+        if args.use_pretrained == "evaluate":
+            train_nn.evaluate(args)
+            print("###  Evaluated from pretrained model  ###")
+        else:
+            print("### Starting to train.  ###")
+            train_nn.fit(args)
+            train_nn.evaluate(args)

train_nn.py

@@ -0,0 +1,300 @@
+import os
+import logging
+import json
+
+import torch
+import torch.nn.functional as F
+from torch.utils import data
+from sklearn.metrics import accuracy_score
+from tqdm import tqdm
+
+from models import CNN, LSTM, Transformer
+from configs import CnnConfig, LstmConfig, TransformerConfig
+from utils import (
+    seed_everything,
+    AverageMeter,
+    EarlyStopping,
+    load_label_map,
+    get_experiment_name,
+    load_json,
+)
+from dataset import KeypointsDataset, FeaturesDatset
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+def train(dataloader, model, optimizer, device):
+    model.train()
+
+    losses = AverageMeter()
+    accuracy = AverageMeter()
+
+    pbar = tqdm(dataloader, desc="Training")
+    for batch in pbar:
+        input_data = batch["data"].to(device)
+        label = batch["label"].to(device)
+
+        optimizer.zero_grad()
+        preds = model(input_data)
+        loss = F.cross_entropy(preds, label)
+        loss.backward()
+        optimizer.step()
+
+        losses.update(loss.item())
+        preds = preds.detach().cpu()
+        accuracy.update(
+            accuracy_score(
+                label.cpu().numpy(),
+                torch.argmax(torch.softmax(preds, dim=-1), dim=-1).numpy(),
+            )
+        )
+        pbar.set_postfix(loss=losses.avg, accuracy=accuracy.avg)
+
+    torch.cuda.empty_cache()
+    loss_avg = losses.avg
+    accuracy_avg = accuracy.avg
+    return loss_avg, accuracy_avg
+
+
+@torch.no_grad()
+def validate(dataloader, model, device):
+    model.eval()
+
+    losses = AverageMeter()
+    accuracy = AverageMeter()
+
+    pbar = tqdm(dataloader, desc="Eval")
+    for batch in pbar:
+        input_data = batch["data"].to(device)
+        label = batch["label"].to(device)
+
+        preds = model(input_data)
+        loss = F.cross_entropy(preds, label)
+
+        losses.update(loss.item())
+        preds = preds.detach().cpu()
+        accuracy.update(
+            accuracy_score(
+                label.cpu().numpy(),
+                torch.argmax(torch.softmax(preds, dim=-1), dim=-1).numpy(),
+            )
+        )
+        pbar.set_postfix(loss=losses.avg, accuracy=accuracy.avg)
+
+    torch.cuda.empty_cache()
+    loss_avg = losses.avg
+    accuracy_avg = accuracy.avg
+    return loss_avg, accuracy_avg
+
+
+def change_max_pos_embd(args, new_mpe_size, n_classes):
+    config = TransformerConfig(
+        size=args.transformer_size, max_position_embeddings=new_mpe_size
+    )
+    if args.use_cnn:
+        config.input_size = CnnConfig.output_dim
+    model = Transformer(config=config, n_classes=n_classes)
+    model = model.to(device)
+    return model
+
+
+def pretrained_name(args):
+    load_modelName = args.dataset
+    if args.use_cnn:
+        load_modelName += "_use_cnn"
+    else:
+        load_modelName += "_no_cnn"
+    if args.model == "lstm":
+        load_modelName += "_lstm.pth"
+    elif args.model == "transformer":
+        load_modelName += "_transformer"
+        if args.transformer_size == "large":
+            load_modelName += "_large.pth"
+        elif args.transformer_size == "small":
+            load_modelName += "_small.pth"
+    return load_modelName
+
+
+def load_pretrained(args, n_classes, model, optimizer=None, scheduler=None):
+    load_modelName = pretrained_name(args)
+    pretrained_model_links = load_json("pretrained_links.json")
+
+    if not os.path.isfile(load_modelName):
+        link = pretrained_model_links[load_modelName]
+        torch.hub.download_url_to_file(link, load_modelName, progress=True)
+
+    if args.model == "transformer":
+        model = change_max_pos_embd(args, new_mpe_size=256, n_classes=n_classes)
+
+    ckpt = torch.load(load_modelName)
+    model.load_state_dict(ckpt["model"])
+    if args.use_pretrained == "resume_training":
+        optimizer.load_state_dict(ckpt["optimizer"])
+        scheduler.load_state_dict(ckpt["scheduler"])
+
+    return model, optimizer, scheduler
+
+
+def fit(args):
+    exp_name = get_experiment_name(args)
+    logging_path = os.path.join(args.save_path, exp_name) + ".log"
+    logging.basicConfig(filename=logging_path, level=logging.INFO, format="%(message)s")
+    seed_everything(args.seed)
+    label_map = load_label_map(args.dataset)
+
+    if args.use_cnn:
+        train_dataset = FeaturesDatset(
+            features_dir=os.path.join(args.data_dir, f"{args.dataset}_train_features"),
+            label_map=label_map,
+            mode="train",
+        )
+        val_dataset = FeaturesDatset(
+            features_dir=os.path.join(args.data_dir, f"{args.dataset}_val_features"),
+            label_map=label_map,
+            mode="val",
+        )
+
+    else:
+        train_dataset = KeypointsDataset(
+            keypoints_dir=os.path.join(
+                args.data_dir, f"{args.dataset}_train_keypoints"
+            ),
+            use_augs=args.use_augs,
+            label_map=label_map,
+            mode="train",
+            max_frame_len=169,
+        )
+        val_dataset = KeypointsDataset(
+            keypoints_dir=os.path.join(args.data_dir, f"{args.dataset}_val_keypoints"),
+            use_augs=False,
+            label_map=label_map,
+            mode="val",
+            max_frame_len=169,
+        )
+
+    train_dataloader = data.DataLoader(
+        train_dataset,
+        batch_size=args.batch_size,
+        shuffle=True,
+        num_workers=4,
+        pin_memory=True,
+    )
+    val_dataloader = data.DataLoader(
+        val_dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=4,
+        pin_memory=True,
+    )
+
+    n_classes = 50
+    if args.dataset == "include":
+        n_classes = 263
+
+    if args.model == "lstm":
+        config = LstmConfig()
+        if args.use_cnn:
+            config.input_size = CnnConfig.output_dim
+        model = LSTM(config=config, n_classes=n_classes)
+    else:
+        config = TransformerConfig(size=args.transformer_size)
+        if args.use_cnn:
+            config.input_size = CnnConfig.output_dim
+        model = Transformer(config=config, n_classes=n_classes)
+
+    model = model.to(device)
+    optimizer = torch.optim.AdamW(
+        model.parameters(), lr=args.learning_rate, weight_decay=0.01
+    )
+    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
+        optimizer, mode="max", factor=0.2
+    )
+
+    if args.use_pretrained == "resume_training":
+        model, optimizer, scheduler = load_pretrained(
+            args, n_classes, model, optimizer, scheduler
+        )
+
+    model_path = os.path.join(args.save_path, exp_name) + ".pth"
+    es = EarlyStopping(patience=15, mode="max")
+    for epoch in range(args.epochs):
+        print(f"Epoch: {epoch+1}/{args.epochs}")
+        train_loss, train_acc = train(train_dataloader, model, optimizer, device)
+        val_loss, val_acc = validate(val_dataloader, model, device)
+        logging.info(
+            "Epoch: {}, train loss: {}, train acc: {}, val loss: {}, val acc: {}".format(
+                epoch + 1, train_loss, train_acc, val_loss, val_acc
+            )
+        )
+        scheduler.step(val_acc)
+        es(
+            model_path=model_path,
+            epoch_score=val_acc,
+            model=model,
+            optimizer=optimizer,
+            scheduler=scheduler,
+        )
+        if es.early_stop:
+            print("Early stopping")
+            break
+
+    print("### Training Complete ###")
+
+
+def evaluate(args):
+    label_map = load_label_map(args.dataset)
+    n_classes = 50
+    if args.dataset == "include":
+        n_classes = 263
+
+    if args.use_cnn:
+        dataset = FeaturesDatset(
+            features_dir=os.path.join(args.data_dir, f"{args.dataset}_test_features"),
+            label_map=label_map,
+            mode="test",
+        )
+
+    else:
+        dataset = KeypointsDataset(
+            keypoints_dir=os.path.join(args.data_dir, f"{args.dataset}_test_keypoints"),
+            use_augs=False,
+            label_map=label_map,
+            mode="test",
+            max_frame_len=169,
+        )
+
+    dataloader = data.DataLoader(
+        dataset,
+        batch_size=args.batch_size,
+        shuffle=False,
+        num_workers=4,
+        pin_memory=True,
+    )
+
+    if args.model == "lstm":
+        config = LstmConfig()
+        if args.use_cnn:
+            config.input_size = CnnConfig.output_dim
+        model = LSTM(config=config, n_classes=n_classes)
+    else:
+        config = TransformerConfig(size=args.transformer_size)
+        if args.use_cnn:
+            config.input_size = CnnConfig.output_dim
+        model = Transformer(config=config, n_classes=n_classes)
+
+    model = model.to(device)
+
+    if args.use_pretrained == "evaluate":
+        model, _, _ = load_pretrained(args, n_classes, model)
+        print("### Model loaded ###")
+
+    else:
+        exp_name = get_experiment_name(args)
+        model_path = os.path.join(args.save_path, exp_name) + ".pth"
+        ckpt = torch.load(model_path)
+        model.load_state_dict(ckpt["model"])
+        print("### Model loaded ###")
+
+    test_loss, test_acc = validate(dataloader, model, device)
+    print("Evaluation Results:")
+    print(f"Loss: {test_loss}, Accuracy: {test_acc}")

train_xgb.py

@@ -0,0 +1,170 @@
+import os
+
+from sklearn.metrics import accuracy_score
+import pandas as pd
+import numpy as np
+
+from models import Xgboost
+from configs import XgbConfig
+from utils import get_experiment_name, load_label_map
+from augment import (
+    plus7rotation,
+    minus7rotation,
+    gaussSample,
+    cutout,
+    upsample,
+    downsample,
+)
+from tqdm.auto import tqdm
+
+
+def flatten(arr, max_seq_len=200):
+    arr = np.array(arr)
+    arr = np.pad(arr, ((0, max_seq_len - arr.shape[0]), (0, 0)), "constant")
+    arr = arr.flatten()
+    return arr
+
+
+def combine_xy(x, y):
+    x, y = np.array(x), np.array(y)
+    _, length = x.shape
+    x = x.reshape((-1, length, 1))
+    y = y.reshape((-1, length, 1))
+    return np.concatenate((x, y), -1).astype(np.float32)
+
+
+def split_xy(data):
+    value_x, value_y = [], []
+    for row in data:
+        row = np.asarray(row)
+        if row.shape == ():
+            continue
+        value_x.append(row[:, 0])
+        value_y.append(row[:, 1])
+    value_x, value_y = np.asarray(value_x), np.asarray(value_y)
+    return value_x, value_y
+
+
+def augment_sample(df, augs):
+    df = df.copy()
+    pose = combine_xy(df.pose_x, df.pose_y)
+    h1 = combine_xy(df.hand1_x, df.hand1_y)
+    h2 = combine_xy(df.hand2_x, df.hand2_y)
+    input_df = pd.DataFrame.from_dict(
+        {
+            "uid": df.uid,
+            "pose": pose.tolist(),
+            "hand1": h1.tolist(),
+            "hand2": h2.tolist(),
+            "label": df.label,
+        }
+    )
+    augmented_samples = []
+    for augmentation in augs:
+        df_augmented = augmentation(input_df)
+        pose_x, pose_y = split_xy(df_augmented.pose)
+        hand1_x, hand1_y = split_xy(df_augmented.hand1)
+        hand2_x, hand2_y = split_xy(df_augmented.hand2)
+        save_df = pd.Series(
+            {
+                "uid": df.uid + "_" + augmentation.__name__,
+                "label": df.label,
+                "pose_x": pose_x.tolist(),
+                "pose_y": pose_y.tolist(),
+                "hand1_x": hand1_x.tolist(),
+                "hand1_y": hand1_y.tolist(),
+                "hand2_x": hand2_x.tolist(),
+                "hand2_y": hand2_y.tolist(),
+                "n_frames": df.n_frames,
+            }
+        )
+        augmented_samples.append(save_df)
+
+    return pd.concat(augmented_samples, axis=0)
+
+
+def preprocess(df, use_augs, label_map, mode):
+    feature_cols = ["pose_x", "pose_y", "hand1_x", "hand1_y", "hand2_x", "hand2_y"]
+    x, y = [], []
+    i = 0
+    no_of_videos = df.shape[0]
+    pbar = tqdm(total=no_of_videos, desc=f"Processing {mode} file....")
+    while i < no_of_videos:
+        if use_augs and mode == "train":
+            augs = [
+                plus7rotation,
+                minus7rotation,
+                gaussSample,
+                cutout,
+                upsample,
+                downsample,
+            ]
+            augmented_rows = augment_sample(df.iloc[i], augs)
+            df = pd.concat([df, augmented_rows], axis=0)
+        row = df.loc[i, feature_cols]
+        flatten_features = np.hstack(list(map(flatten, row.values)))
+        x.append(flatten_features)
+        y.append(label_map[df.loc[i, "label"]])
+        i += 1
+        pbar.update(1)
+    x = np.stack(x)
+    y = np.array(y)
+    return x, y
+
+
+def load_dataframe(files):
+    series = []
+    for file_path in files:
+        series.append(pd.read_json(file_path, typ="series"))
+    return pd.concat(series, axis=0)
+
+
+def fit(args):
+    train_files = sorted(
+        glob.glob(
+            os.path.join(args.data_dir, f"{args.dataset}_train_keypoints", "*.json")
+        )
+    )
+    val_files = sorted(
+        glob.glob(
+            os.path.join(args.data_dir, f"{args.dataset}_val_keypoints", "*.json")
+        )
+    )
+
+    train_df = load_dataframe(train_files)
+    val_df = load_dataframe(val_files)
+
+    label_map = load_label_map(args.dataset)
+    x_train, y_train = preprocess(train_df, args.use_augs, label_map, "train")
+    x_val, y_val = preprocess(val_df, args.use_augs, label_map, "val")
+
+    config = XgbConfig()
+    model = Xgboost(config=config)
+    model.fit(x_train, y_train, x_val, y_val)
+
+    exp_name = get_experiment_name(args)
+    save_path = os.path.join(args.save_dir, exp_name, ".pickle.dat")
+    model.save(save_path)
+
+
+def evaluate(args):
+    test_files = sorted(
+        glob.glob(
+            os.path.join(args.data_dir, f"{args.dataset}_test_keypoints", "*.json")
+        )
+    )
+
+    test_df = load_dataframe(test_files)
+
+    label_map = load_label_map(args.dataset)
+    x_test, y_test = preprocess(test_df, args.use_augs, label_map, "test")
+
+    exp_name = get_experiment_name(args)
+    config = XgbConfig()
+    model = Xgboost(config=config)
+    load_path = os.path.join(args.save_dir, exp_name, ".pickle.dat")
+    model.load(load_path)
+    print("### Model loaded ###")
+
+    test_preds = model(x_test)
+    print("Test accuracy:", accuracy_score(y_test, test_preds))

utils.py

@@ -0,0 +1,103 @@
+import torch
+import random
+import os
+import numpy as np
+import json
+
+
+def seed_everything(seed):
+    random.seed(seed)
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+    torch.backends.cudnn.deterministic = True
+    np.random.seed(seed)
+    os.environ["PYTHONHASHSEED"] = str(seed)
+
+
+def load_json(path):
+    with open(path, "r") as f:
+        json_file = json.load(f)
+    return json_file
+
+
+def load_label_map(dataset):
+    file_path = f"label_maps/label_map_{dataset}.json"
+    return load_json(file_path)
+
+
+def get_experiment_name(args):
+    exp_name = ""
+    if args.use_cnn:
+        exp_name += "cnn_"
+    if args.use_augs:
+        exp_name += "augs_"
+    exp_name += args.model
+    return exp_name
+
+
+class AverageMeter:
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+
+class EarlyStopping:
+    def __init__(self, patience=5, mode="min", delta=0.0):
+        self.patience = patience
+        self.counter = 0
+        self.mode = mode
+        self.best_score = None
+        self.early_stop = False
+        self.delta = delta
+        if self.mode == "min":
+            self.val_score = np.inf
+        else:
+            self.val_score = -np.inf
+
+    def __call__(self, model_path, epoch_score, model, optimizer, scheduler=None):
+
+        if self.mode == "min":
+            score = -1.0 * epoch_score
+        else:
+            score = np.copy(epoch_score)
+
+        if self.best_score is None:
+            self.best_score = score
+            self.save_checkpoint(epoch_score, model, optimizer, scheduler, model_path)
+        elif score <= self.best_score + self.delta:
+            self.counter += 1
+            if self.counter >= self.patience:
+                self.early_stop = True
+        else:
+            self.best_score = score
+            self.save_checkpoint(epoch_score, model, optimizer, scheduler, model_path)
+            self.counter = 0
+
+    def save_checkpoint(self, epoch_score, model, optimizer, scheduler, model_path):
+        if epoch_score not in [-np.inf, np.inf, -np.nan, np.nan]:
+            print(
+                "Validation score improved ({} --> {}). Saving model!".format(
+                    self.val_score, epoch_score
+                )
+            )
+            torch.save(
+                {
+                    "model": model.state_dict(),
+                    "optimizer": optimizer.state_dict(),
+                    "scheduler": scheduler.state_dict() if scheduler else scheduler,
+                    "score": epoch_score,
+                },
+                model_path,
+            )
+        self.val_score = epoch_score