{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":14774,"databundleVersionId":875431,"sourceType":"competition"},{"sourceId":517629,"sourceType":"modelInstanceVersion","modelInstanceId":408055,"modelId":425923}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import cv2\nimport numpy as np\nimport os\nimport pandas as pd\nimport random\nimport timm\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom albumentations import Compose, Normalize, HorizontalFlip, Rotate, ColorJitter\nfrom albumentations.pytorch import ToTensorV2\nfrom sklearn.metrics import cohen_kappa_score\nfrom sklearn.model_selection import train_test_split\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau, CosineAnnealingLR\nfrom torch.utils.data import DataLoader, Dataset\nfrom tqdm import tqdm\nfrom scipy.optimize import minimize\n\n# ==============================================================================\n# 0. 全局设置 (Global Settings)\n# ==============================================================================\ndef set_seed(seed=42):\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    if torch.cuda.is_available():\n        torch.cuda.manual_seed_all(seed)\n        torch.backends.cudnn.deterministic = True\n        torch.backends.cudnn.benchmark = False\nset_seed(42)\n\n# ==============================================================================\n# 1. 预处理函数 (Preprocessing Functions)\n# ==============================================================================\ndef circle_crop(image):\n    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)\n    _, thresh = cv2.threshold(gray, 10, 255, cv2.THRESH_BINARY)\n    contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)\n    if not contours:\n        return image\n    cnt = max(contours, key=cv2.contourArea)\n    x, y, w, h = cv2.boundingRect(cnt)\n    cropped_image = image[y:y + h, x:x + w]\n    return cropped_image\n\ndef apply_ben_graham_preprocessing(image, sigmaX=30):\n    blurred_image = cv2.GaussianBlur(image, (0, 0), sigmaX)\n    processed_image = cv2.addWeighted(image, 4, blurred_image, -4, 128)\n    return processed_image\n\ndef preprocess_directory(input_dir, output_dir, image_size):\n    if not os.path.exists(output_dir):\n        os.makedirs(output_dir)\n    image_files = [f for f in os.listdir(input_dir) if f.lower().endswith(('.png', '.jpg', '.jpeg'))]\n    for filename in tqdm(image_files, desc=f\"Preprocessing {os.path.basename(input_dir)}\"):\n        input_path = os.path.join(input_dir, filename)\n        output_path = os.path.join(output_dir, filename)\n        image = cv2.imread(input_path)\n        if image is None:\n            continue\n        processed_image = circle_crop(image)\n        processed_image = cv2.resize(processed_image, (image_size, image_size))\n        processed_image = apply_ben_graham_preprocessing(processed_image)\n        cv2.imwrite(output_path, processed_image)\n\n# ==============================================================================\n# 2. PyTorch 模型与数据类 (Model and Data Classes)\n# ==============================================================================\ndef replace_batchnorm_with_groupnorm(module, num_groups=32):\n    for name, child in module.named_children():\n        if isinstance(child, nn.BatchNorm2d):\n            num_channels = child.num_features\n            if num_channels % num_groups == 0:\n                setattr(module, name, nn.GroupNorm(num_groups=num_groups, num_channels=num_channels))\n        else:\n            replace_batchnorm_with_groupnorm(child, num_groups)\n\nclass EfficientNetModel(nn.Module):\n    def __init__(self, pretrained_model_path=None):\n        super().__init__()\n        self.model = timm.create_model('efficientnet_b6', pretrained=False)\n        replace_batchnorm_with_groupnorm(self.model)\n        in_features = self.model.classifier.in_features\n        self.model.classifier = nn.Sequential(\n            nn.Dropout(p=0.5),\n            nn.Linear(in_features, 512),\n            nn.Mish(),\n            nn.Dropout(p=0.5),\n            nn.Linear(512, 1)\n        )\n        if pretrained_model_path and os.path.exists(pretrained_model_path):\n            self.model.load_state_dict(torch.load(pretrained_model_path), strict=False)\n        elif pretrained_model_path:\n            print(f\"警告: 预训练模型路径未找到: {pretrained_model_path}\")\n\n    def forward(self, x):\n        return self.model(x)\n\nclass CustomImageDataset(Dataset):\n    def __init__(self, img_dir, labels_df, transform=None):\n        self.img_dir = img_dir\n        self.labels_df = labels_df\n        self.transform = transform\n    def __len__(self):\n        return len(self.labels_df)\n    def __getitem__(self, idx):\n        img_id = self.labels_df.iloc[idx, 0]\n        label = float(self.labels_df.iloc[idx, 1])\n        img_path = os.path.join(self.img_dir, img_id + '.png')\n        try:\n            img = cv2.imread(img_path)\n            img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n            if self.transform:\n                img = self.transform(image=img)[\"image\"]\n            return img, label\n        except Exception as e:\n            return self.__getitem__((idx + 1) % len(self))\n\ndef get_transforms(data_type='train'):\n    if data_type == 'train':\n        return Compose([\n            Rotate(limit=40, p=0.5),\n            ColorJitter(brightness=0.1, contrast=0.1, saturation=0.1, hue=0.1, p=0.3),\n            HorizontalFlip(p=0.5),\n            Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n            ToTensorV2(),\n        ])\n    else:\n        return Compose([\n            Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),\n            ToTensorV2(),\n        ])\n\n# ==============================================================================\n# 3. 训练与预测函数 (Training and Prediction Functions)\n# ==============================================================================\ndef train(model, train_loader, valid_loader, device, num_epochs, lr, image_size):\n    criterion = nn.MSELoss()\n    optimizer = optim.AdamW(model.parameters(), lr=lr, weight_decay=1e-5)\n    scheduler = CosineAnnealingLR(optimizer, T_max=num_epochs, eta_min=1e-6)\n\n    best_valid_kappa = 0.0\n    epochs_no_improve = 0\n    patience = 5\n    save_path = os.path.join(os.getcwd(), f'best_model_kappa_{image_size}.pth')\n\n    for epoch in range(num_epochs):\n        model.train()\n        running_loss = 0.0\n        \n        for images, labels in tqdm(train_loader, desc=\"训练中\"):\n            images = images.to(device)\n            labels = labels.to(device).float().view(-1, 1)\n            optimizer.zero_grad()\n            outputs = model(images)\n            loss = criterion(outputs, labels)\n            loss.backward()\n            optimizer.step()\n            running_loss += loss.item()\n\n        train_loss = running_loss / len(train_loader)\n        \n        model.eval()\n        all_valid_labels = []\n        all_valid_preds = []\n        with torch.no_grad():\n            for images, labels in tqdm(valid_loader, desc=\"验证中\"):\n                images = images.to(device)\n                outputs = model(images).squeeze()\n                all_valid_labels.extend(labels.cpu().numpy())\n                preds_np = outputs.cpu().numpy()\n                all_valid_preds.extend(np.atleast_1d(preds_np))\n        \n        rounded_preds = np.round(all_valid_preds).astype(int)\n        valid_kappa = cohen_kappa_score(all_valid_labels, rounded_preds, weights='quadratic')\n\n        print(f\"\\nEpoch {epoch + 1}/{num_epochs}\")\n        print(f\"  训练损失: {train_loss:.4f}\")\n        print(f\"  验证 Kappa (四舍五入): {valid_kappa:.4f}\")\n        print(f\"  当前学习率: {optimizer.param_groups[0]['lr']:.6f}\")\n\n        scheduler.step()\n\n        if valid_kappa > best_valid_kappa:\n            best_valid_kappa = valid_kappa\n            epochs_no_improve = 0\n            torch.save(model.state_dict(), save_path)\n            print(f\"  -> 已保存新最佳模型! 验证 Kappa: {best_valid_kappa:.4f}\")\n        else:\n            epochs_no_improve += 1\n        \n        if epochs_no_improve >= patience:\n            print(f\"\\n在 {patience} 个 epoch 没有提升后触发早停。\")\n            break\n\nclass OptimizedRounder:\n    def __init__(self):\n        self.coef_ = 0\n    def _kappa_loss(self, coef, X, y):\n        X_p = self.predict(X, coef)\n        ll = cohen_kappa_score(y, X_p, weights='quadratic')\n        return -ll\n    def fit(self, X, y):\n        loss_partial = lambda coef: self._kappa_loss(coef, X, y)\n        initial_coef = [0.5, 1.5, 2.5, 3.5]\n        self.coef_ = minimize(loss_partial, initial_coef, method='nelder-mead')['x']\n        print(f\"优化后的阈值: {self.coef_}\")\n    def predict(self, X, coef):\n        X_p = np.copy(X)\n        for i, pred in enumerate(X_p):\n            if pred < coef[0]: X_p[i] = 0\n            elif pred >= coef[0] and pred < coef[1]: X_p[i] = 1\n            elif pred >= coef[1] and pred < coef[2]: X_p[i] = 2\n            elif pred >= coef[2] and pred < coef[3]: X_p[i] = 3\n            else: X_p[i] = 4\n        return X_p.astype(int)\n\ndef get_validation_predictions(loader, model, device):\n    model.eval()\n    all_labels = []\n    all_preds = []\n    with torch.no_grad():\n        for images, labels in tqdm(loader, desc=\"获取验证集预测结果\"):\n            images = images.to(device)\n            outputs = model(images).squeeze()\n            all_labels.extend(labels.cpu().numpy())\n            preds_np = outputs.cpu().numpy()\n            all_preds.extend(np.atleast_1d(preds_np))\n    return np.array(all_preds), np.array(all_labels)\n\ndef load_and_predict_test_data(test_csv_path, img_dir, model, device, coefficients):\n    test_df = pd.read_csv(test_csv_path)\n    predictions = []\n    test_transform = get_transforms('test')\n    rounder = OptimizedRounder()\n    model.eval()\n    for img_id in tqdm(test_df['id_code'], desc=\"在测试集上预测\"):\n        img_path = os.path.join(img_dir, img_id + '.png')\n        if not os.path.exists(img_path):\n            predictions.append(-1) \n            continue\n        img = cv2.imread(img_path)\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        image_tensor = test_transform(image=img)[\"image\"].unsqueeze(0).to(device)\n        with torch.no_grad():\n            pred = model(image_tensor).squeeze().cpu().numpy()\n        predictions.append(pred)\n    if -1 in predictions:\n        valid_preds = [p for p in predictions if p != -1]\n        rounded_valid_preds = np.round(valid_preds).astype(int)\n        if len(rounded_valid_preds) > 0: mode_val = pd.Series(rounded_valid_preds).mode()[0]\n        else: mode_val = 0\n        predictions = [mode_val if p == -1 else p for p in predictions]\n    test_preds_rounded = rounder.predict(np.array(predictions), coefficients)\n    test_df['diagnosis'] = test_preds_rounded\n    submission_file = os.path.join(os.getcwd(), 'submission.csv')\n    test_df.to_csv(submission_file, index=False)\n    print(f\"\\n预测结果已保存至 {submission_file}\")\n\n\n# ==============================================================================\n# 4. 主执行逻辑 (Main Execution Logic)\n# ==============================================================================\nif __name__ == \"__main__\":\n    # --- 步骤 1: 定义路径和参数 ---\n    print(\"--- 步骤 1: 初始化路径和参数 ---\")\n    IMAGE_SIZE = 512\n    BATCH_SIZE = 4 \n    NUM_EPOCHS = 30\n    LEARNING_RATE = 1e-4\n    NUM_WORKERS = 2\n    \n    COMPETITION_DATA_DIR = 'aptos2019-blindness-detection'\n    # <--- MODIFICATION: Updated the pretrained model path\n    PRETRAINED_MODEL_PATH = '/kaggle/input/efficientnet-b6/pytorch/default/1/pytorch_model_effb6.bin'\n    \n    base_input_path = os.path.join('/kaggle/input', COMPETITION_DATA_DIR)\n    source_train_dir = os.path.join(base_input_path, 'train_images')\n    source_test_dir = os.path.join(base_input_path, 'test_images')\n    train_label_file = os.path.join(base_input_path, 'train.csv')\n    test_label_file = os.path.join(base_input_path, 'test.csv')\n\n    working_dir = os.getcwd()\n    dest_train_dir = os.path.join(working_dir, f'train_images_preprocessed_{IMAGE_SIZE}')\n    dest_test_dir = os.path.join(working_dir, f'test_images_preprocessed_{IMAGE_SIZE}')\n    model_path = os.path.join(working_dir, f'best_model_kappa_{IMAGE_SIZE}.pth')\n\n    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    print(f\"使用的设备: {device}\")\n\n    # --- 步骤 2: 执行数据预处理 (如果需要) ---\n    print(\"\\n--- 步骤 2: 检查预处理数据 ---\")\n    if not os.path.exists(dest_train_dir) or not os.path.exists(dest_test_dir):\n        print(\"未找到预处理数据。开始预处理...\")\n        preprocess_directory(source_train_dir, dest_train_dir, IMAGE_SIZE)\n        preprocess_directory(source_test_dir, dest_test_dir, IMAGE_SIZE)\n    else:\n        print(\"已找到预处理数据。跳过预处理。\")\n\n    # --- 步骤 3: 执行模型训练 (如果需要) ---\n    print(\"\\n--- 步骤 3: 检查现有模型 ---\")\n    run_training = True\n    if os.path.exists(model_path):\n        print(f\"在 {model_path} 找到模型文件。跳过训练。\")\n        run_training = False\n\n    if run_training:\n        print(\"开始模型训练...\")\n        all_labels_df = pd.read_csv(train_label_file)\n        train_df, valid_df = train_test_split(\n            all_labels_df, test_size=0.2, random_state=42, stratify=all_labels_df['diagnosis']\n        )\n        train_dataset = CustomImageDataset(dest_train_dir, train_df, transform=get_transforms('train'))\n        valid_dataset = CustomImageDataset(dest_train_dir, valid_df, transform=get_transforms('valid'))\n        train_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True, num_workers=NUM_WORKERS, pin_memory=True)\n        valid_loader = DataLoader(valid_dataset, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS, pin_memory=True)\n        model = EfficientNetModel(pretrained_model_path=PRETRAINED_MODEL_PATH).to(device)\n        train(model, train_loader, valid_loader, device, NUM_EPOCHS, LEARNING_RATE, IMAGE_SIZE)\n    \n    # --- 步骤 4: 优化舍入阈值 ---\n    print(\"\\n--- 步骤 4: 优化舍入阈值 ---\")\n    if not os.path.exists(model_path):\n         print(\"错误: 未找到模型文件。无法进行阈值优化。\")\n    else:\n        all_labels_df = pd.read_csv(train_label_file)\n        _, valid_df_for_opt = train_test_split(\n            all_labels_df, test_size=0.2, random_state=42, stratify=all_labels_df['diagnosis']\n        )\n        valid_dataset_for_opt = CustomImageDataset(dest_train_dir, valid_df_for_opt, transform=get_transforms('valid'))\n        valid_loader_for_opt = DataLoader(valid_dataset_for_opt, batch_size=BATCH_SIZE, shuffle=False, num_workers=NUM_WORKERS)\n        \n        model_for_opt = EfficientNetModel(pretrained_model_path=None).to(device)\n        model_for_opt.load_state_dict(torch.load(model_path, map_location=device))\n        \n        val_preds, val_labels = get_validation_predictions(valid_loader_for_opt, model_for_opt, device)\n        opt_rounder = OptimizedRounder()\n        opt_rounder.fit(val_preds, val_labels)\n        optimized_coefficients = opt_rounder.coef_\n        \n        optimized_preds = opt_rounder.predict(val_preds, optimized_coefficients)\n        final_val_kappa = cohen_kappa_score(val_labels, optimized_preds, weights='quadratic')\n        print(f\"最终优化后的验证集 Kappa: {final_val_kappa:.4f}\")\n    \n        # --- 步骤 5: 加载最佳模型并进行预测 ---\n        print(\"\\n--- 步骤 5: 加载最佳模型并在测试集上预测 ---\")\n        model_for_pred = EfficientNetModel(pretrained_model_path=None).to(device)\n        model_for_pred.load_state_dict(torch.load(model_path, map_location=device))\n        load_and_predict_test_data(test_label_file, dest_test_dir, model_for_pred, device, optimized_coefficients)\n\n    print(\"\\n--- 脚本执行结束 ---\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}