main.py

import numpy as np
import torch
from torch.utils.data import TensorDataset, DataLoader, random_split, Dataset, WeightedRandomSampler
import tonic
import tonic.transforms as TT
from tqdm import tqdm
import os
import random
from typing import List as list
from typing import Dict as dict
from typing import Tuple as tuple
from lib.train import Trainer
from lib.quantize import Quantizer
from lib.utils import TRAIN_FLAGS
import torch.nn.functional as F
import torchvision.transforms as T
import cv2
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import OneHotEncoder

os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

class TransformedTensorDataset(Dataset):
    def __init__(self, data_tensor, target_tensor, transform=None):
        self.data_tensor = data_tensor
        self.target_tensor = target_tensor
        self.transform = transform

    def __getitem__(self, index):
        x = self.data_tensor[index]
        y = self.target_tensor[index]

        # Convert to PIL image for torchvision transforms if needed
        if self.transform:
            # x shape: [C, H, W] expected by transforms — add channel if needed
            if x.ndim == 2:  # likely [H, W]
                x = x.unsqueeze(0)  # make it [1, H, W]
            x = self.transform(x)

        return x, y

    def __len__(self):
        return len(self.data_tensor)


class Controller:
    def __init__(self) -> None:
        self.args = TRAIN_FLAGS()

        # Set all random seeds
        seed = self.args.seed
        np.random.seed(seed)
        torch.manual_seed(seed)
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
        random.seed(seed)
        generator = torch.Generator().manual_seed(seed)
        
        # Ensure deterministic behavior
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
        torch.use_deterministic_algorithms(True)
        torch.cuda.set_device(self.args.gpu)

        X_train = np.load(os.path.join(self.args.data_dir, "train_frames_20k_ETS_DVS_multi.npy"))
        y_train = np.load(os.path.join(self.args.data_dir, "train_labels_20k_ETS_DVS_multi.npy"))
        X_test = np.load(os.path.join(self.args.data_dir, "test_frames_20k_ETS_DVS_multi.npy"))
        y_test = np.load(os.path.join(self.args.data_dir, "test_labels_20k_ETS_DVS_multi.npy"))

        # Convert NumPy arrays to PyTorch tensors
        test_frames_tensor = torch.tensor(X_test, dtype=torch.float32)
        test_labels_tensor = torch.tensor(y_test, dtype=torch.long)
        train_frames_tensor = torch.tensor(X_train, dtype=torch.float32)
        train_labels_tensor = torch.tensor(y_train, dtype=torch.long)

        online_transform = T.Compose([
            T.RandomResizedCrop((128, 128), scale=(0.8, 1.0)),
            T.RandomHorizontalFlip(p=0.5),
            T.RandomRotation(degrees=15),
            T.Lambda(lambda x: add_speckle_noise(x, std=0.2)),
            T.RandomErasing(p=0.5, scale=(0.1, 0.3), ratio=(0.3, 3.0)),
        ])

        # Wrap with custom dataset that supports transforms
        full_train_dataset = TransformedTensorDataset(train_frames_tensor, train_labels_tensor, transform=None)
        test_dataset = TransformedTensorDataset(test_frames_tensor, test_labels_tensor, transform=None)  

        # Split train/val
        train_size = int(0.8 * len(full_train_dataset))
        val_size = len(full_train_dataset) - train_size
        train_dataset, val_dataset = random_split(full_train_dataset, [train_size, val_size], generator=generator)
       
        # DataLoaders
        self.train_dataloader = DataLoader(train_dataset,
            batch_size=self.args.batch_size, num_workers=self.args.workers,
            pin_memory=True, persistent_workers=False, shuffle=True)

        self.dev_dataloader = DataLoader(val_dataset,
            batch_size=self.args.batch_size, num_workers=self.args.workers,
            pin_memory=True, persistent_workers=True, shuffle=False)

        self.test_dataloader = DataLoader(test_dataset,
            batch_size=self.args.batch_size, num_workers=self.args.workers,
            pin_memory=True, persistent_workers=True, shuffle=False)
        
        
        self.args.gestures = ['hand clapping', 'right hand wave', 'left hand wave', 'right arm CW', 
                              'right arm CCW', 'left arm CW', 'left arm CCW', 'arm roll', 
                              'air drums', 'air guitar', 'other']
        self.args.num_classes = len(self.args.gestures)

    def train_model(self) -> None:
        trainer = Trainer(self.args, len(self.train_dataloader))
        trainer.set_dataloaders(self.train_dataloader, self.dev_dataloader,
                                self.test_dataloader)
        trainer.run()

        return None

    def zero_channel_dropout(self) -> None:
        self.train_dataloader.dataset.channel_dropout = 0
        return None

    def quantize_model(self) -> None:
        quantizer = Quantizer(self.args, len(self.train_dataloader))
        quantizer.set_dataloaders(self.train_dataloader, self.dev_dataloader,
                                  self.test_dataloader)
        self.zero_channel_dropout()
        quantizer.run()
        return None
    

def main() -> None:
    controller = Controller()
    if controller.args.train:
        controller.train_model()
    if controller.args.quantize:
        controller.quantize_model()


if __name__ == "__main__":
    main()