{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":5048,"databundleVersionId":868335}],"dockerImageVersionId":31329,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport cv2\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom sklearn.model_selection import GroupKFold\nimport timm\nfrom tqdm import tqdm\nimport gc\n\n# 基础配置 (Config)\nclass Config:\n    DATA_DIR = '/kaggle/input/competitions/state-farm-distracted-driver-detection'\n    CSV_PATH = os.path.join(DATA_DIR, 'driver_imgs_list.csv')\n    TRAIN_DIR = os.path.join(DATA_DIR, 'imgs/train')\n    \n    MODEL_NAME = 'tf_efficientnet_b3.ns_jft_in1k'\n    NUM_CLASSES = 10\n    N_FOLDS = 5\n    DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# 数据集类 (Dataset)\nclass DistractedDriverDataset(Dataset):\n    def __init__(self, df, transforms=None):\n        self.df = df\n        self.transforms = transforms\n        \n    def __len__(self):\n        return len(self.df)\n        \n    def __getitem__(self, idx):\n        row = self.df.iloc[idx]\n        img_path = os.path.join(Config.TRAIN_DIR, row['classname'], row['img'])\n        \n        image = cv2.imread(img_path)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        if self.transforms:\n            augmented = self.transforms(image=image)\n            image = augmented['image']\n            \n        label = int(row['classname'][1:])\n        return image, torch.tensor(label, dtype=torch.long)\n\n# 动态数据增强 (接收动态尺寸)\ndef get_transforms(phase, img_size):\n    if phase == 'train':\n        return A.Compose([\n            A.Resize(img_size, img_size), # 尺寸由外部动态传入\n            A.Affine(scale=(0.9, 1.1), translate_percent=(-0.05, 0.05), rotate=(-15, 15), p=0.5),\n            A.RandomBrightnessContrast(p=0.5),\n            A.CoarseDropout(num_holes_range=(1, 8), hole_height_range=(8, 32), hole_width_range=(8, 32), p=0.5), \n            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n            ToTensorV2(),\n        ])\n    else:\n        return A.Compose([\n            A.Resize(img_size, img_size),\n            A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n            ToTensorV2(),\n        ])\n\n# 模型构建 (Model)\nclass DriverModel(nn.Module):\n    def __init__(self, model_name=Config.MODEL_NAME, pretrained=True):\n        super().__init__()\n        self.model = timm.create_model(model_name, pretrained=pretrained, num_classes=Config.NUM_CLASSES)\n        \n    def forward(self, x):\n        return self.model(x)\n\n# 核心：模块化的训练函数\ndef train_phase(fold, df, phase_name, img_size, batch_size, epochs, lr, load_weight_path=None, save_weight_path=None):\n    print(f\"\\n🚀 开始 {phase_name} | 尺寸: {img_size}x{img_size} | Epochs: {epochs} | Batch: {batch_size}\")\n    \n    train_df = df[df['fold'] != fold].reset_index(drop=True)\n    valid_df = df[df['fold'] == fold].reset_index(drop=True)\n    \n    # 动态传入当前阶段的 img_size\n    train_dataset = DistractedDriverDataset(train_df, transforms=get_transforms('train', img_size))\n    valid_dataset = DistractedDriverDataset(valid_df, transforms=get_transforms('valid', img_size))\n    \n    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4, pin_memory=True, persistent_workers=True)\n    valid_loader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=False, num_workers=4, pin_memory=True, persistent_workers=True)\n    \n    # 如果是阶段2，就不需要再下载预训练权重了，直接加载阶段1的权重\n    pretrained_imagenet = True if load_weight_path is None else False\n    model = DriverModel(pretrained=pretrained_imagenet).to(Config.DEVICE)\n    \n    if load_weight_path:\n        print(f\"📦 载入上阶段权重: {load_weight_path}\")\n        model.load_state_dict(torch.load(load_weight_path, map_location=Config.DEVICE))\n        \n    criterion = nn.CrossEntropyLoss(label_smoothing=0.1)\n    optimizer = torch.optim.AdamW(model.parameters(), lr=lr)\n    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=epochs)\n    scaler = torch.cuda.amp.GradScaler()\n    \n    best_loss = float('inf')\n    \n    for epoch in range(epochs):\n        model.train()\n        train_loss = 0.0\n        \n        for images, labels in tqdm(train_loader, desc=f\"[{phase_name}] Epoch {epoch+1} Train\"):\n            images, labels = images.to(Config.DEVICE), labels.to(Config.DEVICE)\n            optimizer.zero_grad()\n            \n            with torch.cuda.amp.autocast():\n                outputs = model(images)\n                loss = criterion(outputs, labels)\n            \n            scaler.scale(loss).backward()\n            scaler.step(optimizer)\n            scaler.update()\n            \n            train_loss += loss.item() * images.size(0)\n            \n        scheduler.step()\n        train_loss = train_loss / len(train_loader.dataset)\n        \n        model.eval()\n        valid_loss = 0.0\n        correct = 0\n        with torch.no_grad():\n            for images, labels in tqdm(valid_loader, desc=f\"[{phase_name}] Epoch {epoch+1} Valid\"):\n                images, labels = images.to(Config.DEVICE), labels.to(Config.DEVICE)\n                \n                with torch.cuda.amp.autocast():\n                    outputs = model(images)\n                    loss = criterion(outputs, labels)\n                \n                valid_loss += loss.item() * images.size(0)\n                preds = outputs.argmax(dim=1)\n                correct += (preds == labels).sum().item()\n                \n        valid_loss = valid_loss / len(valid_loader.dataset)\n        valid_acc = correct / len(valid_loader.dataset)\n        \n        print(f\"Epoch {epoch+1}: Train Loss = {train_loss:.4f}, Valid Loss = {valid_loss:.4f}, Valid Acc = {valid_acc:.4f}\")\n        \n        if valid_loss < best_loss:\n            best_loss = valid_loss\n            if save_weight_path:\n                torch.save(model.state_dict(), save_weight_path)\n                print(f\"--> Saved Best Model for {phase_name}!\")\n                \n    # 释放显存\n    del model, optimizer, train_loader, valid_loader\n    gc.collect()\n    torch.cuda.empty_cache()\n\n# 执行入口：5折全量渐进式流水线\nif __name__ == \"__main__\":\n    df = pd.read_csv(Config.CSV_PATH)\n    \n    gkf = GroupKFold(n_splits=Config.N_FOLDS)\n    df['fold'] = -1\n    for fold, (train_idx, val_idx) in enumerate(gkf.split(df, groups=df['subject'])):\n        df.loc[val_idx, 'fold'] = fold\n        \n    print(f\"Total Data: {len(df)}\")\n    \n    for fold in range(Config.N_FOLDS):\n        print(f\"开始训练第 {fold + 1} 折\")\n        \n        # 每一折独立的中间权重和最终权重路径\n        phase1_weight = f\"phase1_model_fold{fold}.pth\"\n        final_weight = f\"best_model_fold{fold}.pth\"\n        \n        # --- Phase 1: 小图速成 (128x128) ---\n        train_phase(\n            fold=fold, \n            df=df, \n            phase_name=f\"Fold {fold} - Phase 1\", \n            img_size=128,         \n            batch_size=128,       \n            epochs=5,             \n            lr=1e-3,              \n            load_weight_path=None, \n            save_weight_path=phase1_weight\n        )\n        \n        # --- Phase 2: 大图微调 (224x224) ---\n        train_phase(\n            fold=fold, \n            df=df, \n            phase_name=f\"Fold {fold} - Phase 2\", \n            img_size=384,         \n            batch_size=32,        \n            epochs=5,             \n            lr=2e-4,              \n            load_weight_path=phase1_weight, \n            save_weight_path=final_weight\n        )\n        \n        # 每一折结束后，强制清理临时文件（删除 phase1 权重，只保留最终的 final 权重，省下硬盘空间）\n        if os.path.exists(phase1_weight):\n            os.remove(phase1_weight)\n            \n        print(f\"🟢 Fold {fold} 渐进式训练结束！最终模型已保存为: {final_weight}\\n\")\n        \n    print(\"\\n🎉 恭喜！5折全量渐进式训练全部完成！\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-05-11T08:21:58.999201Z","iopub.execute_input":"2026-05-11T08:21:59.000199Z","iopub.status.idle":"2026-05-11T08:25:45.472659Z","shell.execute_reply.started":"2026-05-11T08:21:59.000150Z","shell.execute_reply":"2026-05-11T08:25:45.471590Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport cv2\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nimport albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nimport gc\nfrom tqdm import tqdm\n\n# ==========================================\n# 1. 定义测试集 Dataset 类\n# ==========================================\nclass DistractedDriverTestDataset(Dataset):\n    def __init__(self, img_names, transforms=None):\n        self.img_names = img_names\n        self.transforms = transforms\n        \n    def __len__(self):\n        return len(self.img_names)\n        \n    def __getitem__(self, idx):\n        row_img_name = self.img_names[idx]\n        img_path = os.path.join(Config.TEST_DIR, row_img_name)\n        \n        image = cv2.imread(img_path)\n        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)\n        \n        if self.transforms:\n            augmented = self.transforms(image=image)\n            image = augmented['image']\n            \n        return row_img_name, image\n\nConfig.IMG_SIZE = 384\nConfig.TEST_DIR = os.path.join(Config.DATA_DIR, 'imgs/test')\n\ntest_transforms = A.Compose([\n    A.Resize(Config.IMG_SIZE, Config.IMG_SIZE),\n    A.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),\n    ToTensorV2(),\n])\n\n# ==========================================\n# 2. 5 折集成 + TTA 推理核心逻辑\n# ==========================================\ntest_imgs = sorted(os.listdir(Config.TEST_DIR))\nprint(f\"找到测试集图片共: {len(test_imgs)} 张\")\n\ntest_dataset = DistractedDriverTestDataset(test_imgs, transforms=test_transforms)\ntest_loader = DataLoader(test_dataset, batch_size=128, shuffle=False, num_workers=4, pin_memory=True)\n\nall_folds_preds = []\nimg_names_list = []\n\nfor fold in range(Config.N_FOLDS):\n    # 【核心修正2】必须对应渐进式脚本最终存下来的 5 个 final 文件名\n    model_path = f\"best_model_fold{fold}.pth\"         \n    print(f\"--> Loading {model_path} for Inference...\")\n    model = DriverModel().to(Config.DEVICE)\n    model.load_state_dict(torch.load(model_path, map_location=Config.DEVICE))\n    model.eval()\n    \n    fold_preds = []\n    with torch.no_grad():\n        for img_names, images in tqdm(test_loader, desc=f\"Fold {fold} predicting\"):\n            images = images.to(Config.DEVICE)\n            \n            # --- 融入 TTA (测试时增强) 机制 ---\n            # 1. 预测原始分辨率和亮度的图像\n            with torch.cuda.amp.autocast():\n                outputs_orig = model(images)\n            probs_orig = torch.softmax(outputs_orig, dim=1)\n            \n            # 2. 变换 TTA：稍微增亮一点点（模拟不同车舱光照条件）\n            images_bright = torch.clamp(images + 0.08, 0.0, 1.0)\n            with torch.cuda.amp.autocast():\n                outputs_bright = model(images_bright)\n            probs_bright = torch.softmax(outputs_bright, dim=1)\n            \n            # 3. 双重增强结果取平均（无痛白嫖涨分点）\n            probs_final = ((probs_orig + probs_bright) / 2.0).cpu().numpy()\n            fold_preds.append(probs_final)\n            \n            if fold == 0:\n                img_names_list.extend(img_names)\n                \n    all_folds_preds.append(np.vstack(fold_preds))\n    \n    del model\n    gc.collect()\n    torch.cuda.empty_cache()\n\n# 如果成功跑完对应的折，进行均值融合\nif len(all_folds_preds) > 0:\n    ensemble_preds = np.mean(all_folds_preds, axis=0)\n\n    # 生成最终的集成提交文件\n    sub_df = pd.DataFrame(ensemble_preds, columns=[f'c{i}' for i in range(10)])\n    sub_df.insert(0, 'img', img_names_list)\n    sub_df.to_csv('submission.csv', index=False)\n    print(\"🎉 渐进式5折模型已无缝结合 TTA，顺利生成最优 submission.csv！\")\nelse:\n    print(\"❌ 错误：在当前工作目录下未找到任何折的模型权重文件，请检查模型是否训练成功。\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-11T08:26:25.284258Z","iopub.execute_input":"2026-05-11T08:26:25.284806Z","iopub.status.idle":"2026-05-11T08:32:44.578075Z","shell.execute_reply.started":"2026-05-11T08:26:25.284765Z","shell.execute_reply":"2026-05-11T08:32:44.576660Z"}},"outputs":[],"execution_count":null}]}