{"cells":[{"cell_type":"markdown","metadata":{},"source":"# CMI 2025 - MTFN Model Submission\n\nMulti-Modal Temporal Fusion Network for gesture classification"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport numpy as np\nimport polars as pl\nimport os\nfrom pathlib import Path\nimport zipfile\nimport warnings\nwarnings.filterwarnings('ignore')\n\nimport kaggle_evaluation.cmi_inference_server"},{"cell_type":"markdown","metadata":{},"source":"## MTFN Model Definition"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"class MultiScaleConv1D(nn.Module):\n    \"\"\"Multi-scale 1D convolution for temporal feature extraction\"\"\"\n    def __init__(self, in_channels, out_channels, scales=[10, 30, 100, 300]):\n        super().__init__()\n        self.scales = scales\n        self.convs = nn.ModuleList([\n            nn.Conv1d(in_channels, out_channels // len(scales), \n                     kernel_size=scale, padding=scale//2)\n            for scale in scales\n        ])\n        self.bn = nn.BatchNorm1d(out_channels)\n        self.relu = nn.ReLU(inplace=True)\n        \n    def forward(self, x):\n        features = []\n        for conv in self.convs:\n            features.append(conv(x))\n        x = torch.cat(features, dim=1)\n        x = self.bn(x)\n        x = self.relu(x)\n        return x"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"class CrossModalAttention(nn.Module):\n    \"\"\"Cross-modal attention between IMU and TOF sensors\"\"\"\n    def __init__(self, imu_dim, tof_dim, hidden_dim=256):\n        super().__init__()\n        self.imu_proj = nn.Linear(imu_dim, hidden_dim)\n        self.tof_proj = nn.Linear(tof_dim, hidden_dim)\n        self.attention = nn.MultiheadAttention(\n            embed_dim=hidden_dim,\n            num_heads=8,\n            dropout=0.1,\n            batch_first=True\n        )\n        self.norm = nn.LayerNorm(hidden_dim)\n        \n    def forward(self, imu_features, tof_features):\n        B, T, _ = imu_features.shape\n        imu_proj = self.imu_proj(imu_features)\n        tof_proj = self.tof_proj(tof_features)\n        attended, _ = self.attention(imu_proj, tof_proj, tof_proj)\n        return self.norm(attended + imu_proj)"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"class MTFN(nn.Module):\n    \"\"\"Multi-Modal Temporal Fusion Network\"\"\"\n    def __init__(self, imu_dim=7, tof_dim=320, thermal_dim=5, hidden_dim=256, num_classes=18):\n        super().__init__()\n        \n        # Multi-scale CNN for each modality\n        self.imu_cnn = MultiScaleConv1D(imu_dim, hidden_dim)\n        self.tof_cnn = MultiScaleConv1D(tof_dim, hidden_dim)\n        self.thermal_cnn = MultiScaleConv1D(thermal_dim, hidden_dim // 4)\n        \n        # Cross-modal attention\n        self.cross_attention = CrossModalAttention(hidden_dim, hidden_dim, hidden_dim)\n        \n        # Temporal modeling with bidirectional LSTM\n        self.lstm = nn.LSTM(\n            input_size=hidden_dim + hidden_dim // 4,  # fused + thermal\n            hidden_size=hidden_dim,\n            num_layers=3,\n            dropout=0.3,\n            batch_first=True,\n            bidirectional=True\n        )\n        \n        # Classification head\n        self.classifier = nn.Sequential(\n            nn.Linear(hidden_dim * 2, hidden_dim),  # *2 for bidirectional\n            nn.ReLU(),\n            nn.Dropout(0.5),\n            nn.Linear(hidden_dim, num_classes)\n        )\n        \n    def forward(self, imu, tof, thermal):\n        B, T, _ = imu.shape\n        \n        # Extract multi-scale features for each modality\n        # Convert to [batch, channels, time] for Conv1D\n        imu_feat = self.imu_cnn(imu.transpose(1, 2)).transpose(1, 2)  # [B, T, hidden_dim]\n        tof_feat = self.tof_cnn(tof.transpose(1, 2)).transpose(1, 2)  # [B, T, hidden_dim]\n        thermal_feat = self.thermal_cnn(thermal.transpose(1, 2)).transpose(1, 2)  # [B, T, hidden_dim//4]\n        \n        # Cross-modal attention between IMU and TOF\n        fused_feat = self.cross_attention(imu_feat, tof_feat)  # [B, T, hidden_dim]\n        \n        # Combine with thermal features\n        combined = torch.cat([fused_feat, thermal_feat], dim=2)  # [B, T, hidden_dim + hidden_dim//4]\n        \n        # Temporal modeling\n        lstm_out, _ = self.lstm(combined)  # [B, T, hidden_dim*2]\n        \n        # Global average pooling\n        pooled = torch.mean(lstm_out, dim=1)  # [B, hidden_dim*2]\n        \n        # Classification\n        output = self.classifier(pooled)  # [B, num_classes]\n        return output"},{"cell_type":"markdown","metadata":{},"source":"## Model Loading and Global Variables"},{"cell_type":"code","metadata":{},"outputs":[],"source":"# Global model variable\nmodel = None\ndevice = None\n\ndef load_model():\n    \"\"\"Load MTFN model (executed only once)\"\"\"\n    global model, device\n    if model is None:\n        print(\"Loading model on cuda...\")\n        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n        \n        # Try different possible model paths\n        possible_paths = [\n            \"/kaggle/input/cmi-mtfn-model-weights/mtfn_best.pth\",\n            \"/kaggle/input/cmi-mtfn-model/mtfn_best.pth\",\n            \"/kaggle/input/mtfn-model-weights/mtfn_model.pth\",\n            \"./mtfn_model.pth\",\n            \"./mtfn_best.pth\"\n        ]\n        \n        model_path = None\n        for path in possible_paths:\n            if os.path.exists(path):\n                model_path = path\n                break\n        \n        # Check for ZIP files\n        if model_path is None:\n            zip_paths = [\n                \"/kaggle/input/cmi-mtfn-model-weights/mtfn_model_kaggle.zip\",\n                \"/kaggle/input/cmi-mtfn-model/mtfn_model_kaggle.zip\",\n                \"/kaggle/input/mtfn-model-weights/mtfn_model.zip\"\n            ]\n            for zip_path in zip_paths:\n                if os.path.exists(zip_path):\n                    with zipfile.ZipFile(zip_path, 'r') as zip_ref:\n                        zip_ref.extractall(\"/tmp/\")\n                    # Look for extracted model file\n                    for extracted_name in [\"mtfn_best.pth\", \"mtfn_model.pth\"]:\n                        extracted_path = f\"/tmp/{extracted_name}\"\n                        if os.path.exists(extracted_path):\n                            model_path = extracted_path\n                            break\n                    break\n        \n        if model_path is None:\n            raise FileNotFoundError(\"MTFN model file not found. Please add the model dataset.\")\n        \n        # Initialize model and load weights\n        model = MTFN(\n            imu_dim=7,\n            tof_dim=320,\n            thermal_dim=5,\n            hidden_dim=256,\n            num_classes=18\n        )\n        \n        # Load state dict with exact matching (no strict=False)\n        state_dict = torch.load(model_path, map_location=device)\n        model.load_state_dict(state_dict)  # This should work with correct model structure\n        model.to(device)\n        model.eval()\n        \n        print(f\"Model loaded successfully on {device}!\")\n\n# Load model on startup\nload_model()"},{"cell_type":"markdown","metadata":{},"source":"## Gesture Mapping and Feature Extraction"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"# Define the 18 gesture classes\nGESTURE_NAMES = [\n    'Above ear - pull hair',\n    'Cheek - pinch skin',\n    'Eyebrow - pull hair',\n    'Eyelash - pull hair',\n    'Forehead - pull hairline',\n    'Forehead - scratch',\n    'Neck - pinch skin',\n    'Neck - scratch',\n    'Write name on leg',\n    'Wave hello',\n    'Glasses on/off',\n    'Text on phone',\n    'Write name in air',\n    'Feel around in tray and pull out an object',\n    'Scratch knee/leg skin',\n    'Pull air toward your face',\n    'Drink from bottle/cup',\n    'Pinch knee/leg skin'\n]\n\ndef extract_features(sequence_df):\n    \"\"\"Extract IMU, TOF, and Thermal features from the sequence\"\"\"\n    # IMU features (accel_x, accel_y, accel_z, gyro_x, gyro_y, gyro_z, quat_w)\n    imu_cols = ['accel_x', 'accel_y', 'accel_z', 'gyro_x', 'gyro_y', 'gyro_z', 'quat_w']\n    imu_data = sequence_df.select(imu_cols).to_numpy().astype(np.float32)\n    \n    # TOF features\n    tof_cols = [f'tof_{i}' for i in range(320)]\n    tof_data = sequence_df.select(tof_cols).to_numpy().astype(np.float32)\n    \n    # Thermal features\n    thermal_cols = [f'thermal_{i}' for i in range(5)]\n    thermal_data = sequence_df.select(thermal_cols).to_numpy().astype(np.float32)\n    \n    return imu_data, tof_data, thermal_data"},{"cell_type":"markdown","metadata":{},"source":"## Prediction Function"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"def predict(sequence: pl.DataFrame, demographics: pl.DataFrame) -> str:\n    \"\"\"\n    Predict gesture from sensor sequence\n    \n    Args:\n        sequence: Sensor data DataFrame\n        demographics: Demographics DataFrame (not used in this model)\n        \n    Returns:\n        Predicted gesture name\n    \"\"\"\n    try:\n        global model, device\n        \n        # Extract sensor features\n        imu_data, tof_data, thermal_data = extract_features(sequence)\n        \n        # Convert to tensors and add batch dimension\n        imu_tensor = torch.from_numpy(imu_data).unsqueeze(0).to(device)  # [1, T, 7]\n        tof_tensor = torch.from_numpy(tof_data).unsqueeze(0).to(device)  # [1, T, 320]\n        thermal_tensor = torch.from_numpy(thermal_data).unsqueeze(0).to(device)  # [1, T, 5]\n        \n        # Model prediction\n        with torch.no_grad():\n            output = model(imu_tensor, tof_tensor, thermal_tensor)  # [1, 18]\n            prediction_idx = torch.argmax(output, dim=1).item()\n        \n        return GESTURE_NAMES[prediction_idx]\n        \n    except Exception as e:\n        print(f\"Error in prediction: {e}\")\n        # Return default gesture in case of error\n        return 'Wave hello'"},{"cell_type":"markdown","metadata":{},"source":"## Run Inference Server"},{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"# Create inference server\ninference_server = kaggle_evaluation.cmi_inference_server.CMIInferenceServer(predict)\n\n# Run server or local test based on environment\nif os.getenv('KAGGLE_IS_COMPETITION_RERUN'):\n    # Competition environment\n    print(\"Running in competition environment...\")\n    inference_server.serve()\nelse:\n    # Local test environment\n    print(\"Running local test...\")\n    inference_server.run_local_gateway(\n        data_paths=(\n            '/kaggle/input/cmi-detect-behavior-with-sensor-data/test.csv',\n            '/kaggle/input/cmi-detect-behavior-with-sensor-data/test_demographics.csv',\n        )\n    )"}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.8.5"}},"nbformat":4,"nbformat_minor":4}