Major updates : NTU added along with GCN+Transformer

Author: abhijaysingh

config/model.json

@@ -1,9 +1,14 @@
 {
-    "d_model" : 512,
-    "nhead" : 16,
-    "num_layers" : 8,
-    "num_features" : 99,
-    "dropout" : 0.5,
-    "dim_feedforward" : 2048,
-    "num_classes" : 2
+    "gcn_num_features" : 3,
+    "gcn_hidden_dim1" : 32,
+    "gcn_hidden_dim2" : 64,
+    "gcn_output_dim" : 128,
+    
+    "transformer_d_model" : 128,
+    "transformer_nhead" : 4,
+    "transformer_num_layers" : 2,
+    "transformer_num_features" : 128,
+    "transformer_dropout" : 0.3,
+    "transformer_dim_feedforward" : 256,
+    "transformer_num_classes" : 2
 }

data_mgmt/dataloaders/gcn_transformer.py

@@ -0,0 +1,99 @@
+import torch
+from torch_geometric.data import Batch
+from torch.utils.data.dataloader import default_collate
+from typing import Any, List, Mapping, Sequence, Tuple
+
+
+class Collater:
+    """
+    Collates the batch of data
+
+    Parameters
+    ----------
+    dataset : torch.utils.data.Dataset
+        Dataset to collate
+    """
+
+    def __init__(self, dataset):
+        self.dataset = dataset
+
+    def __call__(self, batch) -> Any:
+        """
+        Collates the batch of data
+
+        Parameters
+        ----------
+        batch : List[Any]
+            Batch of data
+
+        Returns
+        -------
+        Any
+            Collated batch of data
+        """
+        elem = batch[0]
+
+        if isinstance(elem, torch.Tensor):
+            return default_collate(batch)
+        elif isinstance(elem, float):
+            return torch.tensor(batch, dtype=torch.float)
+        elif isinstance(elem, int):
+            return torch.tensor(batch)
+        elif isinstance(elem, str):
+            return batch
+        elif isinstance(elem, Mapping):
+            return {key: self([data[key] for data in batch]) for key in elem}
+        elif isinstance(elem, Sequence) and not isinstance(elem, str):
+            return [self(s) for s in zip(*batch)]
+
+        raise TypeError(f"DataLoader found invalid type: '{type(elem)}'")
+
+    def collate_fn(self, batch: List[Any]) -> Any:
+        """
+        Collates the batch of data
+
+        Parameters
+        ----------
+        batch : List[Any]
+            Batch of data
+
+        Returns
+        -------
+        Any
+            Collated batch of data
+        """
+        batched_graphs = [item["poses"] for item in batch]
+        labels = [item["label"] for item in batch]
+
+        for i in range(len(batched_graphs)):
+            batched_graphs[i] = Batch.from_data_list(batched_graphs[i])
+
+        labels = torch.tensor(labels, dtype=torch.long)
+
+        return batched_graphs, labels
+
+
+class DataLoader(torch.utils.data.DataLoader):
+    """
+    Dataloader for the single view case
+
+    Parameters
+    ----------
+    dataset : torch.utils.data.Dataset
+        Dataset to load
+    batch_size : int
+        Batch size, by default 1
+    shuffle : bool, optional
+        Whether to shuffle the dataset, by default False
+    """
+
+    def __init__(self, dataset, batch_size: int = 1, shuffle: bool = False, **kwargs):
+        self.collator = Collater(dataset)
+
+        super().__init__(
+            dataset,
+            batch_size,
+            shuffle,
+            collate_fn=self.collator.collate_fn,
+            **kwargs,
+        )

data_mgmt/dataloader.py data_mgmt/dataloaders/transformer.py

@@ -61,14 +61,14 @@ def collate_fn(self, batch: List[Any]) -> Any:
         Any
             Collated batch of data
         """
-        poses = [item[0] for item in batch]
-        labels = [item[1] for item in batch]
-
+        poses = [item["keypoints"] for item in batch]
+        labels = [item["label"] for item in batch]
+        
         max_length = max([item.shape[0] for item in poses])
         masks = [torch.ones(item.shape[0]) for item in poses]
+
         for i, item in enumerate(masks):
             masks[i] = torch.nn.functional.pad(item, (0, max_length - item.shape[0]))
-
         poses = [torch.nn.functional.pad(item, (0, 0, 0, max_length - item.shape[0])) for item in poses]
 
         poses = torch.stack(poses)
@@ -77,7 +77,6 @@ def collate_fn(self, batch: List[Any]) -> Any:
 
         return poses, masks, labels
 
-
 class DataLoader(torch.utils.data.DataLoader):
     """
     Dataloader for the single view case

data_mgmt/datasets/ntu_dataset.py

@@ -0,0 +1,279 @@
+import torch
+import numpy as np
+import os
+import regex as re
+
+from torch_geometric.data import Data
+from torch_geometric.data import Dataset
+
+from typing import Dict
+
+label_action = [
+    {"id": 0, "A043": "falling"},
+    {"id" : 1, "A008" : "sitting down"},
+    {"id": 1, "A026": "hopping (one foot jumping)"},
+]
+
+file_name_regex = r"S(\d{3})C001P(\d{3})R(\d{3})A(\d{3})"
+file_name_regex = re.compile(file_name_regex)
+
+
+def get_label(file_name: str) -> int:
+    """
+    Returns the label of the file
+
+    Parameters
+    ----------
+    file_name : str
+        Name of the file
+
+    Returns
+    -------
+    int
+        Label of the file
+    """
+    label = file_name[-4:]
+    for i in label_action:
+        if label in i:
+            return i["id"]
+    return -1
+
+
+def is_valid_file(file_name: str, skip: int = 11) -> bool:
+    """
+    Checks if the file is a valid file
+
+    Parameters
+    ----------
+    file_name : str
+        Name of the file
+    skip : int, optional
+        Number of frames to skip, by default 11
+
+    Returns
+    -------
+    bool
+        True if the file is valid, False otherwise
+    """
+    npy_file = file_name.endswith(".npy")
+    file_name = file_name.split("/")[-1].split(".")[0]
+
+    if file_name_regex.match(file_name) is None or get_label(file_name) == -1:
+        return False
+
+    return npy_file
+
+
+def get_edge_index():
+    POSE_CONNECTIONS = [
+        (3, 2),
+        (20, 8),
+        (8, 9),
+        (9, 10),
+        (10, 11),
+        (11, 24),
+        (11, 23),
+        (20, 4),
+        (4, 5),
+        (5, 6),
+        (6, 7),
+        (7, 21),
+        (7, 22),
+        (0, 1),
+        (1, 20),
+        (0, 16),
+        (0, 12),
+        (16, 17),
+        (17, 18),
+        (18, 19),
+        (12, 13),
+        (13, 14),
+        (14, 15),
+    ]
+    edge_index = torch.tensor(POSE_CONNECTIONS, dtype=torch.long).t().contiguous()
+
+    return edge_index
+
+
+def get_multiview_files(dataset_folder: str) -> list:
+    """
+    Returns a list of files that have multiple views
+
+    Parameters
+    ----------
+    dataset_folder : str
+        Path to the dataset folder
+
+    Returns
+    -------
+    list
+        List of files that have multiple views
+    """
+    multiview_files = []
+
+    for root, dirs, files in os.walk(dataset_folder):
+        for file in files:
+            if is_valid_file(file):
+                file_name = file.split("/")[-1].split(".")[0]
+
+                file_name = file_name.split("C001")
+                other_views = [
+                    file_name[0] + "C002" + file_name[1],
+                    file_name[0] + "C003" + file_name[1],
+                ]
+
+                not_exist = False
+                for view in other_views:
+                    if not os.path.exists(os.path.join(root, view + ".skeleton.npy")):
+                        not_exist = True
+                        break
+                if not_exist:
+                    continue
+
+                other_views.append(file_name[0] + "C001" + file_name[1])
+                for i in range(len(other_views)):
+                    other_views[i] = os.path.join(
+                        root, other_views[i] + ".skeleton.npy"
+                    )
+                multiview_files.append(other_views)
+
+    return multiview_files
+
+
+class NTUDataset(Dataset):
+    """
+    Dataset class for the keypoint dataset
+    """
+
+    def __init__(
+        self, dataset_folder: str, skip: int = 11, occlude: bool = False
+    ) -> None:
+        super().__init__(None, None, None)
+        self.dataset_folder = dataset_folder
+        self.edge_index = get_edge_index()
+
+        self.poses = []
+        self.labels = []
+        self.keypoints = []
+
+        self.occluded_kps = np.array([23, 24, 10, 11, 9, 8, 4, 5, 6, 7, 21, 22])
+
+        self.multi_view_files = get_multiview_files(dataset_folder)
+        for files in self.multi_view_files:
+            rand_view = np.random.randint(3)
+
+            for idx, file in enumerate(files):
+                file_data = np.load(file, allow_pickle=True).item()
+                frames = file_data["skel_body0"]
+
+                if occlude and idx == rand_view:
+                    frames = self._occlude_keypoints(frames)
+                pose_graphs = self._create_pose_graph(frames)
+
+                if "C001" in file:
+                    kps = self._get_flattened_keypoints(torch.tensor(frames))
+                    self.keypoints.append(kps)
+                    self.poses.append(pose_graphs)
+
+            file_name = files[0].split("/")[-1].split(".")[0]
+            self.labels.append(get_label(file_name))
+
+    def _create_pose_graph(self, keypoints: torch.Tensor) -> Data:
+        """
+        Creates a Pose Graph from the given keypoints and edge index
+
+        Parameters
+        ----------
+        keypoints : torch.Tensor
+            Keypoints of the pose
+        edge_index : torch.Tensor
+            Edge index of the pose
+
+        Returns
+        -------
+        Data
+            Pose Graph
+        """
+        pose_graphs = []
+        for t in range(keypoints.shape[0]):
+            pose_graph = Data(
+                x=torch.tensor(keypoints[t, :, :], dtype=torch.float),
+                edge_index=self.edge_index,
+            )
+            pose_graphs.append(pose_graph)
+
+        return pose_graphs
+
+    def _get_flattened_keypoints(self, keypoints: torch.Tensor) -> torch.Tensor:
+        """
+        Returns the flattened keypoints
+
+        Parameters
+        ----------
+        keypoints : torch.Tensor
+            Keypoints
+
+        Returns
+        -------
+        torch.Tensor
+            Flattened keypoints
+        """
+        return keypoints.reshape(keypoints.shape[0], -1)
+
+    def _occlude_keypoints(
+        self, frames: torch.Tensor, mask_prob: float = 0.2
+    ) -> torch.Tensor:
+        """
+        Occludes the keypoints of the pose
+
+        Parameters
+        ----------
+        frames : torch.Tensor
+            Keypoints of the pose
+        mask_prob : float, optional
+            Probability of masking the frames, by default 0.5
+
+        Returns
+        -------
+        torch.Tensor
+            Occluded frames
+        """
+        index = np.random.randint(3)
+        if index == 0:
+            mask_indices = np.arange(0, frames.shape[0] // 2)
+        elif index == 1:
+            mask_indices = np.arange(frames.shape[0] // 2, frames.shape[0])
+        else:
+            mask_indices = np.arange(frames.shape[0])
+
+        masked_kps = frames[mask_indices]
+        masked_kps[:, self.occluded_kps, :] = -1
+        frames[mask_indices] = masked_kps
+
+        return frames
+
+    def len(self) -> int:
+        """
+        Returns the number of samples in the dataset
+
+        Returns
+        -------
+        int : len
+            Number of samples in the dataset
+        """
+        return len(self.labels)
+
+    def get(self, index: int) -> Dict[str, torch.Tensor]:
+        """
+        Returns the sample at the given index
+
+        Returns
+        -------
+        Dict[str, torch.Tensor] : sample
+            Sample at the given index
+        """
+        keypoints = self.keypoints[index]
+        poses = self.poses[index]
+        label = self.labels[index]
+
+        return {"keypoints": keypoints, "poses": poses, "label": label}

data_mgmt/dataset.py data_mgmt/datasets/ur_dataset.py

@@ -1,16 +1,15 @@
 import torch
 from torch.utils.data import Dataset 
+from torch_geometric.data import Data
 import numpy as np
 import os
 
-# from dataloader import DataLoader
-
 from typing import Tuple
 
 def get_label(file_name: str) -> int:
     if "adl" in file_name:
-        return 0
-    return 1
+        return 1
+    return 0
 
 
 def is_valid_file(file_name: str, skip: int = 11) -> bool:
@@ -30,23 +29,68 @@ def is_valid_file(file_name: str, skip: int = 11) -> bool:
         True if the file is valid, False otherwise
     """
     npy_file = file_name.endswith(".npy")
-    cam0 = "cam0" in file_name
     skip_frame_num = file_name.split("/")[-1].split("-")[-2] == str(skip)
 
-    return npy_file and cam0 and skip_frame_num
-
+    return npy_file and skip_frame_num
 
-class KeypointsDataset(Dataset):
+def get_edge_index():
+    """
+    Returns the edge index of the pose graph
+    
+    Returns
+    -------
+    torch.Tensor
+        Edge index of the pose graph
+    """
+    POSE_CONNECTIONS = [
+        (0, 1),
+        (1, 2),
+        (2, 3),
+        (3, 7),  # Head to left shoulder
+        (0, 4),
+        (4, 5),
+        (5, 6),
+        (6, 8),  # Head to right shoulder
+        (9, 10),
+        (11, 12),  # Left and right shoulder
+        (11, 13),
+        (13, 15),
+        (15, 17),
+        (15, 19),
+        (15, 21),  # Left arm
+        (12, 14),
+        (14, 16),
+        (16, 18),
+        (16, 20),
+        (16, 22),  # Right arm
+        (11, 23),
+        (12, 24),
+        (23, 24),  # Torso
+        (23, 25),
+        (25, 27),
+        (27, 29),
+        (29, 31),  # Left leg
+        (24, 26),
+        (26, 28),
+        (28, 30),
+        (30, 32),  # Right leg
+    ]
+    edge_index = torch.tensor(POSE_CONNECTIONS, dtype=torch.long).t().contiguous()
+
+    return edge_index
+
+class URDataset(Dataset):
     """
     Dataset class for the keypoint dataset
     """
 
     def __init__(self, dataset_folder: str, skip: int = 11) -> None:
         self.dataset_folder = dataset_folder
+        self.edge_index = get_edge_index()
 
+        self.keypoints = []
         self.poses = []
         self.labels = []
-        self.file_names = []
 
         for root, dirs, files in os.walk(dataset_folder):
             for file in files:
@@ -55,11 +99,38 @@ def __init__(self, dataset_folder: str, skip: int = 11) -> None:
 
                     kps = np.load(file_path)
                     kps = kps[:, :, :3]
+                    pose_graphs = self._create_pose_graph(torch.tensor(kps))
                     kps = self._get_flattened_keypoints(torch.tensor(kps))
 
-                    self.poses.append(kps)
+                    self.poses.append(pose_graphs)
+                    self.keypoints.append(kps)
                     self.labels.append(get_label(file_path))
-                    self.file_names.append(file_path)
+
+    def _create_pose_graph(self, keypoints: torch.Tensor) -> Data:
+        """
+        Creates a Pose Graph from the given keypoints and edge index
+
+        Parameters
+        ----------
+        keypoints : torch.Tensor
+            Keypoints of the pose
+        edge_index : torch.Tensor
+            Edge index of the pose
+
+        Returns
+        -------
+        Data
+            Pose Graph
+        """
+        pose_graphs = []
+        for t in range(keypoints.shape[0]):
+            pose_graph = Data(
+                x=torch.tensor(keypoints[t, :, :], dtype=torch.float),
+                edge_index=self.edge_index,
+            )
+            pose_graphs.append(pose_graph)
+
+        return pose_graphs
 
     def _get_flattened_keypoints(self, keypoints: torch.Tensor) -> torch.Tensor:
         """
@@ -86,7 +157,7 @@ def __len__(self) -> int:
         int : len
             Number of samples in the dataset
         """
-        return len(self.poses)
+        return len(self.keypoints)
 
     def __getitem__(self, index: int) -> Tuple[torch.Tensor, int]:
         """
@@ -97,22 +168,8 @@ def __getitem__(self, index: int) -> Tuple[torch.Tensor, int]:
         dict : {kps, label, file_name}
             A dictionary containing the keypoint array, label and file name
         """
-        poses = self.poses[index]
+        keypoints = self.keypoints[index]
         label = self.labels[index]
-        return poses, label
-    
-# if __name__ == "__main__":
-#     dataset = KeypointsDataset("../data", skip=3)
-    
-#     train_dataset, test_dataset = torch.utils.data.random_split(dataset, [int(0.8 * len(dataset)), len(dataset) - int(0.8 * len(dataset))])
-
-#     train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True)
-#     test_loader = DataLoader(test_dataset, batch_size=2, shuffle=True)
-
-#     for batch in train_loader:
-#         print(batch[0].shape)
-#         print(batch[0])
-#         print(batch[1].shape)
-#         print(batch[1])
-#         print(batch[2].shape)
-#         break
+        poses = self.poses[index]
+
+        return {"keypoints": keypoints, "label": label, "poses": poses}

main.py

@@ -4,12 +4,15 @@
 from collections import Counter
 
 from trainer import Trainer
-from model import get_model
+from model import get_transformer_model, get_gcn_transformer_model
 from utils.logger import Logger
 from utils.model_config import ModelConfig
-from data_mgmt.dataset import KeypointsDataset
+from data_mgmt.datasets.ur_dataset import URDataset
+from data_mgmt.datasets.ntu_dataset import NTUDataset
 
-def parse_args():
+from typing import Tuple
+
+def parse_args() -> argparse.Namespace:
     parser = argparse.ArgumentParser(description="Train the model")
 
     parser.add_argument("--lr", type=float, default=1e-4, help="Learning rate")
@@ -19,6 +22,24 @@ def parse_args():
         default="./data",
         help="Path to the dataset folder",
     )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="transformer",
+        help="Model to use for training, transformer or gcn_transformer",
+    )
+    parser.add_argument(
+        "--dataset_type",
+        type=str,
+        default="ur",
+        help="Type of dataset to use, ntu or ur",    
+    )
+    parser.add_argument(
+        "--skip",
+        type=int,
+        default=11,
+        help="Number of frames to skip",
+    )
     parser.add_argument("--epochs", type=int, default=50, help="Number of epochs")
     parser.add_argument("--batch_size", type=int, default=32, help="Batch size")
     parser.add_argument(
@@ -42,14 +63,34 @@ def parse_args():
         default="./config/model.json",
         help="Path to the model config file",
     )
+    parser.add_argument(
+        "--occlude",
+        action="store_true",
+        help="Whether to occlude the input or not",
+    )
     args = parser.parse_args()
 
+    if args.dataset_type not in ["ntu", "ur"]:
+        raise ValueError("Dataset type should be either ntu or ur")
+    
+    if args.model not in ["transformer", "gcn_transformer"]:
+        raise ValueError("Model should be either transformer or gcn_transformer")
+    
+    if args.dataset_type == "ur":
+        if args.skip % 2 == 0:
+            raise ValueError("Skip frames should be odd")
+        if args.skip > 11:
+            raise ValueError("Skip frames should be less than 11")
+
     return args
 
 
-def load_dataset(dataset_folder, logger):
+def load_dataset(args : argparse.Namespace, logger : Logger) -> Tuple[torch.utils.data.Dataset, torch.utils.data.Dataset, torch.utils.data.Dataset]:
     np.random.seed(42)
-    dataset = KeypointsDataset(dataset_folder, skip=3)
+    if args.dataset_type == "ntu":
+        dataset = NTUDataset(args.dataset, occlude=args.occlude)
+    elif args.dataset_type == "ur":
+        dataset = URDataset(args.dataset, skip=args.skip)
 
     if len(dataset) > 0:
         logger.info("Dataset loaded successfully.")
@@ -88,19 +129,24 @@ def main():
     logger.info("\n")
     logger.info("Loading the dataset...")
     train_dataset, val_dataset, test_dataset = load_dataset(
-        args.dataset, logger
+        args, logger
     )
 
     logger.info(f"Training dataset size: {len(train_dataset)}")
     logger.info(f"Validation dataset size: {len(val_dataset)}")
     logger.info(f"Testing dataset size: {len(test_dataset)}")
 
     model_config = ModelConfig(args.model_config).get_config()
-    model, (train_dataloader, val_dataloader, test_dataloader) = get_model(
-        model_config, args, (train_dataset, val_dataset, test_dataset)
-    )
-
-    trainer = Trainer(model, lr=args.lr, logger=logger)
+    if args.model == "transformer":
+        model, (train_dataloader, val_dataloader, test_dataloader) = get_transformer_model(
+            model_config, args, (train_dataset, val_dataset, test_dataset)
+        )
+    elif args.model == "gcn_transformer":
+        model, (train_dataloader, val_dataloader, test_dataloader) = get_gcn_transformer_model(
+            model_config, args, (train_dataset, val_dataset, test_dataset)
+        )
+    
+    trainer = Trainer(model, lr=args.lr, logger=logger, model_type=args.model)
     logger.info(f"Batch size: {args.batch_size}")
     logger.info(f"Number of epochs: {args.epochs}")
     logger.info(f"Learning rate: {args.lr}")

model.py

@@ -2,16 +2,18 @@
 from torch.utils.data import Dataset
 
 from models.transformer import Transformer
-from data_mgmt.dataloader import DataLoader
+from models.action_recognizer import ActionRecognizer
+from data_mgmt.dataloaders.transformer import DataLoader as TransformerDataLoader
+from data_mgmt.dataloaders.gcn_transformer import DataLoader as GCNTransformerDataLoader
 
 from typing import Dict, Tuple
 
-def get_model(
+def get_transformer_model(
     config: Dict,
     args: argparse.Namespace,
     dataset: Tuple[Dataset, Dataset, Dataset],
 ) -> Tuple[
-    Transformer, Tuple[DataLoader, DataLoader, DataLoader]
+    Transformer, Tuple[TransformerDataLoader, TransformerDataLoader, TransformerDataLoader]
 ]:
     """
     Returns the model and the dataloader
@@ -31,19 +33,69 @@ def get_model(
         Model and the dataloaders
     """
     train_dataset, val_dataset, test_dataset = dataset
-    train_loader = DataLoader(
+    train_loader = TransformerDataLoader(
         train_dataset, batch_size=args.batch_size, shuffle=True
     )
-    val_loader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=True)
-    test_loader = DataLoader(
+    val_loader = TransformerDataLoader(val_dataset, batch_size=args.batch_size, shuffle=True)
+    test_loader = TransformerDataLoader(
         test_dataset, batch_size=args.batch_size, shuffle=True
     )
 
     return Transformer(
-        d_model=config["d_model"],
-        nhead=config["nhead"],
-        num_layers=config["num_layers"],
-        num_features=config["num_features"],
-        dropout=config["dropout"],
-        dim_ff=config["dim_feedforward"],
+        d_model=config["transformer_d_model"],
+        nhead=config["transformer_nhead"],
+        num_layers=config["transformer_num_layers"],
+        num_features=config["transformer_num_features"],
+        dropout=config["transformer_dropout"],
+        dim_ff=config["transformer_dim_feedforward"],
+        num_classes=config["transformer_num_classes"],
+        dataset=args.dataset_type,
     ), (train_loader, val_loader, test_loader)
+
+def get_gcn_transformer_model(
+    config: Dict,
+    args: argparse.Namespace,
+    dataset: Tuple[Dataset, Dataset, Dataset],
+) -> Tuple[
+    ActionRecognizer, Tuple[GCNTransformerDataLoader, GCNTransformerDataLoader, GCNTransformerDataLoader]
+]:
+    """
+    Returns the model and the dataloader
+
+    Parameters
+    ----------
+    config : Dict
+        Configuration for the model
+    args : argparse.Namespace
+        Arguments passed to the program
+    dataset : Tuple[Dataset, Dataset, Dataset]
+        Dataset to use for training, validation and testing
+
+    Returns
+    -------
+    Tuple[ActionRecognizer, Tuple[DataLoader, DataLoader, DataLoader]]
+        Model and the dataloaders
+    """
+    train_dataset, val_dataset, test_dataset = dataset
+    train_loader = GCNTransformerDataLoader(
+        train_dataset, batch_size=args.batch_size, shuffle=True
+    )
+    val_loader = GCNTransformerDataLoader(val_dataset, batch_size=args.batch_size, shuffle=True)
+    test_loader = GCNTransformerDataLoader(
+        test_dataset, batch_size=args.batch_size, shuffle=True
+    )
+
+    return ActionRecognizer(
+        gcn_num_features=config["gcn_num_features"],
+        gcn_hidden_dim1=config["gcn_hidden_dim1"],
+        gcn_hidden_dim2=config["gcn_hidden_dim2"],
+        gcn_output_dim=config["gcn_output_dim"],
+        transformer_d_model=config["transformer_d_model"],
+        transformer_nhead=config["transformer_nhead"],
+        transformer_num_layers=config["transformer_num_layers"],
+        transformer_num_features=config["transformer_num_features"],
+        transformer_dropout=config["transformer_dropout"],
+        transformer_dim_feedforward=config["transformer_dim_feedforward"],
+        transformer_num_classes=config["transformer_num_classes"],
+        dataset=args.dataset_type,
+    ), (train_loader, val_loader, test_loader)

models/action_recognizer.py

@@ -0,0 +1,87 @@
+import torch
+import torch.nn as nn
+from torch_geometric.data import Batch
+
+from models.transformer import Transformer
+from models.gcn import PoseGCN
+
+
+class ActionRecognizer(nn.Module):
+    def __init__(
+        self,
+        gcn_num_features: int,
+        gcn_hidden_dim1: int,
+        gcn_hidden_dim2: int,
+        gcn_output_dim: int,
+        transformer_d_model: int,
+        transformer_nhead: int,
+        transformer_num_layers: int,
+        transformer_num_features: int,
+        transformer_dropout: float = 0.1,
+        transformer_dim_feedforward: int = 2048,
+        transformer_num_classes: int = 2,
+        dataset: str = "ntu",
+    ) -> None:
+        """
+        Parameters
+        ----------
+        gcn_num_features : int
+            Number of features in the input sequence
+        gcn_hidden_dim1 : int
+            Dimension of the first hidden layer of the GCN
+        gcn_hidden_dim2 : int
+            Dimension of the second hidden layer of the GCN
+        gcn_output_dim : int
+            Dimension of the output layer of the GCN
+        transformer_d_model : int
+            Dimension of the input embedding
+        transformer_nhead : int
+            Number of attention heads
+        transformer_num_layers : int
+            Number of transformer encoder layers
+        transformer_num_features : int
+            Number of features in the input sequence
+        transformer_dropout : float, optional
+            Dropout rate, by default 0.1
+        transformer_dim_feedforward : int, optional
+            Dimension of the feedforward network, by default 2048
+        """
+        super(ActionRecognizer, self).__init__()
+
+        self.gcn = PoseGCN(
+            gcn_num_features, gcn_hidden_dim1, gcn_hidden_dim2, gcn_output_dim
+        )
+        self.transformer = Transformer(
+            transformer_d_model,
+            transformer_nhead,
+            transformer_num_layers,
+            transformer_num_features,
+            transformer_dropout,
+            transformer_dim_feedforward,
+            num_classes=transformer_num_classes,
+        )
+        self.num_classes = transformer_num_classes
+        self.dataset = dataset
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    def forward(self, batch: torch.Tensor) -> torch.Tensor:
+        """
+        Parameters
+        ----------
+        kps : torch.Tensor
+            Input sequence of keypoints
+
+        Returns
+        -------
+        torch.Tensor
+            Classification of the input sequence of keypoints
+        """
+        outputs = []
+
+        for item in batch:
+            view_embedding = self.gcn(item)
+
+            output = self.transformer(view_embedding.unsqueeze(0).to(self.device))
+            outputs.append(output)
+
+        return torch.stack(outputs).squeeze(1)

models/gcn.py

@@ -0,0 +1,56 @@
+import torch
+import torch.nn.functional as F
+from torch_geometric.nn import GCNConv
+from torch_geometric.data import Data, Batch
+from torch_geometric.nn import global_mean_pool
+
+
+class PoseGCN(torch.nn.Module):
+    def __init__(
+        self, num_features: int, hidden_dim1: int, hidden_dim2: int, output_dim: int
+    ) -> None:
+        """
+        Parameters
+        ----------
+        num_features : int
+            Number of features in the input sequence
+        hidden_dim1 : int
+            Dimension of the first hidden layer of the GCN
+        hidden_dim2 : int
+            Dimension of the second hidden layer of the GCN
+        output_dim : int
+            Dimension of the output layer of the GCN
+        """
+        super(PoseGCN, self).__init__()
+        self.conv1 = GCNConv(num_features, hidden_dim1)
+        self.conv2 = GCNConv(hidden_dim1, hidden_dim2)
+        self.conv3 = GCNConv(hidden_dim2, output_dim)
+
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    def forward(self, data: Batch) -> torch.Tensor:
+        """
+        Parameters
+        ----------
+        data : Data
+            Pose Graph
+
+        Returns
+        -------
+        torch.Tensor
+            Output of the GCN of shape (batch_size, output_dim)
+        """
+        x, edge_index, batch = (
+            data.x.to(self.device),
+            data.edge_index.to(self.device),
+            data.batch.to(self.device),
+        )
+
+        x = self.conv1(x, edge_index)
+        x = torch.relu(x)
+        x = self.conv2(x, edge_index)
+        x = torch.relu(x)
+        x = self.conv3(x, edge_index)
+
+        x = global_mean_pool(x, batch)
+        return x

models/transformer.py

@@ -51,6 +51,7 @@ def __init__(
         dropout: float = 0.1,
         dim_ff: int = 2048,
         num_classes: int = 2,
+        dataset: str = "ntu",
     ) -> None:
         """
         Parameters
@@ -74,6 +75,7 @@ def __init__(
         self.num_layers = num_layers
         self.num_features = num_features
         self.num_classes = num_classes
+        self.dataset = dataset
 
         self.pos_encoding = get_positional_encoding(
             1000, d_model
@@ -89,7 +91,7 @@ def __init__(
         )
         self.decoder = nn.Linear(self.d_model, self.num_classes)
 
-    def forward(self, x: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    def forward(self, x: torch.Tensor, mask: torch.Tensor = None) -> torch.Tensor:
         """
         Parameters
         ----------

trainer.py

@@ -22,6 +22,7 @@ def __init__(
         model: nn.Module,
         lr: float = 5e-5,
         logger: Logger = None,
+        model_type: str = "transformer",
     ) -> None:
         """
         Parameters
@@ -37,6 +38,8 @@ def __init__(
         self.logger = logger
         self.model = model
         self.lr = lr
+        self.model_type = model_type
+
         self.writer = SummaryWriter()
 
         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -148,8 +151,12 @@ def train_one_epoch(
         epoch_correct = 0
         epoch_count = 0
         for idx, batch in enumerate(iter(train_loader)):
-            predictions = self.model(batch[0].float().to(self.device), batch[1].to(self.device))
-            labels = batch[2].to(self.device)
+            if self.model_type == "transformer":
+                predictions = self.model(batch[0].float().to(self.device), batch[1].to(self.device))
+                labels = batch[2].to(self.device)
+            elif self.model_type == "gcn_transformer":
+                predictions = self.model(batch[0])
+                labels = batch[1].to(self.device)
 
             loss = self.criterion(predictions, labels)
             self.writer.add_scalar("Training loss per batch", loss, idx)
@@ -191,8 +198,12 @@ def evaluate(
             val_epoch_count = 0
 
             for idx, batch in enumerate(iter(val_loader)):
-                predictions = self.model(batch[0].float().to(self.device), batch[1].to(self.device))
-                labels = batch[2].to(self.device)
+                if self.model_type == "transformer":
+                    predictions = self.model(batch[0].float().to(self.device), batch[1].to(self.device))
+                    labels = batch[2].to(self.device)
+                elif self.model_type == "gcn_transformer":
+                    predictions = self.model(batch[0])
+                    labels = batch[1].to(self.device)
 
                 val_loss = self.criterion(predictions, labels)
                 self.writer.add_scalar("Validation loss per batch", val_loss, idx)
@@ -224,6 +235,7 @@ def test(
             Tuple containing the test epoch loss, test epoch correct
             and test epoch count
         """
+        output_path = os.path.join(output_path, self.model.dataset)
         if not os.path.exists(output_path):
             os.makedirs(output_path)
 
@@ -237,14 +249,19 @@ def test(
             + ".pt"
         )
         torch.save(self.best_model.state_dict(), os.path.join(output_path, file_name))
+        
         self.best_model.to(self.device)
         self.best_model.eval()
         with torch.no_grad():
             predictions = []
             labels = []
             for idx, batch in enumerate(iter(test_loader)):
-                predictions.extend(self.best_model(batch[0].float().to(self.device), batch[1].to(self.device)).argmax(axis=1).tolist())
-                labels.extend(batch[2].tolist())
+                if self.model_type == "transformer":
+                    predictions.extend(self.best_model(batch[0].float().to(self.device), batch[1].to(self.device)).argmax(axis=1).tolist())
+                    labels.extend(batch[2].tolist())
+                elif self.model_type == "gcn_transformer":
+                    predictions.extend(self.best_model(batch[0]).argmax(axis=1).tolist())
+                    labels.extend(batch[1].tolist())
 
             self.logger.info(f"Predictions: {predictions}")
             self.logger.info(f"Labels: {labels}")
@@ -253,11 +270,17 @@ def test(
             precision, recall, f1_score, _ = precision_recall_fscore_support(
                 labels, predictions, average="weighted"
             )
+            cm = confusion_matrix(labels, predictions)
+            tn, fp, fn, tp = cm.ravel()
+            sensitivity = tp / (tp + fn)
+            specificity = tn / (tn + fp)
+            geometric_mean = (sensitivity * specificity) ** 0.5
 
             self.logger.info(f"Accuracy: {accuracy:.4f}")
-            self.logger.info(f"Precision: {precision}")
-            self.logger.info(f"Recall: {recall}")
-            self.logger.info(f"F1 Score: {f1_score}")
+            self.logger.info(f"Precision: {precision:.4f}")
+            self.logger.info(f"Recall: {recall:.4f}")
+            self.logger.info(f"F1 Score: {f1_score:.4f}")
+            self.logger.info(f"G-Mean: {geometric_mean:.4f}")
 
             plt.figure(figsize=(8, 6))
             colors = cycle(["aqua", "darkorange"])
@@ -288,7 +311,6 @@ def test(
             plt.savefig(os.path.join(output_path, file_name))
             plt.show()
 
-            cm = confusion_matrix(labels, predictions)
             ax = sns.heatmap(
                 cm,
                 annot=True,
@@ -319,6 +341,7 @@ def _plot_losses(self, output_path: str) -> None:
         -------
         None
         """
+        output_path = os.path.join(output_path, self.model.dataset)
         if not os.path.exists(output_path):
             os.makedirs(output_path)