{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":102335,"databundleVersionId":12518947,"sourceType":"competition"},{"sourceId":245653776,"sourceType":"kernelVersion"},{"sourceId":452815,"sourceType":"modelInstanceVersion","modelInstanceId":359040,"modelId":380352}],"dockerImageVersionId":31041,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Inference","metadata":{}},{"cell_type":"code","source":"import os, json, joblib, numpy as np, pandas as pd\nimport random, math\nfrom pathlib import Path\nimport warnings \nwarnings.filterwarnings(\"ignore\")\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim import Adam\nfrom torch.optim.lr_scheduler import CosineAnnealingWarmRestarts\n\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.preprocessing import StandardScaler, LabelEncoder\nfrom sklearn.utils.class_weight import compute_class_weight\nfrom sklearn.model_selection import StratifiedKFold\nfrom timm.scheduler import CosineLRScheduler\nfrom scipy.signal import firwin\nimport polars as pl\nfrom tqdm import tqdm\nfrom transformers import BertModel, BertConfig\n\nimport numpy as np\nfrom scipy.fft import rfft, rfftfreq\n\n# ================================\nimport kagglehub\nmetric = kagglehub.package_import('jiazhuang/cmi-2025-metric')\n\ndef get_competition_score(true, pred):\n    assert len(true) == len(pred)\n    N = len(true)\n    true = pd.DataFrame({'id': range(N), 'gesture': true})\n    pred = pd.DataFrame({'id': range(N), 'gesture': pred})\n    return metric.score(true, pred, 'id')\n\n# ================================\n# TTA Helpers\n# ================================\ndef tta_jitter(x, sigma=0.01):\n    noise = np.random.randn(*x.shape) * sigma\n    return x + noise\n\n\n# Configuration\nTRAIN = True                     # ← set to True when you want to train\nRAW_DIR = Path(\"/kaggle/input/cmi-detect-behavior-with-sensor-data\")\nPRETRAINED_DIR = Path(\"/kaggle/input/cmi3-models-p\") # used when TRAIN=False\nEXPORT_DIR = Path(\"./\")                                    # artefacts will be saved here\nBATCH_SIZE = 64\nPAD_PERCENTILE = 100\nmaxlen = PAD_PERCENTILE\nLR_INIT = 1e-3\nWD = 3e-3\n# MIXUP_ALPHA = 0.4\nPATIENCE = 40\nFOLDS = 5\nrandom_state = 42\nepochs_warmup = 20\nwarmup_lr_init = 1.822126131809773e-05\nlr_min = 3.810323058740104e-09\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"▶ imports ready · pytorch {torch.__version__} · device: {device}\")\n\n# ================================\n# Model Components\n# ================================\nmean = torch.tensor([\n    0,  0, 0, 0, 0,\n    0,  9.0319e-03,  1.0849e+00, -2.6186e-03,  3.7651e-03,\n    -5.3660e-03, -2.8177e-03,  1.3318e-03, -1.5876e-04,  6.3495e-01,\n     6.2877e-01,  6.0607e-01,  6.2142e-01,  6.3808e-01,  6.5420e-01,\n     7.4102e-03, -3.4159e-03, -7.5237e-03, -2.6034e-02,  2.9704e-02,\n    -3.1546e-02, -2.0610e-03, -4.6986e-03, -4.7216e-03, -2.6281e-02,\n     1.5799e-02,  1.0016e-02,\n     0, 0, 0, 1  # rot_x, rot_y, rot_z, rot_w\n], dtype=torch.float32).view(1, -1, 1).to(device)         \n\nstd = torch.tensor([\n    1, 1, 1, 1, 1, 1, 0.2067, 0.8583, 0.3162,\n    0.2668, 0.2917, 0.2341, 0.3023, 0.3281, 1.0264, 0.8838, 0.8686, 1.0973,\n    1.0267, 0.9018, 0.4658, 0.2009, 0.2057, 1.2240, 0.9535, 0.6655, 0.2941,\n    0.3421, 0.8156, 0.6565, 1.1034, 1.5577,\n    1, 1, 1, 1  # rot_x, rot_y, rot_z, rot_w\n], dtype=torch.float32).view(1, -1, 1).to(device) + 1e-8  \n\nclass ImuFeatureExtractor(nn.Module):\n    def __init__(self, fs=100., add_quaternion=False, K=10):\n        super().__init__()\n        self.fs = fs\n        self.add_quaternion = add_quaternion\n        self.K = K  # Number of frequency components to keep\n\n        k = 15\n        self.lpf = nn.Conv1d(6, 6, kernel_size=k, padding=k//2,\n                             groups=6, bias=False)\n        nn.init.kaiming_uniform_(self.lpf.weight, a=math.sqrt(5))\n\n        self.lpf_acc  = nn.Conv1d(3, 3, k, padding=k//2, groups=3, bias=False)\n        self.lpf_gyro = nn.Conv1d(3, 3, k, padding=k//2, groups=3, bias=False)\n\n    def forward(self, imu):\n        B, C, T = imu.shape\n        acc  = imu[:, 0:3, :]                 # acc_x, acc_y, acc_z\n        gyro = imu[:, 3:6, :]                 # gyro_x, gyro_y, gyro_z\n        quat = imu[:, 6:10, :]                # rot_x, rot_y, rot_z, rot_w\n        extra = imu[:, 10:, :]                 \n\n        # 1) magnitude\n        acc_mag  = torch.norm(acc,  dim=1, keepdim=True)          # (B,1,T)\n        gyro_mag = torch.norm(gyro, dim=1, keepdim=True)\n\n        # 2) jerk \n        jerk = F.pad(acc[:, :, 1:] - acc[:, :, :-1], (1,0))       # (B,3,T)\n        gyro_delta = F.pad(gyro[:, :, 1:] - gyro[:, :, :-1], (1,0))\n\n        # 3) energy\n        acc_pow  = acc ** 2\n        gyro_pow = gyro ** 2\n\n        # 4) LPF / HPF \n        acc_lpf  = self.lpf_acc(acc)\n        acc_hpf  = acc - acc_lpf\n        gyro_lpf = self.lpf_gyro(gyro)\n        gyro_hpf = gyro - gyro_lpf\n\n        # 5) FFT Amplitude Spectrum\n        acc_fft = torch.from_numpy(np.abs(rfft(acc.cpu().numpy(), axis=-1)))[:, :, :self.K].to(acc.device)\n        gyro_fft = torch.from_numpy(np.abs(rfft(gyro.cpu().numpy(), axis=-1)))[:, :, :self.K].to(gyro.device)\n\n        # 6) Spectral Entropy\n        def spectral_entropy(signal):\n            psd = np.abs(rfft(signal, axis=-1)) ** 2\n            psd_norm = psd / np.sum(psd, axis=-1, keepdims=True)\n            se = -np.sum(psd_norm * np.log(psd_norm + 1e-12), axis=-1)\n            return torch.from_numpy(se).to(signal.device)\n\n        acc_entropy = spectral_entropy(acc.cpu().numpy())\n        gyro_entropy = spectral_entropy(gyro.cpu().numpy())\n\n        # 7) Spectral Centroid\n        def spectral_centroid(signal):\n            freqs = rfftfreq(signal.shape[-1], 1/self.fs)\n            centroid = np.sum(freqs * np.abs(rfft(signal, axis=-1)), axis=-1) / np.sum(np.abs(rfft(signal, axis=-1)), axis=-1)\n            return torch.from_numpy(centroid).to(signal.device)\n\n        acc_centroid = spectral_centroid(acc.cpu().numpy())\n        gyro_centroid = spectral_centroid(gyro.cpu().numpy())\n\n        # 8) Dominant Frequency\n        def dominant_frequency(signal):\n            freqs = rfftfreq(signal.shape[-1], 1/self.fs)\n            dom_freq = freqs[np.argmax(np.abs(rfft(signal, axis=-1)), axis=-1)]\n            return torch.from_numpy(dom_freq).to(signal.device)\n\n        acc_dom_freq = dominant_frequency(acc.cpu().numpy())\n        gyro_dom_freq = dominant_frequency(gyro.cpu().numpy())\n\n        features = [\n            acc, gyro, quat,\n            acc_mag, gyro_mag,\n            jerk, gyro_delta,\n            acc_pow, gyro_pow,\n            acc_lpf, acc_hpf,\n            gyro_lpf, gyro_hpf,\n            acc_fft, gyro_fft,\n            acc_entropy.unsqueeze(1), gyro_entropy.unsqueeze(1),\n            acc_centroid.unsqueeze(1), gyro_centroid.unsqueeze(1),\n            acc_dom_freq.unsqueeze(1), gyro_dom_freq.unsqueeze(1)\n        ]\n        return torch.cat(features, dim=1)  # (B, C_out, T)\n\n\nclass SEBlock(nn.Module):\n    def __init__(self, channels, reduction=8):\n        super().__init__()\n        self.squeeze = nn.AdaptiveAvgPool1d(1)\n        self.excitation = nn.Sequential(\n            nn.Linear(channels, channels // reduction, bias=False),\n            nn.ReLU(inplace=True),\n            nn.Linear(channels // reduction, channels, bias=False),\n            nn.Sigmoid()\n        )\n    \n    def forward(self, x):\n        b, c, _ = x.size()\n        y = self.squeeze(x).view(b, c)\n        y = self.excitation(y).view(b, c, 1)\n        return x * y.expand_as(x)\n\nclass ResidualSECNNBlock(nn.Module):\n    def __init__(self, in_channels, out_channels, kernel_size, pool_size=2, dropout=0.3, weight_decay=1e-4):\n        super().__init__()\n        \n        # First conv block\n        self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size, padding=kernel_size//2, bias=False)\n        self.bn1 = nn.BatchNorm1d(out_channels)\n        \n        # Second conv block\n        self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size, padding=kernel_size//2, bias=False)\n        self.bn2 = nn.BatchNorm1d(out_channels)\n        \n        # SE block\n        self.se = SEBlock(out_channels)\n        \n        # Shortcut connection\n        self.shortcut = nn.Sequential()\n        if in_channels != out_channels:\n            self.shortcut = nn.Sequential(\n                nn.Conv1d(in_channels, out_channels, 1, bias=False),\n                nn.BatchNorm1d(out_channels)\n            )\n        \n        self.pool = nn.MaxPool1d(pool_size)\n        self.dropout = nn.Dropout(dropout)\n        \n    def forward(self, x):\n        shortcut = self.shortcut(x)\n        \n        # First conv\n        out = F.relu(self.bn1(self.conv1(x)))\n        # Second conv\n        out = self.bn2(self.conv2(out))\n        \n        # SE block\n        out = self.se(out)\n        \n        # Add shortcut\n        out += shortcut\n        out = F.relu(out)\n        \n        # Pool and dropout\n        out = self.pool(out)\n        out = self.dropout(out)\n        \n        return out\n\nclass AttentionLayer(nn.Module):\n    def __init__(self, hidden_dim):\n        super().__init__()\n        self.attention = nn.Linear(hidden_dim, 1)\n        \n    def forward(self, x):\n        # x shape: (batch, seq_len, hidden_dim)\n        scores = torch.tanh(self.attention(x))  # (batch, seq_len, 1)\n        weights = F.softmax(scores.squeeze(-1), dim=1)  # (batch, seq_len)\n        context = torch.sum(x * weights.unsqueeze(-1), dim=1)  # (batch, hidden_dim)\n        return context\n\nclass CBAMBlock(nn.Module):\n    def __init__(self, channels, reduction=16, kernel_size=7):\n        super().__init__()\n        # Channel attention\n        self.avg_pool = nn.AdaptiveAvgPool1d(1)\n        self.max_pool = nn.AdaptiveMaxPool1d(1)\n        self.shared_mlp = nn.Sequential(\n            nn.Linear(channels, channels // reduction, bias=False),\n            nn.ReLU(),\n            nn.Linear(channels // reduction, channels, bias=False)\n        )\n        self.sigmoid_channel = nn.Sigmoid()\n\n        # Spatial attention\n        self.conv_spatial = nn.Conv1d(2, 1, kernel_size=kernel_size, padding=kernel_size//2, bias=False)\n        self.sigmoid_spatial = nn.Sigmoid()\n\n    def forward(self, x):\n        # Channel attention\n        b, c, t = x.size()\n        avg_out = self.shared_mlp(self.avg_pool(x).view(b, c))\n        max_out = self.shared_mlp(self.max_pool(x).view(b, c))\n        channel_att = self.sigmoid_channel(avg_out + max_out).view(b, c, 1)\n        x = x * channel_att.expand_as(x)\n\n        # Spatial attention\n        avg_out = torch.mean(x, dim=1, keepdim=True)\n        max_out, _ = torch.max(x, dim=1, keepdim=True)\n        spatial_att = self.sigmoid_spatial(self.conv_spatial(torch.cat([avg_out, max_out], dim=1)))\n        x = x * spatial_att.expand_as(x)\n\n        return x\n\nclass MultiModalModel(nn.Module):\n    def __init__(self, imu_dim, tof_dim, thm_dim, d_model, n_classes):\n        super().__init__()\n        self.imu_dim = imu_dim\n        self.tof_dim = tof_dim\n        self.thm_dim = thm_dim\n        self.d_model = d_model\n        self.n_classes = n_classes\n\n        # IMU deep branch\n        self.imu_block1 = ResidualSECNNBlock(imu_dim, 64, 3)\n        self.imu_block2 = ResidualSECNNBlock(64, 128, 5)\n        \n        # TOF lighter branch - 加强CNN结构\n        # 更新TOF分支以适应增加的统计特征维度\n        self.tof_branch = nn.Sequential(\n            # 第一层卷积：将TOF特征数映射到中间通道数\n            nn.Conv1d(tof_dim, 64, kernel_size=3, padding=1, bias=False),\n            nn.BatchNorm1d(64),  # 批归一化\n            nn.ReLU(inplace=True),                    # 激活函数\n            nn.MaxPool1d(2, ceil_mode=True),          # 下采样（减半时间维度）\n            nn.Dropout(0.3),       # 防止过拟合\n\n            # 第二层卷积：保持通道数，进一步提取特征\n            nn.Conv1d(64, 128, kernel_size=3, padding=1, bias=False),\n            nn.BatchNorm1d(128),\n            nn.ReLU(inplace=True),\n            nn.MaxPool1d(2, ceil_mode=True),          # 再次下采样\n            nn.Dropout(0.3)\n        )\n\n        # THM lighter branch\n        self.thm_conv1 = nn.Conv1d(thm_dim, 64, 3, padding=1, bias=False)\n        self.thm_bn1 = nn.BatchNorm1d(64)\n        self.thm_pool1 = nn.MaxPool1d(2)\n        self.thm_drop1 = nn.Dropout(0.3)\n        \n        self.thm_conv2 = nn.Conv1d(64, 128, 3, padding=1, bias=False)\n        self.thm_bn2 = nn.BatchNorm1d(128)\n        self.thm_pool2 = nn.MaxPool1d(2)\n        self.thm_drop2 = nn.Dropout(0.3)\n\n        # CBAM Block for each modality\n        self.cbam_imu = CBAMBlock(128)\n        self.cbam_tof = CBAMBlock(128)\n        self.cbam_thm = CBAMBlock(128)\n        \n        # BERT for fusion\n        bert_config = BertConfig(hidden_size=384, num_hidden_layers=4, num_attention_heads=8, intermediate_size=512)  # 128*3 for three modalities\n        self.bert = BertModel(bert_config)\n        \n        # Attentionjiapin\n        self.attention = AttentionLayer(384)  # 128*3 for three modalities\n        \n        # Dense layers\n        self.dense1 = nn.Linear(384, 256, bias=False)\n        self.bn_dense1 = nn.BatchNorm1d(256)\n        self.drop1 = nn.Dropout(0.3)\n        \n        self.dense2 = nn.Linear(256, 128, bias=False)\n        self.bn_dense2 = nn.BatchNorm1d(128)\n        self.drop2 = nn.Dropout(0.3)\n        \n        self.classifier = nn.Linear(128, n_classes)\n        \n    def forward(self, imu, tof, thm):\n        # IMU branch\n        x1 = self.imu_block1(imu)\n        x1 = self.imu_block2(x1)\n        x1 = self.cbam_imu(x1)\n        \n        # TOF branch - 使用加强的CNN结构\n        x2 = self.tof_branch(tof)\n        x2 = self.cbam_tof(x2)\n\n        # THM branch\n        x3 = F.relu(self.thm_bn1(self.thm_conv1(thm)))\n        x3 = self.thm_drop1(self.thm_pool1(x3))\n        x3 = F.relu(self.thm_bn2(self.thm_conv2(x3)))\n        x3 = self.thm_drop2(self.thm_pool2(x3))\n        x3 = self.cbam_thm(x3)\n\n        # Concatenate branches\n        merged = torch.cat([x1, x2, x3], dim=1).transpose(1, 2)  # (batch, seq_len, 384)\n        \n        # BERT for fusion\n        bert_output = self.bert(inputs_embeds=merged).last_hidden_state\n        \n        # Attention\n        attended = self.attention(bert_output)\n        \n        # Dense layers\n        x = F.relu(self.bn_dense1(self.dense1(attended)))\n        x = self.drop1(x)\n        x = F.relu(self.bn_dense2(self.dense2(x)))\n        x = self.drop2(x)\n        \n        # Classification\n        logits = (self.classifier(x))\n        return logits\n\n\n# ================================\n# Data Handling\n# ================================\n\ndef pad_sequences_torch(sequences, maxlen, padding='post', truncating='post', value=0.0):\n    result = []\n    for seq in sequences:\n        if len(seq) >= maxlen:\n            seq = seq[:maxlen] if truncating=='post' else seq[-maxlen:]\n        else:\n            pad_len = maxlen - len(seq)\n            pad_array = np.full((pad_len, seq.shape[1]), value)\n            seq = np.concatenate([seq, pad_array], axis=0) if padding=='post' else np.concatenate([pad_array, seq], axis=0)\n        result.append(seq)\n    return np.array(result, dtype=np.float32)\n\n# 添加TOF统计特征提取函数\ndef extract_tof_statistics(tof_data, tof_cols):\n    \"\"\"\n    从原始TOF数据中提取统计特征\n    包括全局统计和区域统计\n    \"\"\"\n    B, T, D = tof_data.shape\n    tof_reshaped = tof_data.reshape(B*T, D)  # (B*T, D)\n    \n    # 全局统计特征: 均值、标准差、最小值、最大值\n    tof_mean = np.mean(tof_reshaped, axis=1, keepdims=True)  # (B*T, 1)\n    tof_std = np.std(tof_reshaped, axis=1, keepdims=True)    # (B*T, 1)\n    tof_min = np.min(tof_reshaped, axis=1, keepdims=True)    # (B*T, 1)\n    tof_max = np.max(tof_reshaped, axis=1, keepdims=True)    # (B*T, 1)\n    \n    # 区域统计特征: 根据实际的TOF维度(320)划分为合适区域\n    # 将320个像素划分为20个区域，每个区域16个像素\n    num_regions = 20\n    pixels_per_region = D // num_regions  # 320 // 20 = 16\n    \n    tof_regions = tof_reshaped.reshape(B*T, num_regions, pixels_per_region)  # (B*T, 20, 16)\n    region_means = np.mean(tof_regions, axis=2)     # (B*T, 20)\n    region_stds = np.std(tof_regions, axis=2)       # (B*T, 20)\n    region_mins = np.min(tof_regions, axis=2)       # (B*T, 20)\n    region_maxs = np.max(tof_regions, axis=2)       # (B*T, 20)\n    \n    # 合并所有统计特征\n    stats_features = np.concatenate([\n        tof_mean, tof_std, tof_min, tof_max,           # 全局统计 (4)\n        region_means, region_stds, region_mins, region_maxs  # 区域统计 (20*4=80)\n    ], axis=1)  # (B*T, 84)\n    \n    # 重塑回(B, T, 84)\n    stats_features = stats_features.reshape(B, T, 84)\n    \n    return stats_features.astype(np.float32)\n\ndef preprocess_sequence(df_seq: pd.DataFrame, feature_cols: list, scaler: StandardScaler):\n    mat = df_seq[feature_cols].ffill().bfill().fillna(0).values\n    return scaler.transform(mat).astype('float32')\n\nclass CMI3Dataset(Dataset):\n    def __init__(self,\n                 X_list,\n                 y_list,\n                 maxlen,\n                 mode=\"train\",\n                 imu_dim=7,\n                 tof_dim=64,  # 添加TOF维度参数\n                 augment=None):\n        self.X_list = X_list\n        self.mode = mode\n        self.y_list = y_list\n        self.maxlen = maxlen\n        self.imu_dim = imu_dim     \n        self.tof_dim = tof_dim    # 存储TOF维度\n        self.augment = augment   \n\n    def pad_sequences_torch(self, seq, maxlen, padding='post', truncating='post', value=0.0):\n\n        if seq.shape[0] >= maxlen:\n            if truncating == 'post':\n                seq = seq[:maxlen]\n            else:  # 'pre'\n                seq = seq[-maxlen:]\n        else:\n            pad_len = maxlen - seq.shape[0]\n            if padding == 'post':\n                seq = np.concatenate([seq, np.full((pad_len, seq.shape[1]), value)])\n            else:  # 'pre'\n                seq = np.concatenate([np.full((pad_len, seq.shape[1]), value), seq])\n        return seq  \n        \n    def __getitem__(self, index):\n        X = self.X_list[index]\n        y = self.y_list[index]\n\n        # ---------- (A)  Augmentation ----------\n        if self.mode == \"train\" and self.augment is not None:\n            X = self.augment(X, self.imu_dim)     \n\n        X = self.pad_sequences_torch(X, self.maxlen, 'pre', 'pre')\n        return X, y\n    \n    def __len__(self):\n        return len(self.X_list)\n\n\nclass EarlyStopping:\n    \"\"\"Early stopping utility\"\"\"\n    def __init__(self, patience=7, min_delta=0, restore_best_weights=True):\n        self.patience = patience\n        self.min_delta = min_delta\n        self.restore_best_weights = restore_best_weights\n        self.best_loss = None\n        self.counter = 0\n        self.best_weights = None\n        \n    def __call__(self, val_loss, model):\n        if self.best_loss is None:\n            self.best_loss = val_loss\n            self.save_checkpoint(model)\n        elif val_loss < self.best_loss - self.min_delta:\n            self.best_loss = val_loss\n            self.counter = 0\n            self.save_checkpoint(model)\n        else:\n            self.counter += 1\n            \n        if self.counter >= self.patience:\n            if self.restore_best_weights:\n                model.load_state_dict(self.best_weights)\n            return True\n        return False\n    \n    def save_checkpoint(self, model):\n        self.best_weights = model.state_dict().copy()\n\nclass EMA:\n    def __init__(self, model, decay=0.999):\n        self.decay = decay\n        self.shadow = {}\n        self.backup = {}\n\n        for name, param in model.named_parameters():\n            if param.requires_grad:\n                self.shadow[name] = param.data.clone()\n\n    def update(self, model):\n        for name, param in model.named_parameters():\n            if param.requires_grad:\n                assert name in self.shadow\n                new_average = (1.0 - self.decay) * param.data + self.decay * self.shadow[name]\n                self.shadow[name] = new_average.clone()\n\n    def apply_shadow(self, model):\n        self.backup = {}\n        for name, param in model.named_parameters():\n            if param.requires_grad:\n                self.backup[name] = param.data.clone()\n                param.data = self.shadow[name]\n\n    def restore(self, model):\n        for name, param in model.named_parameters():\n            if param.requires_grad and name in self.backup:\n                param.data = self.backup[name]\n        self.backup = {}\n\ndef set_seed(seed: int = 42):\n    random.seed(seed)\n\n    os.environ['PYTHONHASHSEED'] = str(seed)\n\n    np.random.seed(seed)\n\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed) \n    # torch.backends.cudnn.deterministic = True\n    # torch.backends.cudnn.benchmark = False\n    # torch.use_deterministic_algorithms(True)\n\nclass Augment:\n    def __init__(self,\n                 p_jitter=0.8, sigma=0.02, scale_range=[0.9,1.1],\n                 p_dropout=0.3,\n                 p_moda=0.5,          \n                 drift_std=0.005,     \n                 drift_max=0.25):      \n        self.p_jitter  = p_jitter\n        self.sigma     = sigma\n        self.scale_min, self.scale_max = scale_range\n        self.p_dropout = p_dropout\n        self.p_moda    = p_moda\n        self.drift_std = drift_std\n        self.drift_max = drift_max\n\n    # ---------- Jitter & Scaling ----------\n    def jitter_scale(self, x: np.ndarray) -> np.ndarray:\n        noise  = np.random.randn(*x.shape) * self.sigma\n        scale  = np.random.uniform(self.scale_min,\n                                   self.scale_max,\n                                   size=(1, x.shape[1]))\n        return (x + noise) * scale\n\n    # ---------- Sensor Drop-out ----------\n    def sensor_dropout(self,\n                       x: np.ndarray,\n                       imu_dim: int) -> np.ndarray:\n\n        if random.random() < self.p_dropout:\n            x[:, imu_dim:] = 0.0\n        return x\n\n    def motion_drift(self, x: np.ndarray, imu_dim: int) -> np.ndarray:\n\n        T = x.shape[0]\n\n        drift = np.cumsum(\n            np.random.normal(scale=self.drift_std, size=(T, 1)),\n            axis=0\n        )\n        drift = np.clip(drift, -self.drift_max, self.drift_max)   \n\n        x[:, :6] += drift\n\n        if imu_dim > 6:\n            x[:, 6:imu_dim] += drift     \n        return x\n    \n    def random_erasing(self, x: np.ndarray, p_erase=0.5, erase_ratio=0.1) -> np.ndarray:\n        if random.random() < p_erase:\n            T, D = x.shape\n            erase_size = int(T * erase_ratio)\n            start_idx = random.randint(0, T - erase_size)\n            x[start_idx:start_idx + erase_size, :] = 0.0\n        return x\n\n    # ---------- master call ----------\n    def __call__(self,\n                 x: np.ndarray,\n                 imu_dim: int) -> np.ndarray:\n        if random.random() < self.p_jitter:\n            x = self.jitter_scale(x)\n\n        if random.random() < self.p_moda:\n            x = self.motion_drift(x, imu_dim)\n\n        x = self.sensor_dropout(x, imu_dim)\n\n        # Apply random erasing\n        x = self.random_erasing(x)\n\n        return x\n# ================================\nif TRAIN:\n    print(\"▶ TRAIN MODE – loading dataset …\")\n    df = pd.read_csv(RAW_DIR / \"train.csv\")\n\n    # Label encoding\n    le = LabelEncoder()\n    df['gesture_int'] = le.fit_transform(df['gesture'])\n    np.save(EXPORT_DIR / \"gesture_classes.npy\", le.classes_)\n\n    # Feature list\n    meta_cols = {'gesture', 'gesture_int', 'sequence_type', 'behavior', 'orientation',\n                 'row_id', 'subject', 'phase', 'sequence_id', 'sequence_counter'}\n    feature_cols = [c for c in df.columns if c not in meta_cols]\n\n    imu_cols = [c for c in feature_cols if not (c.startswith('thm_') or c.startswith('tof_'))]\n    tof_cols = [c for c in feature_cols if c.startswith('tof_')]\n    thm_cols = [c for c in feature_cols if c.startswith('thm_')]\n    print(f\"  IMU {len(imu_cols)} | TOF {len(tof_cols)} | THM {len(thm_cols)} | total {len(feature_cols)} features\")\n\n    # Global scaler\n    scaler = StandardScaler().fit(df[feature_cols].ffill().bfill().fillna(0).values)\n    joblib.dump(scaler, EXPORT_DIR / \"scaler.pkl\")\n\n    # Build sequences\n    seq_gp = df.groupby('sequence_id')\n    X_list, y_list, id_list = [], [], []\n    lens = []  # 添加这行来收集长度信息\n    for seq_id, seq in seq_gp:\n        mat = preprocess_sequence(seq, feature_cols, scaler)\n        X_list.append(mat)\n        y_list.append(seq['gesture_int'].iloc[0])\n        id_list.append(seq_id)\n        lens.append(len(mat))  # 收集每个序列的长度\n    \n    pad_len = PAD_PERCENTILE#int(np.percentile(lens, PAD_PERCENTILE))\n    print(pad_len)\n    np.save(EXPORT_DIR / \"sequence_maxlen.npy\", pad_len)\n    np.save(EXPORT_DIR / \"feature_cols.npy\", np.array(feature_cols))\n    id_list = np.array(id_list)\n    \n    # 修改这里：先使用pad_sequences_torch处理X_list，然后再进行TOF统计特征提取\n    X_list_padded = pad_sequences_torch(X_list, maxlen=pad_len, padding='pre', truncating='pre')\n    \n    # 在这里添加TOF统计特征\n    original_tof_dim = len(tof_cols)\n    \n    # 提取TOF统计特征\n    tof_data = X_list_padded[:, :, len(imu_cols):len(imu_cols)+len(tof_cols)]  # (N, T, 320)\n    tof_stats = extract_tof_statistics(tof_data, tof_cols)  # (N, T, 84)\n    \n    # 替换原始TOF数据为统计特征\n    other_data = np.concatenate([\n        X_list_padded[:, :, :len(imu_cols)],  # IMU数据\n        X_list_padded[:, :, len(imu_cols)+len(tof_cols):]  # THM数据\n    ], axis=2)  # (N, T, D-len(tof_cols))\n    \n    # 组合新特征：IMU + TOF统计 + THM\n    X_list_enhanced = np.concatenate([other_data, tof_stats], axis=2)\n    \n    X_list_all = X_list_enhanced  # 已经是填充后的数组\n    y_list_all = np.eye(len(le.classes_))[y_list].astype(np.float32)  # One-hot encoding\n\n    augmenter = Augment(\n        p_jitter=0.9844818619033621, sigma=0.03291295776089293, scale_range=(0.7542342630597011,1.1625052821731077),\n        p_dropout=0.41782786013520684,\n        p_moda=0.3910622476959722, drift_std=0.0040285239353308015, drift_max=0.3929358950258158    \n    )  \nEPOCHS = 145\nif TRAIN:\n    # Split\n    skf = StratifiedKFold(n_splits=FOLDS, shuffle=True, random_state=random_state)\n    models = []\n    \n    for fold, (train_idx, val_idx) in enumerate(skf.split(id_list, np.argmax(y_list_all, axis=1))):\n\n        train_list= X_list_all[train_idx]\n        train_y_list= y_list_all[train_idx]\n        val_list = X_list_all[val_idx]\n        val_y_list= y_list_all[val_idx]\n\n        \n        # Data loaders\n        train_dataset = CMI3Dataset(train_list, train_y_list, maxlen, mode=\"train\", imu_dim=len(imu_cols),\n                                augment=augmenter)\n        train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=0, drop_last=True)\n    \n        val_dataset = CMI3Dataset(val_list, val_y_list, maxlen, mode=\"val\")\n        val_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=0, drop_last=True)\n\n    \n        # Model\n        model = MultiModalModel(\n            imu_dim=len(imu_cols),\n            tof_dim=84,  # 更新为新的TOF统计特征维度\n            thm_dim=len(thm_cols),  # Now using actual THM data\n            d_model=64,\n            n_classes=len(le.classes_)\n        ).to(device)\n        ema = EMA(model, decay=0.999)\n        # Optimizer and scheduler\n        optimizer = Adam(model.parameters(), lr=LR_INIT, weight_decay=WD)\n        \n        steps_per_epoch = len(train_loader)\n        nbatch = len(train_loader)\n        warmup = epochs_warmup * nbatch\n        nsteps = EPOCHS * nbatch\n        scheduler = CosineLRScheduler(optimizer,\n                  warmup_t=warmup, warmup_lr_init=warmup_lr_init, warmup_prefix=True,\n                  t_initial=(nsteps - warmup), lr_min=lr_min) \n    \n        early_stopping = EarlyStopping(patience=PATIENCE, restore_best_weights=True)\n    \n        train_loss = 0.0\n        train_acc = None\n        val_acc = None\n        val_best_acc = 0.0\n        i_scheduler = 0\n        val_loss = 0.0\n        \n        # Training loop\n        print(\"▶ Starting training...\")\n        for epoch in range(EPOCHS):\n            model.train()\n            train_preds = []\n            train_targets = []\n            train_loss = 0.0\n            with tqdm(train_loader, desc=f\"Epoch {epoch+1}/{EPOCHS}\") as t:\n                for X, y in t:\n                    X, y = X.float().to(device),y.float().to(device)\n                    optimizer.zero_grad()\n                    \n                    # Ensure input data is correctly shaped\n                    B, T, D = X.shape\n                    imu = X[:, :, :len(imu_cols)].permute(0, 2, 1)  # [B, C, T]\n                    # 修正TOF和THM数据的索引方式，确保TOF只使用统计特征\n                    tof_start_idx = len(imu_cols)\n                    tof_end_idx = len(imu_cols) + 84  # TOF统计特征维度为84\n                    tof = X[:, :, tof_start_idx:tof_end_idx].permute(0, 2, 1)  # TOF统计特征 (84 channels)\n                    thm = X[:, :, -len(thm_cols):].permute(0, 2, 1)  # THM数据\n                    \n                    logits = model(imu, tof, thm)\n\n                    loss = -torch.sum(F.log_softmax(logits, dim=1) * y, dim=1).mean()\n                    loss.backward()\n                    optimizer.step()\n                    ema.update(model)\n                    train_preds.extend(logits.argmax(dim=1).cpu().numpy())\n                    train_targets.extend(y.argmax(dim=1).cpu().numpy())\n                    scheduler.step(i_scheduler)\n                    i_scheduler +=1\n\n                    train_loss += loss.item()\n                    t.set_postfix(loss=train_loss/len(train_loader))\n\n            # Validation phase\n            model.eval()\n            val_loss = 0.0\n            val_preds = []\n            val_targets = []\n            with torch.no_grad():\n                with tqdm(val_loader, desc=\"Validation\") as t:\n                    for X, y in t:\n                        X, y = X.float().to(device),y.float().to(device)\n                        \n                        # Ensure input data is correctly shaped\n                        B, T, D = X.shape\n                        imu = X[:, :, :len(imu_cols)].permute(0, 2, 1)  # [B, C, T]\n                        # 修正TOF和THM数据的索引方式，确保TOF只使用统计特征\n                        tof_start_idx = len(imu_cols)\n                        tof_end_idx = len(imu_cols) + 84  # TOF统计特征维度为84\n                        tof = X[:, :, tof_start_idx:tof_end_idx].permute(0, 2, 1)  # TOF统计特征 (84 channels)\n                        thm = X[:, :, -len(thm_cols):].permute(0, 2, 1)  # THM数据\n                        \n                        logits = model(imu, tof, thm)\n                        loss = -torch.sum(F.log_softmax(logits, dim=1) * y, dim=1).mean()  # 修改为与训练一致的损失计算\n                        val_loss += loss.item()\n                        val_preds.extend(logits.argmax(dim=1).cpu().numpy())\n                        val_targets.extend(y.argmax(dim=1).cpu().numpy())\n                        t.set_postfix(loss=val_loss/len(val_loader))\n\n            val_loss /= len(val_loader)  # 移到循环外\n\n            # Convert numeric predictions to class labels for competition metric\n            train_preds_labels = [le.classes_[pred] for pred in train_preds]\n            train_targets_labels = [le.classes_[target] for target in train_targets]\n            val_preds_labels = [le.classes_[pred] for pred in val_preds]\n            val_targets_labels = [le.classes_[target] for target in val_targets]\n            \n            # Calculate competition scores using only kaggle metric\n            train_f1 = get_competition_score(train_targets_labels, train_preds_labels)\n            val_f1 = get_competition_score(val_targets_labels, val_preds_labels)\n\n            print(f\"Epoch {epoch+1}/{EPOCHS} - Loss: {train_loss/len(train_loader):.4f} - Train F1: {train_f1:.4f} - Val Loss: {val_loss:.4f} - Val F1: {val_f1:.4f}\")\n\n        model.eval()\n        with torch.inference_mode():\n            val_preds = []\n            val_targets = []\n            for X, y in (val_loader):  \n                half = BATCH_SIZE // 2         \n\n                x_front = X[:half]               \n                x_back  = X[half:].clone()      \n                \n                x_back[:, :, 7:] = 0.0    \n                X = torch.cat([x_front, x_back], dim=0)  # (B, C, T)\n                X, y = X.float().to(device), y.to(device)\n                \n                # Ensure input data is correctly shaped\n                B, T, D = X.shape\n                imu = X[:, :, :len(imu_cols)].permute(0, 2, 1)  # [B, C, T]\n                # 修正TOF和THM数据的索引方式，确保TOF只使用统计特征\n                tof_start_idx = len(imu_cols)\n                tof_end_idx = len(imu_cols) + 84  # TOF统计特征维度为84\n                tof = X[:, :, tof_start_idx:tof_end_idx].permute(0, 2, 1)  # TOF统计特征 (84 channels)\n                thm = X[:, :, -len(thm_cols):].permute(0, 2, 1)  # THM数据\n                \n                logits = model(imu, tof, thm)\n                val_preds.extend(logits.argmax(dim=1).cpu().numpy())\n                val_targets.extend(y.argmax(dim=1).cpu().numpy())\n                \n                loss = F.cross_entropy(logits, y)\n                val_loss += loss.item()\n    \n            train_loss /= len(train_loader)\n            val_loss /= len(val_loader)\n        models.append(model)\n        # Save model\n        torch.save({\n            'model_state_dict': model.state_dict(),\n            'imu_dim': len(imu_cols),\n            'tof_dim': len(tof_cols),\n            'thm_dim': len(thm_cols),  # Save THM dimension\n            'n_classes': len(le.classes_),\n            'pad_len': pad_len\n        }, EXPORT_DIR / f\"gesture_two_branch_fold{fold}.pth\")\n        # 计算验证指标\n        val_preds = [le.classes_[pred] for pred in val_preds]\n        val_targets = [le.classes_[target] for target in val_targets]\n        val_f1 = get_competition_score(val_targets, val_preds)\n        \n        # Print validation accuracy and F1 score for the current fold\n        print(f\"fold: {fold} val_f1: {val_f1:.4f}\")\n        print(f\"✔ Training done - Final F1 Score: {val_f1:.4f} - artefacts saved in {EXPORT_DIR}\")\n\nelse:\n    print(\"▶ INFERENCE MODE – loading artefacts from\", PRETRAINED_DIR)\n    feature_cols = np.load(PRETRAINED_DIR / \"feature_cols.npy\", allow_pickle=True).tolist()\n    pad_len = int(np.load(PRETRAINED_DIR / \"sequence_maxlen.npy\"))\n    scaler = joblib.load(PRETRAINED_DIR / \"scaler.pkl\")\n    gesture_classes = np.load(PRETRAINED_DIR / \"gesture_classes.npy\", allow_pickle=True)\n\n    imu_cols = [c for c in feature_cols if not (c.startswith('thm_') or c.startswith('tof_'))]\n    tof_cols = [c for c in feature_cols if c.startswith('tof_')]\n    thm_cols = [c for c in feature_cols if c.startswith('thm_')]  # Load THM columns\n\n    MODELS = [f'gesture_two_branch_fold{i}.pth' for i in range(5)]\n    models = []\n    for path in MODELS:\n        checkpoint = torch.load(PRETRAINED_DIR / path, map_location=device)\n        model = MultiModalModel(\n            checkpoint['imu_dim'],\n            checkpoint['tof_dim'],\n            checkpoint['thm_dim'],  # Load THM dimension\n            d_model=64,\n            n_classes=checkpoint['n_classes']\n        ).to(device)\n        model.load_state_dict(checkpoint['model_state_dict'])\n        model.eval()\n        models.append(model)\n\n    print(\"  model, scaler, pads loaded – ready for evaluation\")\n\n# Make sure gesture_classes exists in both modes\nif TRAIN:\n    gesture_classes = le.classes_\n\n# ================================\n# predict 函数（只用 time_shift）\n# ================================\ndef predict(sequence: pl.DataFrame, demographics: pl.DataFrame, n_tta=8) -> str:\n    global gesture_classes\n    if gesture_classes is None:\n        gesture_classes = np.load(\n            PRETRAINED_DIR / \"gesture_classes.npy\",\n            allow_pickle=True\n        )\n\n    # 1. 原始预处理\n    df_seq = sequence.to_pandas()\n    mat = preprocess_sequence(df_seq, feature_cols, scaler)\n\n    # 2. 仅使用 tta_jitter\n    tta_samples = []\n    for _ in range(n_tta):\n        x_aug = tta_jitter(mat.copy(), sigma=0.01)\n        tta_samples.append(x_aug)\n\n    # 3. Pad 并堆叠成 Tensor\n    pads = pad_sequences_torch(\n        tta_samples,\n        maxlen=pad_len,\n        padding='pre',\n        truncating='pre'\n    )\n    X = torch.FloatTensor(pads).to(device)  # (n_tta, C, T)\n\n    # 4. 多模型 ensemble + TTA 预测\n    with torch.no_grad():\n        probs_accum = []\n        for m in models:\n            # Ensure input data is correctly shaped\n            imu = X[:, :, :len(imu_cols)].permute(0, 2, 1)  # [B, C, T]\n            # 修正TOF和THM数据的索引方式，确保TOF只使用统计特征\n            tof_start_idx = len(imu_cols)\n            tof_end_idx = len(imu_cols) + 84  # TOF统计特征维度为84\n            tof = X[:, :, tof_start_idx:tof_end_idx].permute(0, 2, 1)  # TOF统计特征 (84 channels)\n            thm = X[:, :, -len(thm_cols):].permute(0, 2, 1)  # THM数据\n            \n            logits = m(imu, tof, thm)               # (n_tta, n_classes)\n            probs  = torch.softmax(logits, dim=1)\n            probs_accum.append(probs)\n        # 按模型平均 → (n_tta, n_classes)\n        avg_per_tta = torch.stack(probs_accum).mean(0)\n        # 对所有 TTA 样本平均 → (n_classes,)\n        final_prob = avg_per_tta.mean(0)\n        idx = int(final_prob.argmax().cpu().numpy())\n\n    return str(gesture_classes[idx])\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-08-09T04:39:00.939282Z","iopub.execute_input":"2025-08-09T04:39:00.939563Z"}},"outputs":[],"execution_count":null}]}