From d7f8f0cc8f01640f0a4216ee4ad340041c27b175 Mon Sep 17 00:00:00 2001 From: Haokun Date: Fri, 19 Dec 2025 17:49:01 +0100 Subject: [PATCH 01/11] ResNet34-UNet --- resnet.py | 379 ++++++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 379 insertions(+) create mode 100644 resnet.py diff --git a/resnet.py b/resnet.py new file mode 100644 index 0000000..4251a9b --- /dev/null +++ b/resnet.py @@ -0,0 +1,379 @@ +# -*- coding: utf-8 -*- +""" +ResNet34-UNet（简化版，可直接跑） +- 数据都在同一个目录：Desktop/deep_datachallenge/images +- 训练标签：Y_train_T9NrBYo.csv（flatten + -1 padding） +- 训练井：Well 1–6 +- 测试井：Well 7–11（同目录里筛选） +- 输出：submission.csv（每行一个 patch，flatten，pad 到 160*272 用 -1） + +注意你只需要改： +1) DATA_ROOT 路径 +2) EPOCHS/BATCH_SIZE 等超参数按你显存调整 +""" + +import re +from pathlib import Path + +import numpy as np +import pandas as pd + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.utils.data import Dataset, DataLoader, Subset + +from torchvision.models import resnet34, ResNet34_Weights + + +# ========================= +# 0. 超参数与路径 +# ========================= +DATA_ROOT = Path(r"C:\Users\lenovo\Desktop\deep_datachallenge") +IMAGES_DIR = DATA_ROOT / "images" +Y_TRAIN_CSV = DATA_ROOT / "Y_train_T9NrBYo.csv" + +TARGET_H = 160 +TARGET_W = 272 + +NUM_CLASSES = 3 # 你确认 CSV 里只有 0/1/2 +IGNORE_INDEX = -1 # CSV padding + +BATCH_SIZE = 8 +LR = 1e-3 +WEIGHT_DECAY = 1e-4 +EPOCHS = 20 + +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" + + +# ========================= +# 1. 工具函数 +# ========================= +def parse_well_id(name: str) -> int: + """从 well_1_section_0_patch_0 提取 well id=1""" + m = re.search(r"well_(\d+)_", name) + return int(m.group(1)) if m else -1 + + +def minmax_normalize(x: np.ndarray) -> np.ndarray: + """min-max 归一化；NaN/inf 置 0""" + x = x.astype(np.float32) + x = np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0) + x_min = float(x.min()) + x_max = float(x.max()) + if x_max - x_min < 1e-6: + return np.zeros_like(x, dtype=np.float32) + return (x - x_min) / (x_max - x_min) + + +def pad_to_160x272(img: np.ndarray, fill_value: float = 0.0) -> np.ndarray: + """把 (160,160) 或 (160,272) pad 到 (160,272)""" + h, w = img.shape + assert h == TARGET_H, f"期望高度 {TARGET_H}，但拿到 {h}" + if w == TARGET_W: + return img + if w < TARGET_W: + out = np.full((TARGET_H, TARGET_W), fill_value, dtype=img.dtype) + out[:, :w] = img + return out + # 若更宽，简单裁剪（一般不会发生） + return img[:, :TARGET_W] + + +def decode_mask_from_csv_row(row_values: np.ndarray) -> np.ndarray: + """ + 从 CSV 一行恢复 mask： + - row_values: flatten + -1 padding + - 去掉 -1 后 reshape 成 (160, w) + """ + valid = row_values[row_values != IGNORE_INDEX] + assert len(valid) % TARGET_H == 0, f"mask 有效长度 {len(valid)} 不能被 160 整除" + w = len(valid) // TARGET_H + return valid.reshape(TARGET_H, w).astype(np.int64) + + +def pad_mask_to_160x272(mask: np.ndarray) -> np.ndarray: + """把 (160,w) pad 到 (160,272)，pad 用 -1（ignore）""" + h, w = mask.shape + assert h == TARGET_H + if w == TARGET_W: + return mask + out = np.full((TARGET_H, TARGET_W), IGNORE_INDEX, dtype=np.int64) + out[:, :w] = mask + return out + + +# ========================= +# 2. Dataset（支持按 wells 过滤） +# ========================= +class WellSegDataset(Dataset): + def __init__(self, images_dir: Path, y_csv_path: Path = None, wells=None): + """ + wells: 例如 {1,2,3,4,5,6} 或 {7,8,9,10,11} + y_csv_path=None 表示无标签（测试） + """ + self.images_dir = images_dir + self.has_label = y_csv_path is not None + + all_paths = sorted(images_dir.glob("*.npy")) + all_names = [p.stem for p in all_paths] + + if wells is not None: + keep = [] + for p, n in zip(all_paths, all_names): + w = parse_well_id(n) + if w in wells: + keep.append((p, n)) + self.image_paths = [x[0] for x in keep] + self.names = [x[1] for x in keep] + else: + self.image_paths = all_paths + self.names = all_names + + if self.has_label: + # CSV index 通常就是 patch 名（不含 .npy） + self.y_df = pd.read_csv(y_csv_path, index_col=0) + else: + self.y_df = None + + def __len__(self): + return len(self.image_paths) + + def __getitem__(self, idx: int): + name = self.names[idx] + img_path = self.image_paths[idx] + + # 读图 + img = np.load(img_path) # (160,160) or (160,272) + raw_w = img.shape[1] # 记录原始宽度（推理时要裁回去） + img = minmax_normalize(img) + img = pad_to_160x272(img, fill_value=0.0) + img_t = torch.from_numpy(img).unsqueeze(0).float() # (1,160,272) + + if not self.has_label: + return {"name": name, "image": img_t, "raw_w": raw_w} + + # 读 mask + row = self.y_df.loc[name].values.astype(np.int64) + mask = decode_mask_from_csv_row(row) # (160,w) + mask = pad_mask_to_160x272(mask) # (160,272) + mask_t = torch.from_numpy(mask).long() + + return {"name": name, "image": img_t, "mask": mask_t, "raw_w": raw_w} + + +# ========================= +# 3. ResNet34-UNet（简化实现） +# ========================= +class ConvRelu(nn.Module): + def __init__(self, in_ch, out_ch): + super().__init__() + self.block = nn.Sequential( + nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False), + nn.BatchNorm2d(out_ch), + nn.ReLU(inplace=True), + ) + + def forward(self, x): + return self.block(x) + + +class UpBlock(nn.Module): + def __init__(self, in_ch, skip_ch, out_ch): + super().__init__() + self.conv1 = ConvRelu(in_ch + skip_ch, out_ch) + self.conv2 = ConvRelu(out_ch, out_ch) + + def forward(self, x, skip): + x = F.interpolate(x, size=skip.shape[-2:], mode="bilinear", align_corners=False) + x = torch.cat([x, skip], dim=1) + x = self.conv1(x) + x = self.conv2(x) + return x + + +class ResNet34UNet(nn.Module): + def __init__(self, num_classes: int): + super().__init__() + backbone = resnet34(weights=ResNet34_Weights.IMAGENET1K_V1) + + # 输入单通道：把第一层卷积改成 1 通道 + old_conv1 = backbone.conv1 + new_conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False) + with torch.no_grad(): + new_conv1.weight[:] = old_conv1.weight.mean(dim=1, keepdim=True) + backbone.conv1 = new_conv1 + + self.enc0 = nn.Sequential(backbone.conv1, backbone.bn1, backbone.relu) # /2 + self.pool0 = backbone.maxpool # /4 + self.enc1 = backbone.layer1 # /4 + self.enc2 = backbone.layer2 # /8 + self.enc3 = backbone.layer3 # /16 + self.enc4 = backbone.layer4 # /32 + + self.center = nn.Sequential(ConvRelu(512, 512), ConvRelu(512, 512)) + self.up4 = UpBlock(512, 256, 256) + self.up3 = UpBlock(256, 128, 128) + self.up2 = UpBlock(128, 64, 64) + self.up1 = UpBlock(64, 64, 64) + + self.head = nn.Conv2d(64, num_classes, kernel_size=1) + + def forward(self, x): + e0 = self.enc0(x) # 64, H/2, W/2 + e1 = self.enc1(self.pool0(e0)) # 64, H/4, W/4 + e2 = self.enc2(e1) # 128, H/8, W/8 + e3 = self.enc3(e2) # 256, H/16, W/16 + e4 = self.enc4(e3) # 512, H/32, W/32 + + c = self.center(e4) + d4 = self.up4(c, e3) + d3 = self.up3(d4, e2) + d2 = self.up2(d3, e1) + d1 = self.up1(d2, e0) + + out = self.head(d1) + out = F.interpolate(out, size=(TARGET_H, TARGET_W), mode="bilinear", align_corners=False) + return out + + +# ========================= +# 4. 训练与验证（最简单） +# ========================= +def train_one_epoch(model, loader, optimizer): + model.train() + total_loss = 0.0 + + for batch in loader: + x = batch["image"].to(DEVICE) # (B,1,160,272) + y = batch["mask"].to(DEVICE) # (B,160,272) 包含 -1 + + logits = model(x) # (B,C,160,272) + loss = F.cross_entropy(logits, y, ignore_index=IGNORE_INDEX) + + optimizer.zero_grad() + loss.backward() + optimizer.step() + + total_loss += loss.item() * x.size(0) + + return total_loss / len(loader.dataset) + + +@torch.no_grad() +def valid_one_epoch(model, loader): + model.eval() + total_loss = 0.0 + + for batch in loader: + x = batch["image"].to(DEVICE) + y = batch["mask"].to(DEVICE) + + logits = model(x) + loss = F.cross_entropy(logits, y, ignore_index=IGNORE_INDEX) + total_loss += loss.item() * x.size(0) + + return total_loss / len(loader.dataset) + + +# ========================= +# 5. 推理并生成提交 CSV +# ========================= +@torch.no_grad() +def predict_and_make_submission(model, images_dir: Path, out_csv_path: Path, test_wells: set): + """ + 从 images_dir 中筛选 test_wells 预测并生成提交 CSV + - 每行：一个 patch + - 长度：160*272 + - 如果原始宽度 < 272，剩余用 -1 padding + """ + model.eval() + + test_ds = WellSegDataset(images_dir, y_csv_path=None, wells=test_wells) + test_loader = DataLoader(test_ds, batch_size=1, shuffle=False, num_workers=0) + + preds_dict = {} + + for batch in test_loader: + name = batch["name"][0] + raw_w = int(batch["raw_w"][0]) # 原始宽度 160 或 272 + x = batch["image"].to(DEVICE) # (1,1,160,272) + + logits = model(x) # (1,C,160,272) + pred_full = torch.argmax(logits, dim=1).squeeze(0).cpu().numpy().astype(np.int64) # (160,272) + + # 裁回原始宽度 + pred = pred_full[:, :raw_w] # (160,raw_w) + + # flatten + pad 到 160*272 + if raw_w < TARGET_W: + padded = np.full((TARGET_H * TARGET_W,), IGNORE_INDEX, dtype=np.int64) + padded[: TARGET_H * raw_w] = pred.flatten() + preds_dict[name] = padded + else: + preds_dict[name] = pred.flatten() + + sub = pd.DataFrame(preds_dict, dtype="int64").T + sub.to_csv(out_csv_path) + print(f"[OK] submission 已保存: {out_csv_path}") + + +# ========================= +# 6. 主函数：严格按井划分训练/验证，测试井预测提交 +# ========================= +def main(): + # ====== (A) 训练井与测试井定义 ====== + TRAIN_WELLS = {1, 2, 3, 4, 5, 6} + TEST_WELLS = {7, 8, 9, 10, 11} + + # 验证集：从训练井里“按井留出”避免泄漏（例：留 well6） + VAL_WELLS = {6} + + # ====== (B) 构建训练集（只读 well1-6） ====== + train_ds_all = WellSegDataset(IMAGES_DIR, Y_TRAIN_CSV, wells=TRAIN_WELLS) + + # 按井划分 train/val + train_indices, val_indices = [], [] + for i, name in enumerate(train_ds_all.names): + w = parse_well_id(name) + if w in VAL_WELLS: + val_indices.append(i) + else: + train_indices.append(i) + + train_ds = Subset(train_ds_all, train_indices) + val_ds = Subset(train_ds_all, val_indices) + + train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=0) + val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=0) + + print(f"训练样本数: {len(train_ds)} | 验证样本数: {len(val_ds)} | val_wells={VAL_WELLS}") + + # ====== (C) 模型与优化器 ====== + model = ResNet34UNet(num_classes=NUM_CLASSES).to(DEVICE) + optimizer = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY) + + # ====== (D) 训练 ====== + best_val = 1e9 + best_path = DATA_ROOT / "best_resnet34_unet.pth" + + for epoch in range(1, EPOCHS + 1): + tr_loss = train_one_epoch(model, train_loader, optimizer) + va_loss = valid_one_epoch(model, val_loader) + + print(f"Epoch {epoch:02d}/{EPOCHS} | train_loss={tr_loss:.4f} | val_loss={va_loss:.4f}") + + if va_loss < best_val: + best_val = va_loss + torch.save(model.state_dict(), best_path) + print(f" -> 保存最优模型: {best_path}") + + # ====== (E) 生成提交（从同一个 images/ 里筛选 well7-11） ====== + out_csv = DATA_ROOT / "submission.csv" + model.load_state_dict(torch.load(best_path, map_location=DEVICE)) + predict_and_make_submission(model, IMAGES_DIR, out_csv, test_wells=TEST_WELLS) + + +if __name__ == "__main__": + main() From 65d776817320b1a8a47bb895c786ada1ead5d461 Mon Sep 17 00:00:00 2001 From: Haokun Date: Fri, 19 Dec 2025 22:28:32 +0100 Subject: [PATCH 02/11] resnet34,score:0.4582 --- resnet.ipynb | 414 +++++++++++++++++++++++++++++++++++++++++++++++++++ resnet.py | 109 +++++--------- 2 files changed, 454 insertions(+), 69 deletions(-) create mode 100644 resnet.ipynb diff --git a/resnet.ipynb b/resnet.ipynb new file mode 100644 index 0000000..e02f1de --- /dev/null +++ b/resnet.ipynb @@ -0,0 +1,414 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "51ec4b7d", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "训练样本数: 2790 | 验证样本数: 1620 | val_wells={6}\n", + "Epoch 01/5 | train_loss=0.1464 | val_loss=0.0906\n", + " -> 保存最优模型: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_resnet34_unet.pth\n", + "Epoch 02/5 | train_loss=0.0871 | val_loss=0.0923\n", + "Epoch 03/5 | train_loss=0.0784 | val_loss=0.0879\n", + " -> 保存最优模型: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_resnet34_unet.pth\n", + "Epoch 04/5 | train_loss=0.0767 | val_loss=0.0893\n", + "Epoch 05/5 | train_loss=0.0745 | val_loss=0.1063\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "C:\\Users\\lenovo\\AppData\\Local\\Temp\\ipykernel_12820\\130010335.py:345: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n", + " model.load_state_dict(torch.load(best_path, map_location=DEVICE))\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[OK] submission 已保存: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\submission.csv\n" + ] + } + ], + "source": [ + "# -*- coding: utf-8 -*-\n", + "\"\"\"\n", + "ResNet34-UNet（简化版，可直接跑）\n", + "- 训练图像目录：X_train_uDRk9z9/images（well1-6）\n", + "- 测试图像目录：X_test_xNbnvIa/images（well7-11）\n", + "- 训练标签：Y_train_T9NrBYo.csv（flatten + -1 padding）\n", + "- 验证：从训练集中按井划分（例：well6 为 val，其余为 train）\n", + "- 输出：submission.csv（每行一个 patch，flatten，pad 到 160*272 用 -1）\n", + "\"\"\"\n", + "\n", + "import re\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "import pandas as pd\n", + "\n", + "import torch\n", + "import torch.nn as nn\n", + "import torch.nn.functional as F\n", + "from torch.utils.data import Dataset, DataLoader, Subset\n", + "\n", + "from torchvision.models import resnet34, ResNet34_Weights\n", + "\n", + "\n", + "# =========================\n", + "# 0. 超参数与路径\n", + "# =========================\n", + "DATA_ROOT = Path(r\"C:\\Users\\lenovo\\Desktop\\deep_datachallenge\") # 改成你的真实路径\n", + "\n", + "TRAIN_IMAGES_DIR = DATA_ROOT / \"X_train_uDRk9z9\" / \"images\"\n", + "TEST_IMAGES_DIR = DATA_ROOT / \"X_test_xNbnvIa\" / \"images\"\n", + "Y_TRAIN_CSV = DATA_ROOT / \"Y_train_T9NrBYo.csv\"\n", + "\n", + "TARGET_H = 160\n", + "TARGET_W = 272\n", + "\n", + "NUM_CLASSES = 3 # 你确认 CSV 里只有 0/1/2\n", + "IGNORE_INDEX = -1 # CSV padding\n", + "\n", + "BATCH_SIZE = 8\n", + "LR = 1e-3\n", + "WEIGHT_DECAY = 1e-4\n", + "EPOCHS = 5 # 训练慢就先设 5，跑通后再加\n", + "\n", + "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", + "\n", + "\n", + "# =========================\n", + "# 1. 工具函数\n", + "# =========================\n", + "def parse_well_id(name: str) -> int:\n", + " \"\"\"从 well_1_section_0_patch_0 提取 well id=1\"\"\"\n", + " m = re.search(r\"well_(\\d+)_\", name)\n", + " return int(m.group(1)) if m else -1\n", + "\n", + "\n", + "def minmax_normalize(x: np.ndarray) -> np.ndarray:\n", + " \"\"\"min-max 归一化；NaN/inf 置 0\"\"\"\n", + " x = x.astype(np.float32)\n", + " x = np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)\n", + " x_min = float(x.min())\n", + " x_max = float(x.max())\n", + " if x_max - x_min < 1e-6:\n", + " return np.zeros_like(x, dtype=np.float32)\n", + " return (x - x_min) / (x_max - x_min)\n", + "\n", + "\n", + "def pad_to_160x272(img: np.ndarray, fill_value: float = 0.0) -> np.ndarray:\n", + " \"\"\"把 (160,160) 或 (160,272) pad 到 (160,272)\"\"\"\n", + " h, w = img.shape\n", + " assert h == TARGET_H, f\"期望高度 {TARGET_H}，但拿到 {h}\"\n", + " if w == TARGET_W:\n", + " return img\n", + " if w < TARGET_W:\n", + " out = np.full((TARGET_H, TARGET_W), fill_value, dtype=img.dtype)\n", + " out[:, :w] = img\n", + " return out\n", + " return img[:, :TARGET_W]\n", + "\n", + "\n", + "def decode_mask_from_csv_row(row_values: np.ndarray) -> np.ndarray:\n", + " \"\"\"\n", + " 从 CSV 一行恢复 mask：\n", + " - row_values: flatten + -1 padding\n", + " - 去掉 -1 后 reshape 成 (160, w)\n", + " \"\"\"\n", + " valid = row_values[row_values != IGNORE_INDEX]\n", + " assert len(valid) % TARGET_H == 0, f\"mask 有效长度 {len(valid)} 不能被 160 整除\"\n", + " w = len(valid) // TARGET_H\n", + " return valid.reshape(TARGET_H, w).astype(np.int64)\n", + "\n", + "\n", + "def pad_mask_to_160x272(mask: np.ndarray) -> np.ndarray:\n", + " \"\"\"把 (160,w) pad 到 (160,272)，pad 用 -1（ignore）\"\"\"\n", + " h, w = mask.shape\n", + " assert h == TARGET_H\n", + " if w == TARGET_W:\n", + " return mask\n", + " out = np.full((TARGET_H, TARGET_W), IGNORE_INDEX, dtype=np.int64)\n", + " out[:, :w] = mask\n", + " return out\n", + "\n", + "\n", + "# =========================\n", + "# 2. Dataset（训练/测试共用）\n", + "# =========================\n", + "class WellSegDataset(Dataset):\n", + " def __init__(self, images_dir: Path, y_csv_path: Path = None):\n", + " \"\"\"\n", + " y_csv_path=None 表示无标签（测试）\n", + " \"\"\"\n", + " self.images_dir = images_dir\n", + " self.has_label = y_csv_path is not None\n", + "\n", + " self.image_paths = sorted(images_dir.glob(\"*.npy\"))\n", + " self.names = [p.stem for p in self.image_paths]\n", + "\n", + " if self.has_label:\n", + " # CSV index 通常就是 patch 名（不含 .npy）\n", + " self.y_df = pd.read_csv(y_csv_path, index_col=0)\n", + " else:\n", + " self.y_df = None\n", + "\n", + " def __len__(self):\n", + " return len(self.image_paths)\n", + "\n", + " def __getitem__(self, idx: int):\n", + " name = self.names[idx]\n", + " img_path = self.image_paths[idx]\n", + "\n", + " img = np.load(img_path) # (160,160) or (160,272)\n", + " raw_w = img.shape[1] # 记录原始宽度（推理时裁回去）\n", + " img = minmax_normalize(img)\n", + " img = pad_to_160x272(img, fill_value=0.0)\n", + " img_t = torch.from_numpy(img).unsqueeze(0).float() # (1,160,272)\n", + "\n", + " if not self.has_label:\n", + " return {\"name\": name, \"image\": img_t, \"raw_w\": raw_w}\n", + "\n", + " row = self.y_df.loc[name].values.astype(np.int64)\n", + " mask = decode_mask_from_csv_row(row) # (160,w)\n", + " mask = pad_mask_to_160x272(mask) # (160,272)\n", + " mask_t = torch.from_numpy(mask).long()\n", + "\n", + " return {\"name\": name, \"image\": img_t, \"mask\": mask_t, \"raw_w\": raw_w}\n", + "\n", + "\n", + "# =========================\n", + "# 3. ResNet34-UNet（简化实现）\n", + "# =========================\n", + "class ConvRelu(nn.Module):\n", + " def __init__(self, in_ch, out_ch):\n", + " super().__init__()\n", + " self.block = nn.Sequential(\n", + " nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False),\n", + " nn.BatchNorm2d(out_ch),\n", + " nn.ReLU(inplace=True),\n", + " )\n", + "\n", + " def forward(self, x):\n", + " return self.block(x)\n", + "\n", + "\n", + "class UpBlock(nn.Module):\n", + " def __init__(self, in_ch, skip_ch, out_ch):\n", + " super().__init__()\n", + " self.conv1 = ConvRelu(in_ch + skip_ch, out_ch)\n", + " self.conv2 = ConvRelu(out_ch, out_ch)\n", + "\n", + " def forward(self, x, skip):\n", + " x = F.interpolate(x, size=skip.shape[-2:], mode=\"bilinear\", align_corners=False)\n", + " x = torch.cat([x, skip], dim=1)\n", + " x = self.conv1(x)\n", + " x = self.conv2(x)\n", + " return x\n", + "\n", + "\n", + "class ResNet34UNet(nn.Module):\n", + " def __init__(self, num_classes: int):\n", + " super().__init__()\n", + " backbone = resnet34(weights=ResNet34_Weights.IMAGENET1K_V1)\n", + "\n", + " # 输入单通道：第一层卷积改成 1 通道（用原权重均值初始化）\n", + " old_conv1 = backbone.conv1\n", + " new_conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)\n", + " with torch.no_grad():\n", + " new_conv1.weight[:] = old_conv1.weight.mean(dim=1, keepdim=True)\n", + " backbone.conv1 = new_conv1\n", + "\n", + " self.enc0 = nn.Sequential(backbone.conv1, backbone.bn1, backbone.relu) # /2\n", + " self.pool0 = backbone.maxpool # /4\n", + " self.enc1 = backbone.layer1 # /4\n", + " self.enc2 = backbone.layer2 # /8\n", + " self.enc3 = backbone.layer3 # /16\n", + " self.enc4 = backbone.layer4 # /32\n", + "\n", + " self.center = nn.Sequential(ConvRelu(512, 512), ConvRelu(512, 512))\n", + " self.up4 = UpBlock(512, 256, 256)\n", + " self.up3 = UpBlock(256, 128, 128)\n", + " self.up2 = UpBlock(128, 64, 64)\n", + " self.up1 = UpBlock(64, 64, 64)\n", + "\n", + " self.head = nn.Conv2d(64, num_classes, kernel_size=1)\n", + "\n", + " def forward(self, x):\n", + " e0 = self.enc0(x)\n", + " e1 = self.enc1(self.pool0(e0))\n", + " e2 = self.enc2(e1)\n", + " e3 = self.enc3(e2)\n", + " e4 = self.enc4(e3)\n", + "\n", + " c = self.center(e4)\n", + " d4 = self.up4(c, e3)\n", + " d3 = self.up3(d4, e2)\n", + " d2 = self.up2(d3, e1)\n", + " d1 = self.up1(d2, e0)\n", + "\n", + " out = self.head(d1)\n", + " out = F.interpolate(out, size=(TARGET_H, TARGET_W), mode=\"bilinear\", align_corners=False)\n", + " return out\n", + "\n", + "\n", + "# =========================\n", + "# 4. 训练与验证（最简单）\n", + "# =========================\n", + "def train_one_epoch(model, loader, optimizer):\n", + " model.train()\n", + " total_loss = 0.0\n", + " for batch in loader:\n", + " x = batch[\"image\"].to(DEVICE)\n", + " y = batch[\"mask\"].to(DEVICE)\n", + "\n", + " logits = model(x)\n", + " loss = F.cross_entropy(logits, y, ignore_index=IGNORE_INDEX)\n", + "\n", + " optimizer.zero_grad()\n", + " loss.backward()\n", + " optimizer.step()\n", + "\n", + " total_loss += loss.item() * x.size(0)\n", + "\n", + " return total_loss / len(loader.dataset)\n", + "\n", + "\n", + "@torch.no_grad()\n", + "def valid_one_epoch(model, loader):\n", + " model.eval()\n", + " total_loss = 0.0\n", + " for batch in loader:\n", + " x = batch[\"image\"].to(DEVICE)\n", + " y = batch[\"mask\"].to(DEVICE)\n", + "\n", + " logits = model(x)\n", + " loss = F.cross_entropy(logits, y, ignore_index=IGNORE_INDEX)\n", + " total_loss += loss.item() * x.size(0)\n", + "\n", + " return total_loss / len(loader.dataset)\n", + "\n", + "\n", + "# =========================\n", + "# 5. 推理并生成提交 CSV\n", + "# =========================\n", + "@torch.no_grad()\n", + "def predict_and_make_submission(model, test_images_dir: Path, out_csv_path: Path):\n", + " \"\"\"\n", + " 对 test_images_dir 全部 npy 预测并生成提交 CSV\n", + " - 每行：一个 patch\n", + " - 长度：160*272\n", + " - 如果原始宽度 < 272，剩余用 -1 padding\n", + " \"\"\"\n", + " model.eval()\n", + "\n", + " test_ds = WellSegDataset(test_images_dir, y_csv_path=None)\n", + " test_loader = DataLoader(test_ds, batch_size=1, shuffle=False, num_workers=0)\n", + "\n", + " preds_dict = {}\n", + "\n", + " for batch in test_loader:\n", + " name = batch[\"name\"][0]\n", + " raw_w = int(batch[\"raw_w\"][0])\n", + " x = batch[\"image\"].to(DEVICE)\n", + "\n", + " logits = model(x)\n", + " pred_full = torch.argmax(logits, dim=1).squeeze(0).cpu().numpy().astype(np.int64) # (160,272)\n", + "\n", + " pred = pred_full[:, :raw_w] # 裁回原始宽度\n", + " if raw_w < TARGET_W:\n", + " padded = np.full((TARGET_H * TARGET_W,), IGNORE_INDEX, dtype=np.int64)\n", + " padded[: TARGET_H * raw_w] = pred.flatten()\n", + " preds_dict[name] = padded\n", + " else:\n", + " preds_dict[name] = pred.flatten()\n", + "\n", + " sub = pd.DataFrame(preds_dict, dtype=\"int64\").T\n", + " sub.to_csv(out_csv_path)\n", + " print(f\"[OK] submission 已保存: {out_csv_path}\")\n", + "\n", + "\n", + "# =========================\n", + "# 6. 主函数：训练(井1-5) + 验证(井6) + 预测test(井7-11目录)\n", + "# =========================\n", + "def main():\n", + " # (A) 构建训练集（well1-6）\n", + " train_ds_all = WellSegDataset(TRAIN_IMAGES_DIR, Y_TRAIN_CSV)\n", + "\n", + " # (B) 按井划分 train/val：well6 做验证\n", + " VAL_WELLS = {6}\n", + " train_indices, val_indices = [], []\n", + " for i, name in enumerate(train_ds_all.names):\n", + " w = parse_well_id(name)\n", + " if w in VAL_WELLS:\n", + " val_indices.append(i)\n", + " else:\n", + " train_indices.append(i)\n", + "\n", + " train_ds = Subset(train_ds_all, train_indices) # well1-5\n", + " val_ds = Subset(train_ds_all, val_indices) # well6\n", + "\n", + " train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)\n", + " val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)\n", + "\n", + " print(f\"训练样本数: {len(train_ds)} | 验证样本数: {len(val_ds)} | val_wells={VAL_WELLS}\")\n", + "\n", + " # (C) 模型与优化器\n", + " model = ResNet34UNet(num_classes=NUM_CLASSES).to(DEVICE)\n", + " optimizer = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)\n", + "\n", + " # (D) 训练\n", + " best_val = 1e9\n", + " best_path = DATA_ROOT / \"best_resnet34_unet.pth\"\n", + "\n", + " for epoch in range(1, EPOCHS + 1):\n", + " tr_loss = train_one_epoch(model, train_loader, optimizer)\n", + " va_loss = valid_one_epoch(model, val_loader)\n", + "\n", + " print(f\"Epoch {epoch:02d}/{EPOCHS} | train_loss={tr_loss:.4f} | val_loss={va_loss:.4f}\")\n", + "\n", + " if va_loss < best_val:\n", + " best_val = va_loss\n", + " torch.save(model.state_dict(), best_path)\n", + " print(f\" -> 保存最优模型: {best_path}\")\n", + "\n", + " # (E) 生成提交（测试目录 well7-11）\n", + " out_csv = DATA_ROOT / \"submission.csv\"\n", + " model.load_state_dict(torch.load(best_path, map_location=DEVICE))\n", + " predict_and_make_submission(model, TEST_IMAGES_DIR, out_csv)\n", + "\n", + "\n", + "if __name__ == \"__main__\":\n", + " main()\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "deep-torch", + "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.10.19" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/resnet.py b/resnet.py index 4251a9b..e2b7975 100644 --- a/resnet.py +++ b/resnet.py @@ -1,15 +1,11 @@ # -*- coding: utf-8 -*- """ ResNet34-UNet（简化版，可直接跑） -- 数据都在同一个目录：Desktop/deep_datachallenge/images +- 训练图像目录：X_train_uDRk9z9/images（well1-6） +- 测试图像目录：X_test_xNbnvIa/images（well7-11） - 训练标签：Y_train_T9NrBYo.csv（flatten + -1 padding） -- 训练井：Well 1–6 -- 测试井：Well 7–11（同目录里筛选） +- 验证：从训练集中按井划分（例：well6 为 val，其余为 train） - 输出：submission.csv（每行一个 patch，flatten，pad 到 160*272 用 -1） - -注意你只需要改： -1) DATA_ROOT 路径 -2) EPOCHS/BATCH_SIZE 等超参数按你显存调整 """ import re @@ -29,8 +25,10 @@ # ========================= # 0. 超参数与路径 # ========================= -DATA_ROOT = Path(r"C:\Users\lenovo\Desktop\deep_datachallenge") -IMAGES_DIR = DATA_ROOT / "images" +DATA_ROOT = Path(r"C:\Users\lenovo\Desktop\deep_datachallenge") # 改成你的真实路径 + +TRAIN_IMAGES_DIR = DATA_ROOT / "X_train_uDRk9z9" / "images" +TEST_IMAGES_DIR = DATA_ROOT / "X_test_xNbnvIa" / "images" Y_TRAIN_CSV = DATA_ROOT / "Y_train_T9NrBYo.csv" TARGET_H = 160 @@ -42,7 +40,7 @@ BATCH_SIZE = 8 LR = 1e-3 WEIGHT_DECAY = 1e-4 -EPOCHS = 20 +EPOCHS = 5 # 训练慢就先设 5，跑通后再加 DEVICE = "cuda" if torch.cuda.is_available() else "cpu" @@ -77,7 +75,6 @@ def pad_to_160x272(img: np.ndarray, fill_value: float = 0.0) -> np.ndarray: out = np.full((TARGET_H, TARGET_W), fill_value, dtype=img.dtype) out[:, :w] = img return out - # 若更宽，简单裁剪（一般不会发生） return img[:, :TARGET_W] @@ -105,31 +102,18 @@ def pad_mask_to_160x272(mask: np.ndarray) -> np.ndarray: # ========================= -# 2. Dataset（支持按 wells 过滤） +# 2. Dataset（训练/测试共用） # ========================= class WellSegDataset(Dataset): - def __init__(self, images_dir: Path, y_csv_path: Path = None, wells=None): + def __init__(self, images_dir: Path, y_csv_path: Path = None): """ - wells: 例如 {1,2,3,4,5,6} 或 {7,8,9,10,11} y_csv_path=None 表示无标签（测试） """ self.images_dir = images_dir self.has_label = y_csv_path is not None - all_paths = sorted(images_dir.glob("*.npy")) - all_names = [p.stem for p in all_paths] - - if wells is not None: - keep = [] - for p, n in zip(all_paths, all_names): - w = parse_well_id(n) - if w in wells: - keep.append((p, n)) - self.image_paths = [x[0] for x in keep] - self.names = [x[1] for x in keep] - else: - self.image_paths = all_paths - self.names = all_names + self.image_paths = sorted(images_dir.glob("*.npy")) + self.names = [p.stem for p in self.image_paths] if self.has_label: # CSV index 通常就是 patch 名（不含 .npy） @@ -144,9 +128,8 @@ def __getitem__(self, idx: int): name = self.names[idx] img_path = self.image_paths[idx] - # 读图 - img = np.load(img_path) # (160,160) or (160,272) - raw_w = img.shape[1] # 记录原始宽度（推理时要裁回去） + img = np.load(img_path) # (160,160) or (160,272) + raw_w = img.shape[1] # 记录原始宽度（推理时裁回去） img = minmax_normalize(img) img = pad_to_160x272(img, fill_value=0.0) img_t = torch.from_numpy(img).unsqueeze(0).float() # (1,160,272) @@ -154,7 +137,6 @@ def __getitem__(self, idx: int): if not self.has_label: return {"name": name, "image": img_t, "raw_w": raw_w} - # 读 mask row = self.y_df.loc[name].values.astype(np.int64) mask = decode_mask_from_csv_row(row) # (160,w) mask = pad_mask_to_160x272(mask) # (160,272) @@ -198,7 +180,7 @@ def __init__(self, num_classes: int): super().__init__() backbone = resnet34(weights=ResNet34_Weights.IMAGENET1K_V1) - # 输入单通道：把第一层卷积改成 1 通道 + # 输入单通道：第一层卷积改成 1 通道（用原权重均值初始化） old_conv1 = backbone.conv1 new_conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False) with torch.no_grad(): @@ -221,11 +203,11 @@ def __init__(self, num_classes: int): self.head = nn.Conv2d(64, num_classes, kernel_size=1) def forward(self, x): - e0 = self.enc0(x) # 64, H/2, W/2 - e1 = self.enc1(self.pool0(e0)) # 64, H/4, W/4 - e2 = self.enc2(e1) # 128, H/8, W/8 - e3 = self.enc3(e2) # 256, H/16, W/16 - e4 = self.enc4(e3) # 512, H/32, W/32 + e0 = self.enc0(x) + e1 = self.enc1(self.pool0(e0)) + e2 = self.enc2(e1) + e3 = self.enc3(e2) + e4 = self.enc4(e3) c = self.center(e4) d4 = self.up4(c, e3) @@ -244,12 +226,11 @@ def forward(self, x): def train_one_epoch(model, loader, optimizer): model.train() total_loss = 0.0 - for batch in loader: - x = batch["image"].to(DEVICE) # (B,1,160,272) - y = batch["mask"].to(DEVICE) # (B,160,272) 包含 -1 + x = batch["image"].to(DEVICE) + y = batch["mask"].to(DEVICE) - logits = model(x) # (B,C,160,272) + logits = model(x) loss = F.cross_entropy(logits, y, ignore_index=IGNORE_INDEX) optimizer.zero_grad() @@ -265,7 +246,6 @@ def train_one_epoch(model, loader, optimizer): def valid_one_epoch(model, loader): model.eval() total_loss = 0.0 - for batch in loader: x = batch["image"].to(DEVICE) y = batch["mask"].to(DEVICE) @@ -281,32 +261,29 @@ def valid_one_epoch(model, loader): # 5. 推理并生成提交 CSV # ========================= @torch.no_grad() -def predict_and_make_submission(model, images_dir: Path, out_csv_path: Path, test_wells: set): +def predict_and_make_submission(model, test_images_dir: Path, out_csv_path: Path): """ - 从 images_dir 中筛选 test_wells 预测并生成提交 CSV + 对 test_images_dir 全部 npy 预测并生成提交 CSV - 每行：一个 patch - 长度：160*272 - 如果原始宽度 < 272，剩余用 -1 padding """ model.eval() - test_ds = WellSegDataset(images_dir, y_csv_path=None, wells=test_wells) + test_ds = WellSegDataset(test_images_dir, y_csv_path=None) test_loader = DataLoader(test_ds, batch_size=1, shuffle=False, num_workers=0) preds_dict = {} for batch in test_loader: name = batch["name"][0] - raw_w = int(batch["raw_w"][0]) # 原始宽度 160 或 272 - x = batch["image"].to(DEVICE) # (1,1,160,272) + raw_w = int(batch["raw_w"][0]) + x = batch["image"].to(DEVICE) - logits = model(x) # (1,C,160,272) + logits = model(x) pred_full = torch.argmax(logits, dim=1).squeeze(0).cpu().numpy().astype(np.int64) # (160,272) - # 裁回原始宽度 - pred = pred_full[:, :raw_w] # (160,raw_w) - - # flatten + pad 到 160*272 + pred = pred_full[:, :raw_w] # 裁回原始宽度 if raw_w < TARGET_W: padded = np.full((TARGET_H * TARGET_W,), IGNORE_INDEX, dtype=np.int64) padded[: TARGET_H * raw_w] = pred.flatten() @@ -320,20 +297,14 @@ def predict_and_make_submission(model, images_dir: Path, out_csv_path: Path, tes # ========================= -# 6. 主函数：严格按井划分训练/验证，测试井预测提交 +# 6. 主函数：训练(井1-5) + 验证(井6) + 预测test(井7-11目录) # ========================= def main(): - # ====== (A) 训练井与测试井定义 ====== - TRAIN_WELLS = {1, 2, 3, 4, 5, 6} - TEST_WELLS = {7, 8, 9, 10, 11} + # (A) 构建训练集（well1-6） + train_ds_all = WellSegDataset(TRAIN_IMAGES_DIR, Y_TRAIN_CSV) - # 验证集：从训练井里“按井留出”避免泄漏（例：留 well6） + # (B) 按井划分 train/val：well6 做验证 VAL_WELLS = {6} - - # ====== (B) 构建训练集（只读 well1-6） ====== - train_ds_all = WellSegDataset(IMAGES_DIR, Y_TRAIN_CSV, wells=TRAIN_WELLS) - - # 按井划分 train/val train_indices, val_indices = [], [] for i, name in enumerate(train_ds_all.names): w = parse_well_id(name) @@ -342,19 +313,19 @@ def main(): else: train_indices.append(i) - train_ds = Subset(train_ds_all, train_indices) - val_ds = Subset(train_ds_all, val_indices) + train_ds = Subset(train_ds_all, train_indices) # well1-5 + val_ds = Subset(train_ds_all, val_indices) # well6 train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=0) val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=0) print(f"训练样本数: {len(train_ds)} | 验证样本数: {len(val_ds)} | val_wells={VAL_WELLS}") - # ====== (C) 模型与优化器 ====== + # (C) 模型与优化器 model = ResNet34UNet(num_classes=NUM_CLASSES).to(DEVICE) optimizer = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY) - # ====== (D) 训练 ====== + # (D) 训练 best_val = 1e9 best_path = DATA_ROOT / "best_resnet34_unet.pth" @@ -369,10 +340,10 @@ def main(): torch.save(model.state_dict(), best_path) print(f" -> 保存最优模型: {best_path}") - # ====== (E) 生成提交（从同一个 images/ 里筛选 well7-11） ====== + # (E) 生成提交（测试目录 well7-11） out_csv = DATA_ROOT / "submission.csv" model.load_state_dict(torch.load(best_path, map_location=DEVICE)) - predict_and_make_submission(model, IMAGES_DIR, out_csv, test_wells=TEST_WELLS) + predict_and_make_submission(model, TEST_IMAGES_DIR, out_csv) if __name__ == "__main__": From 8ac09c9972a92b561c85f3b6f5e383407779f715 Mon Sep 17 00:00:00 2001 From: Haokun Date: Sat, 20 Dec 2025 00:02:59 +0100 Subject: [PATCH 03/11] add 3 new model,look README --- README.md | 90 ++++++++- deeplabv3.ipynb | 391 ++++++++++++++++++++++++++++++++++++ resnet.ipynb | 31 ++- resnet.py | 350 --------------------------------- upernet.ipynb | 511 ++++++++++++++++++++++++++++++++++++++++++++++++ 5 files changed, 1014 insertions(+), 359 deletions(-) create mode 100644 deeplabv3.ipynb delete mode 100644 resnet.py create mode 100644 upernet.ipynb diff --git a/README.md b/README.md index 14a4873..0914cf0 100644 --- a/README.md +++ b/README.md @@ -1 +1,89 @@ -# datachallenge_deep \ No newline at end of file +# datachallenge_deep + +transformer最强 +下面按“类型”把**已经用过**的模型和**还可以用**的模型列出来（只列结构名，不展开细节）。 + +--- + +## 1) U-Net 系列（Encoder–Decoder / Skip Connection） + +### 你已用 + +* **U-Net（keep272）0.4850** +* **ResNet34-U-Net 0.4858865707148295** + +### 仍可用 + +* **ResNet18/50-U-Net** +* **UNet++（Nested U-Net）** +* **Attention U-Net** +* **U-Net 3+** +* **R2U-Net / R2AttU-Net（递归残差）** +* **BCDU-Net** +* **DenseUNet** +* **U-Net + ASPP（在瓶颈加 ASPP）** + +--- + +## 2) DeepLab / ASPP 多尺度语义分割（Atrous） + +### 你已用 + +* **DeepLabV3（ResNet50 backbone）0.46073617997736044** + +### 仍可用 + +* **DeepLabV3+**（比 V3 更强的 decoder） +* **DeepLabV3-MobileNetV3**（更快） +* **LR-ASPP（轻量 ASPP）** + +--- + +## 3) Transformer / ViT 系列语义分割 + +### 你已用 + +* **SegFormer 0.649377** + +### 仍可用 + +* **Swin-UNet / Swin-Transformer segmentation** +* **Mask2Former（更强但更复杂）** +* **UPerNet + Swin / ViT backbone** +* **DPT（Dense Prediction Transformer）** + +--- + +## 4) Pyramid / FPN 多尺度融合（CNN 语义分割头） + +### 你已用 + +* **UPerNet（PSP + FPN）0.5387714111687977** + +### 仍可用 + +* **PSPNet** +* **FPN-Seg** +* **HRNet + OCR（高分辨率）** +* **PAN（Path Aggregation Network）** + +--- + +## 5) “边界/形状友好”的结构（适合细长界面） + +### 你已用 + +* （暂无） + +### 仍可用 + +* **Boundary Head（分割 + 边界分支）** +* **Distance Transform 回归 + 分割** +* **Gated-SCNN（边界引导）** + +--- + +如果你要我下一步直接落地，我建议你从**“还可用”里最值得做、且代码不会爆炸**的两类选一个： + +* **UNet++**（最贴合你任务的细长结构） +* **DeepLabV3+**（比 V3 更强，改动也不算太大） diff --git a/deeplabv3.ipynb b/deeplabv3.ipynb new file mode 100644 index 0000000..093df65 --- /dev/null +++ b/deeplabv3.ipynb @@ -0,0 +1,391 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "35e48916", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "训练样本数: 2790 | 验证样本数: 1620 | val_wells={6}\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Downloading: \"https://download.pytorch.org/models/deeplabv3_resnet50_coco-cd0a2569.pth\" to C:\\Users\\lenovo/.cache\\torch\\hub\\checkpoints\\deeplabv3_resnet50_coco-cd0a2569.pth\n", + "100%|██████████| 161M/161M [00:18<00:00, 9.19MB/s] \n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Epoch 01/5 | train_loss=0.1193 | val_loss=0.1085\n", + " -> 保存最优模型: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_deeplabv3.pth\n", + "Epoch 02/5 | train_loss=0.0891 | val_loss=0.4385\n", + "Epoch 03/5 | train_loss=0.0838 | val_loss=0.0938\n", + " -> 保存最优模型: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_deeplabv3.pth\n", + "Epoch 04/5 | train_loss=0.0797 | val_loss=0.1012\n", + "Epoch 05/5 | train_loss=0.0788 | val_loss=0.1925\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "C:\\Users\\lenovo\\AppData\\Local\\Temp\\ipykernel_3448\\751595344.py:308: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n", + " model.load_state_dict(torch.load(best_path, map_location=DEVICE))\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[OK] submission 已保存: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\submission.csv\n" + ] + } + ], + "source": [ + "# -*- coding: utf-8 -*-\n", + "\"\"\"\n", + "DeepLabV3（torchvision）分割：可直接跑（简单版）\n", + "- 训练图像目录：X_train_uDRk9z9/images（well1-6）\n", + "- 测试图像目录：X_test_xNbnvIa/images（well7-11）\n", + "- 训练标签：Y_train_T9NrBYo.csv（flatten + -1 padding）\n", + "- 验证：按井划分（例：well6 为 val，其余为 train）\n", + "- 输出：submission.csv（每行一个 patch，flatten，pad 到 160*272 用 -1）\n", + "\n", + "你只需要确认/修改：\n", + "1) DATA_ROOT 路径\n", + "2) EPOCHS / BATCH_SIZE（训练慢就先设小）\n", + "\"\"\"\n", + "\n", + "import re\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "import pandas as pd\n", + "\n", + "import torch\n", + "import torch.nn as nn\n", + "import torch.nn.functional as F\n", + "from torch.utils.data import Dataset, DataLoader, Subset\n", + "\n", + "from torchvision.models.segmentation import deeplabv3_resnet50, DeepLabV3_ResNet50_Weights\n", + "\n", + "\n", + "# =========================\n", + "# 0. 超参数与路径\n", + "# =========================\n", + "DATA_ROOT = Path(r\"C:\\Users\\lenovo\\Desktop\\deep_datachallenge\") # 改成你的真实路径\n", + "\n", + "TRAIN_IMAGES_DIR = DATA_ROOT / \"X_train_uDRk9z9\" / \"images\"\n", + "TEST_IMAGES_DIR = DATA_ROOT / \"X_test_xNbnvIa\" / \"images\"\n", + "Y_TRAIN_CSV = DATA_ROOT / \"Y_train_T9NrBYo.csv\"\n", + "\n", + "TARGET_H = 160\n", + "TARGET_W = 272\n", + "\n", + "NUM_CLASSES = 3 # 你确认 CSV 里只有 0/1/2\n", + "IGNORE_INDEX = -1 # CSV padding\n", + "\n", + "BATCH_SIZE = 8\n", + "LR = 1e-3\n", + "WEIGHT_DECAY = 1e-4\n", + "EPOCHS = 5 # 训练慢就先设 5，跑通后再加\n", + "\n", + "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", + "\n", + "\n", + "# =========================\n", + "# 1. 工具函数\n", + "# =========================\n", + "def parse_well_id(name: str) -> int:\n", + " \"\"\"从 well_1_section_0_patch_0 提取 well id=1\"\"\"\n", + " m = re.search(r\"well_(\\d+)_\", name)\n", + " return int(m.group(1)) if m else -1\n", + "\n", + "\n", + "def minmax_normalize(x: np.ndarray) -> np.ndarray:\n", + " \"\"\"min-max 归一化；NaN/inf 置 0\"\"\"\n", + " x = x.astype(np.float32)\n", + " x = np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)\n", + " x_min = float(x.min())\n", + " x_max = float(x.max())\n", + " if x_max - x_min < 1e-6:\n", + " return np.zeros_like(x, dtype=np.float32)\n", + " return (x - x_min) / (x_max - x_min)\n", + "\n", + "\n", + "def pad_to_160x272(img: np.ndarray, fill_value: float = 0.0) -> np.ndarray:\n", + " \"\"\"把 (160,160) 或 (160,272) pad 到 (160,272)\"\"\"\n", + " h, w = img.shape\n", + " assert h == TARGET_H, f\"期望高度 {TARGET_H}，但拿到 {h}\"\n", + " if w == TARGET_W:\n", + " return img\n", + " if w < TARGET_W:\n", + " out = np.full((TARGET_H, TARGET_W), fill_value, dtype=img.dtype)\n", + " out[:, :w] = img\n", + " return out\n", + " return img[:, :TARGET_W]\n", + "\n", + "\n", + "def decode_mask_from_csv_row(row_values: np.ndarray) -> np.ndarray:\n", + " \"\"\"\n", + " 从 CSV 一行恢复 mask：\n", + " - row_values: flatten + -1 padding\n", + " - 去掉 -1 后 reshape 成 (160, w)\n", + " \"\"\"\n", + " valid = row_values[row_values != IGNORE_INDEX]\n", + " assert len(valid) % TARGET_H == 0, f\"mask 有效长度 {len(valid)} 不能被 160 整除\"\n", + " w = len(valid) // TARGET_H\n", + " return valid.reshape(TARGET_H, w).astype(np.int64)\n", + "\n", + "\n", + "def pad_mask_to_160x272(mask: np.ndarray) -> np.ndarray:\n", + " \"\"\"把 (160,w) pad 到 (160,272)，pad 用 -1（ignore）\"\"\"\n", + " h, w = mask.shape\n", + " assert h == TARGET_H\n", + " if w == TARGET_W:\n", + " return mask\n", + " out = np.full((TARGET_H, TARGET_W), IGNORE_INDEX, dtype=np.int64)\n", + " out[:, :w] = mask\n", + " return out\n", + "\n", + "\n", + "# =========================\n", + "# 2. Dataset（训练/测试共用）\n", + "# =========================\n", + "class WellSegDataset(Dataset):\n", + " def __init__(self, images_dir: Path, y_csv_path: Path = None):\n", + " \"\"\"\n", + " y_csv_path=None 表示无标签（测试）\n", + " \"\"\"\n", + " self.images_dir = images_dir\n", + " self.has_label = y_csv_path is not None\n", + "\n", + " self.image_paths = sorted(images_dir.glob(\"*.npy\"))\n", + " self.names = [p.stem for p in self.image_paths]\n", + "\n", + " if self.has_label:\n", + " # CSV index 通常就是 patch 名（不含 .npy）\n", + " self.y_df = pd.read_csv(y_csv_path, index_col=0)\n", + " else:\n", + " self.y_df = None\n", + "\n", + " def __len__(self):\n", + " return len(self.image_paths)\n", + "\n", + " def __getitem__(self, idx: int):\n", + " name = self.names[idx]\n", + " img_path = self.image_paths[idx]\n", + "\n", + " img = np.load(img_path) # (160,160) or (160,272)\n", + " raw_w = img.shape[1] # 记录原始宽度（推理时裁回去）\n", + " img = minmax_normalize(img)\n", + " img = pad_to_160x272(img, fill_value=0.0)\n", + " img_t = torch.from_numpy(img).unsqueeze(0).float() # (1,160,272)\n", + "\n", + " if not self.has_label:\n", + " return {\"name\": name, \"image\": img_t, \"raw_w\": raw_w}\n", + "\n", + " row = self.y_df.loc[name].values.astype(np.int64)\n", + " mask = decode_mask_from_csv_row(row) # (160,w)\n", + " mask = pad_mask_to_160x272(mask) # (160,272)\n", + " mask_t = torch.from_numpy(mask).long()\n", + "\n", + " return {\"name\": name, \"image\": img_t, \"mask\": mask_t, \"raw_w\": raw_w}\n", + "\n", + "\n", + "# =========================\n", + "# 3. DeepLabV3（torchvision）\n", + "# =========================\n", + "class DeepLabV3(nn.Module):\n", + " def __init__(self, num_classes: int):\n", + " super().__init__()\n", + "\n", + " # 载入 DeepLabV3-ResNet50（预训练）\n", + " self.model = deeplabv3_resnet50(weights=DeepLabV3_ResNet50_Weights.DEFAULT)\n", + "\n", + " # 改分类头输出类别数\n", + " self.model.classifier[4] = nn.Conv2d(256, num_classes, kernel_size=1)\n", + "\n", + " # 输入单通道：conv1 从 3 通道改为 1 通道（权重均值初始化）\n", + " old_conv1 = self.model.backbone[\"conv1\"]\n", + " new_conv1 = nn.Conv2d(\n", + " in_channels=1,\n", + " out_channels=old_conv1.out_channels,\n", + " kernel_size=old_conv1.kernel_size,\n", + " stride=old_conv1.stride,\n", + " padding=old_conv1.padding,\n", + " bias=False,\n", + " )\n", + " with torch.no_grad():\n", + " new_conv1.weight[:] = old_conv1.weight.mean(dim=1, keepdim=True)\n", + " self.model.backbone[\"conv1\"] = new_conv1\n", + "\n", + " def forward(self, x):\n", + " out = self.model(x)[\"out\"] # (B,C,H,W)\n", + " # 保证输出回到 (160,272)\n", + " out = F.interpolate(out, size=(TARGET_H, TARGET_W), mode=\"bilinear\", align_corners=False)\n", + " return out\n", + "\n", + "\n", + "# =========================\n", + "# 4. 训练与验证（最简单）\n", + "# =========================\n", + "def train_one_epoch(model, loader, optimizer):\n", + " model.train()\n", + " total_loss = 0.0\n", + " for batch in loader:\n", + " x = batch[\"image\"].to(DEVICE)\n", + " y = batch[\"mask\"].to(DEVICE)\n", + "\n", + " logits = model(x)\n", + " loss = F.cross_entropy(logits, y, ignore_index=IGNORE_INDEX)\n", + "\n", + " optimizer.zero_grad()\n", + " loss.backward()\n", + " optimizer.step()\n", + "\n", + " total_loss += loss.item() * x.size(0)\n", + "\n", + " return total_loss / len(loader.dataset)\n", + "\n", + "\n", + "@torch.no_grad()\n", + "def valid_one_epoch(model, loader):\n", + " model.eval()\n", + " total_loss = 0.0\n", + " for batch in loader:\n", + " x = batch[\"image\"].to(DEVICE)\n", + " y = batch[\"mask\"].to(DEVICE)\n", + "\n", + " logits = model(x)\n", + " loss = F.cross_entropy(logits, y, ignore_index=IGNORE_INDEX)\n", + " total_loss += loss.item() * x.size(0)\n", + "\n", + " return total_loss / len(loader.dataset)\n", + "\n", + "\n", + "# =========================\n", + "# 5. 推理并生成提交 CSV\n", + "# =========================\n", + "@torch.no_grad()\n", + "def predict_and_make_submission(model, test_images_dir: Path, out_csv_path: Path):\n", + " \"\"\"\n", + " 对 test_images_dir 全部 npy 预测并生成提交 CSV\n", + " - 每行：一个 patch\n", + " - 长度：160*272\n", + " - 如果原始宽度 < 272，剩余用 -1 padding\n", + " \"\"\"\n", + " model.eval()\n", + "\n", + " test_ds = WellSegDataset(test_images_dir, y_csv_path=None)\n", + " test_loader = DataLoader(test_ds, batch_size=1, shuffle=False, num_workers=0)\n", + "\n", + " preds_dict = {}\n", + "\n", + " for batch in test_loader:\n", + " name = batch[\"name\"][0]\n", + " raw_w = int(batch[\"raw_w\"][0])\n", + " x = batch[\"image\"].to(DEVICE)\n", + "\n", + " logits = model(x)\n", + " pred_full = torch.argmax(logits, dim=1).squeeze(0).cpu().numpy().astype(np.int64) # (160,272)\n", + "\n", + " pred = pred_full[:, :raw_w] # 裁回原始宽度\n", + " if raw_w < TARGET_W:\n", + " padded = np.full((TARGET_H * TARGET_W,), IGNORE_INDEX, dtype=np.int64)\n", + " padded[: TARGET_H * raw_w] = pred.flatten()\n", + " preds_dict[name] = padded\n", + " else:\n", + " preds_dict[name] = pred.flatten()\n", + "\n", + " sub = pd.DataFrame(preds_dict, dtype=\"int64\").T\n", + " sub.to_csv(out_csv_path)\n", + " print(f\"[OK] submission 已保存: {out_csv_path}\")\n", + "\n", + "\n", + "# =========================\n", + "# 6. 主函数：训练(井1-5) + 验证(井6) + 预测test(井7-11目录)\n", + "# =========================\n", + "def main():\n", + " # (A) 构建训练集（well1-6）\n", + " train_ds_all = WellSegDataset(TRAIN_IMAGES_DIR, Y_TRAIN_CSV)\n", + "\n", + " # (B) 按井划分 train/val：well6 做验证\n", + " VAL_WELLS = {6}\n", + " train_indices, val_indices = [], []\n", + " for i, name in enumerate(train_ds_all.names):\n", + " w = parse_well_id(name)\n", + " if w in VAL_WELLS:\n", + " val_indices.append(i)\n", + " else:\n", + " train_indices.append(i)\n", + "\n", + " train_ds = Subset(train_ds_all, train_indices) # well1-5\n", + " val_ds = Subset(train_ds_all, val_indices) # well6\n", + "\n", + " train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)\n", + " val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)\n", + "\n", + " print(f\"训练样本数: {len(train_ds)} | 验证样本数: {len(val_ds)} | val_wells={VAL_WELLS}\")\n", + "\n", + " # (C) 模型与优化器\n", + " model = DeepLabV3(num_classes=NUM_CLASSES).to(DEVICE)\n", + " optimizer = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)\n", + "\n", + " # (D) 训练\n", + " best_val = 1e9\n", + " best_path = DATA_ROOT / \"best_deeplabv3.pth\"\n", + "\n", + " for epoch in range(1, EPOCHS + 1):\n", + " tr_loss = train_one_epoch(model, train_loader, optimizer)\n", + " va_loss = valid_one_epoch(model, val_loader)\n", + "\n", + " print(f\"Epoch {epoch:02d}/{EPOCHS} | train_loss={tr_loss:.4f} | val_loss={va_loss:.4f}\")\n", + "\n", + " if va_loss < best_val:\n", + " best_val = va_loss\n", + " torch.save(model.state_dict(), best_path)\n", + " print(f\" -> 保存最优模型: {best_path}\")\n", + "\n", + " # (E) 生成提交（测试目录 well7-11）\n", + " out_csv = DATA_ROOT / \"submission.csv\"\n", + " model.load_state_dict(torch.load(best_path, map_location=DEVICE))\n", + " predict_and_make_submission(model, TEST_IMAGES_DIR, out_csv)\n", + "\n", + "\n", + "if __name__ == \"__main__\":\n", + " main()\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "deep-torch", + "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.10.19" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/resnet.ipynb b/resnet.ipynb index e02f1de..73839af 100644 --- a/resnet.ipynb +++ b/resnet.ipynb @@ -3,7 +3,7 @@ { "cell_type": "code", "execution_count": 1, - "id": "51ec4b7d", + "id": "02e3cbc1", "metadata": {}, "outputs": [ { @@ -11,20 +11,35 @@ "output_type": "stream", "text": [ "训练样本数: 2790 | 验证样本数: 1620 | val_wells={6}\n", - "Epoch 01/5 | train_loss=0.1464 | val_loss=0.0906\n", + "Epoch 01/20 | train_loss=0.1348 | val_loss=0.1186\n", " -> 保存最优模型: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_resnet34_unet.pth\n", - "Epoch 02/5 | train_loss=0.0871 | val_loss=0.0923\n", - "Epoch 03/5 | train_loss=0.0784 | val_loss=0.0879\n", + "Epoch 02/20 | train_loss=0.0833 | val_loss=0.4633\n", + "Epoch 03/20 | train_loss=0.0795 | val_loss=0.1263\n", + "Epoch 04/20 | train_loss=0.0759 | val_loss=0.0764\n", " -> 保存最优模型: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_resnet34_unet.pth\n", - "Epoch 04/5 | train_loss=0.0767 | val_loss=0.0893\n", - "Epoch 05/5 | train_loss=0.0745 | val_loss=0.1063\n" + "Epoch 05/20 | train_loss=0.0762 | val_loss=0.0862\n", + "Epoch 06/20 | train_loss=0.0723 | val_loss=0.2669\n", + "Epoch 07/20 | train_loss=0.0719 | val_loss=0.0981\n", + "Epoch 08/20 | train_loss=0.0703 | val_loss=0.4578\n", + "Epoch 09/20 | train_loss=0.0725 | val_loss=0.0795\n", + "Epoch 10/20 | train_loss=0.0686 | val_loss=0.0902\n", + "Epoch 11/20 | train_loss=0.0678 | val_loss=0.0798\n", + "Epoch 12/20 | train_loss=0.0682 | val_loss=0.1063\n", + "Epoch 13/20 | train_loss=0.0683 | val_loss=0.0870\n", + "Epoch 14/20 | train_loss=0.0655 | val_loss=0.0828\n", + "Epoch 15/20 | train_loss=0.0654 | val_loss=0.0841\n", + "Epoch 16/20 | train_loss=0.0668 | val_loss=0.0817\n", + "Epoch 17/20 | train_loss=0.0639 | val_loss=0.0867\n", + "Epoch 18/20 | train_loss=0.0636 | val_loss=0.0818\n", + "Epoch 19/20 | train_loss=0.0634 | val_loss=0.0777\n", + "Epoch 20/20 | train_loss=0.0633 | val_loss=0.0790\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ - "C:\\Users\\lenovo\\AppData\\Local\\Temp\\ipykernel_12820\\130010335.py:345: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n", + "C:\\Users\\lenovo\\AppData\\Local\\Temp\\ipykernel_11248\\3190494213.py:345: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.\n", " model.load_state_dict(torch.load(best_path, map_location=DEVICE))\n" ] }, @@ -79,7 +94,7 @@ "BATCH_SIZE = 8\n", "LR = 1e-3\n", "WEIGHT_DECAY = 1e-4\n", - "EPOCHS = 5 # 训练慢就先设 5，跑通后再加\n", + "EPOCHS = 20 # 训练慢就先设 5，跑通后再加\n", "\n", "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", "\n", diff --git a/resnet.py b/resnet.py deleted file mode 100644 index e2b7975..0000000 --- a/resnet.py +++ /dev/null @@ -1,350 +0,0 @@ -# -*- coding: utf-8 -*- -""" -ResNet34-UNet（简化版，可直接跑） -- 训练图像目录：X_train_uDRk9z9/images（well1-6） -- 测试图像目录：X_test_xNbnvIa/images（well7-11） -- 训练标签：Y_train_T9NrBYo.csv（flatten + -1 padding） -- 验证：从训练集中按井划分（例：well6 为 val，其余为 train） -- 输出：submission.csv（每行一个 patch，flatten，pad 到 160*272 用 -1） -""" - -import re -from pathlib import Path - -import numpy as np -import pandas as pd - -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch.utils.data import Dataset, DataLoader, Subset - -from torchvision.models import resnet34, ResNet34_Weights - - -# ========================= -# 0. 超参数与路径 -# ========================= -DATA_ROOT = Path(r"C:\Users\lenovo\Desktop\deep_datachallenge") # 改成你的真实路径 - -TRAIN_IMAGES_DIR = DATA_ROOT / "X_train_uDRk9z9" / "images" -TEST_IMAGES_DIR = DATA_ROOT / "X_test_xNbnvIa" / "images" -Y_TRAIN_CSV = DATA_ROOT / "Y_train_T9NrBYo.csv" - -TARGET_H = 160 -TARGET_W = 272 - -NUM_CLASSES = 3 # 你确认 CSV 里只有 0/1/2 -IGNORE_INDEX = -1 # CSV padding - -BATCH_SIZE = 8 -LR = 1e-3 -WEIGHT_DECAY = 1e-4 -EPOCHS = 5 # 训练慢就先设 5，跑通后再加 - -DEVICE = "cuda" if torch.cuda.is_available() else "cpu" - - -# ========================= -# 1. 工具函数 -# ========================= -def parse_well_id(name: str) -> int: - """从 well_1_section_0_patch_0 提取 well id=1""" - m = re.search(r"well_(\d+)_", name) - return int(m.group(1)) if m else -1 - - -def minmax_normalize(x: np.ndarray) -> np.ndarray: - """min-max 归一化；NaN/inf 置 0""" - x = x.astype(np.float32) - x = np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0) - x_min = float(x.min()) - x_max = float(x.max()) - if x_max - x_min < 1e-6: - return np.zeros_like(x, dtype=np.float32) - return (x - x_min) / (x_max - x_min) - - -def pad_to_160x272(img: np.ndarray, fill_value: float = 0.0) -> np.ndarray: - """把 (160,160) 或 (160,272) pad 到 (160,272)""" - h, w = img.shape - assert h == TARGET_H, f"期望高度 {TARGET_H}，但拿到 {h}" - if w == TARGET_W: - return img - if w < TARGET_W: - out = np.full((TARGET_H, TARGET_W), fill_value, dtype=img.dtype) - out[:, :w] = img - return out - return img[:, :TARGET_W] - - -def decode_mask_from_csv_row(row_values: np.ndarray) -> np.ndarray: - """ - 从 CSV 一行恢复 mask： - - row_values: flatten + -1 padding - - 去掉 -1 后 reshape 成 (160, w) - """ - valid = row_values[row_values != IGNORE_INDEX] - assert len(valid) % TARGET_H == 0, f"mask 有效长度 {len(valid)} 不能被 160 整除" - w = len(valid) // TARGET_H - return valid.reshape(TARGET_H, w).astype(np.int64) - - -def pad_mask_to_160x272(mask: np.ndarray) -> np.ndarray: - """把 (160,w) pad 到 (160,272)，pad 用 -1（ignore）""" - h, w = mask.shape - assert h == TARGET_H - if w == TARGET_W: - return mask - out = np.full((TARGET_H, TARGET_W), IGNORE_INDEX, dtype=np.int64) - out[:, :w] = mask - return out - - -# ========================= -# 2. Dataset（训练/测试共用） -# ========================= -class WellSegDataset(Dataset): - def __init__(self, images_dir: Path, y_csv_path: Path = None): - """ - y_csv_path=None 表示无标签（测试） - """ - self.images_dir = images_dir - self.has_label = y_csv_path is not None - - self.image_paths = sorted(images_dir.glob("*.npy")) - self.names = [p.stem for p in self.image_paths] - - if self.has_label: - # CSV index 通常就是 patch 名（不含 .npy） - self.y_df = pd.read_csv(y_csv_path, index_col=0) - else: - self.y_df = None - - def __len__(self): - return len(self.image_paths) - - def __getitem__(self, idx: int): - name = self.names[idx] - img_path = self.image_paths[idx] - - img = np.load(img_path) # (160,160) or (160,272) - raw_w = img.shape[1] # 记录原始宽度（推理时裁回去） - img = minmax_normalize(img) - img = pad_to_160x272(img, fill_value=0.0) - img_t = torch.from_numpy(img).unsqueeze(0).float() # (1,160,272) - - if not self.has_label: - return {"name": name, "image": img_t, "raw_w": raw_w} - - row = self.y_df.loc[name].values.astype(np.int64) - mask = decode_mask_from_csv_row(row) # (160,w) - mask = pad_mask_to_160x272(mask) # (160,272) - mask_t = torch.from_numpy(mask).long() - - return {"name": name, "image": img_t, "mask": mask_t, "raw_w": raw_w} - - -# ========================= -# 3. ResNet34-UNet（简化实现） -# ========================= -class ConvRelu(nn.Module): - def __init__(self, in_ch, out_ch): - super().__init__() - self.block = nn.Sequential( - nn.Conv2d(in_ch, out_ch, 3, padding=1, bias=False), - nn.BatchNorm2d(out_ch), - nn.ReLU(inplace=True), - ) - - def forward(self, x): - return self.block(x) - - -class UpBlock(nn.Module): - def __init__(self, in_ch, skip_ch, out_ch): - super().__init__() - self.conv1 = ConvRelu(in_ch + skip_ch, out_ch) - self.conv2 = ConvRelu(out_ch, out_ch) - - def forward(self, x, skip): - x = F.interpolate(x, size=skip.shape[-2:], mode="bilinear", align_corners=False) - x = torch.cat([x, skip], dim=1) - x = self.conv1(x) - x = self.conv2(x) - return x - - -class ResNet34UNet(nn.Module): - def __init__(self, num_classes: int): - super().__init__() - backbone = resnet34(weights=ResNet34_Weights.IMAGENET1K_V1) - - # 输入单通道：第一层卷积改成 1 通道（用原权重均值初始化） - old_conv1 = backbone.conv1 - new_conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False) - with torch.no_grad(): - new_conv1.weight[:] = old_conv1.weight.mean(dim=1, keepdim=True) - backbone.conv1 = new_conv1 - - self.enc0 = nn.Sequential(backbone.conv1, backbone.bn1, backbone.relu) # /2 - self.pool0 = backbone.maxpool # /4 - self.enc1 = backbone.layer1 # /4 - self.enc2 = backbone.layer2 # /8 - self.enc3 = backbone.layer3 # /16 - self.enc4 = backbone.layer4 # /32 - - self.center = nn.Sequential(ConvRelu(512, 512), ConvRelu(512, 512)) - self.up4 = UpBlock(512, 256, 256) - self.up3 = UpBlock(256, 128, 128) - self.up2 = UpBlock(128, 64, 64) - self.up1 = UpBlock(64, 64, 64) - - self.head = nn.Conv2d(64, num_classes, kernel_size=1) - - def forward(self, x): - e0 = self.enc0(x) - e1 = self.enc1(self.pool0(e0)) - e2 = self.enc2(e1) - e3 = self.enc3(e2) - e4 = self.enc4(e3) - - c = self.center(e4) - d4 = self.up4(c, e3) - d3 = self.up3(d4, e2) - d2 = self.up2(d3, e1) - d1 = self.up1(d2, e0) - - out = self.head(d1) - out = F.interpolate(out, size=(TARGET_H, TARGET_W), mode="bilinear", align_corners=False) - return out - - -# ========================= -# 4. 训练与验证（最简单） -# ========================= -def train_one_epoch(model, loader, optimizer): - model.train() - total_loss = 0.0 - for batch in loader: - x = batch["image"].to(DEVICE) - y = batch["mask"].to(DEVICE) - - logits = model(x) - loss = F.cross_entropy(logits, y, ignore_index=IGNORE_INDEX) - - optimizer.zero_grad() - loss.backward() - optimizer.step() - - total_loss += loss.item() * x.size(0) - - return total_loss / len(loader.dataset) - - -@torch.no_grad() -def valid_one_epoch(model, loader): - model.eval() - total_loss = 0.0 - for batch in loader: - x = batch["image"].to(DEVICE) - y = batch["mask"].to(DEVICE) - - logits = model(x) - loss = F.cross_entropy(logits, y, ignore_index=IGNORE_INDEX) - total_loss += loss.item() * x.size(0) - - return total_loss / len(loader.dataset) - - -# ========================= -# 5. 推理并生成提交 CSV -# ========================= -@torch.no_grad() -def predict_and_make_submission(model, test_images_dir: Path, out_csv_path: Path): - """ - 对 test_images_dir 全部 npy 预测并生成提交 CSV - - 每行：一个 patch - - 长度：160*272 - - 如果原始宽度 < 272，剩余用 -1 padding - """ - model.eval() - - test_ds = WellSegDataset(test_images_dir, y_csv_path=None) - test_loader = DataLoader(test_ds, batch_size=1, shuffle=False, num_workers=0) - - preds_dict = {} - - for batch in test_loader: - name = batch["name"][0] - raw_w = int(batch["raw_w"][0]) - x = batch["image"].to(DEVICE) - - logits = model(x) - pred_full = torch.argmax(logits, dim=1).squeeze(0).cpu().numpy().astype(np.int64) # (160,272) - - pred = pred_full[:, :raw_w] # 裁回原始宽度 - if raw_w < TARGET_W: - padded = np.full((TARGET_H * TARGET_W,), IGNORE_INDEX, dtype=np.int64) - padded[: TARGET_H * raw_w] = pred.flatten() - preds_dict[name] = padded - else: - preds_dict[name] = pred.flatten() - - sub = pd.DataFrame(preds_dict, dtype="int64").T - sub.to_csv(out_csv_path) - print(f"[OK] submission 已保存: {out_csv_path}") - - -# ========================= -# 6. 主函数：训练(井1-5) + 验证(井6) + 预测test(井7-11目录) -# ========================= -def main(): - # (A) 构建训练集（well1-6） - train_ds_all = WellSegDataset(TRAIN_IMAGES_DIR, Y_TRAIN_CSV) - - # (B) 按井划分 train/val：well6 做验证 - VAL_WELLS = {6} - train_indices, val_indices = [], [] - for i, name in enumerate(train_ds_all.names): - w = parse_well_id(name) - if w in VAL_WELLS: - val_indices.append(i) - else: - train_indices.append(i) - - train_ds = Subset(train_ds_all, train_indices) # well1-5 - val_ds = Subset(train_ds_all, val_indices) # well6 - - train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=0) - val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=0) - - print(f"训练样本数: {len(train_ds)} | 验证样本数: {len(val_ds)} | val_wells={VAL_WELLS}") - - # (C) 模型与优化器 - model = ResNet34UNet(num_classes=NUM_CLASSES).to(DEVICE) - optimizer = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY) - - # (D) 训练 - best_val = 1e9 - best_path = DATA_ROOT / "best_resnet34_unet.pth" - - for epoch in range(1, EPOCHS + 1): - tr_loss = train_one_epoch(model, train_loader, optimizer) - va_loss = valid_one_epoch(model, val_loader) - - print(f"Epoch {epoch:02d}/{EPOCHS} | train_loss={tr_loss:.4f} | val_loss={va_loss:.4f}") - - if va_loss < best_val: - best_val = va_loss - torch.save(model.state_dict(), best_path) - print(f" -> 保存最优模型: {best_path}") - - # (E) 生成提交（测试目录 well7-11） - out_csv = DATA_ROOT / "submission.csv" - model.load_state_dict(torch.load(best_path, map_location=DEVICE)) - predict_and_make_submission(model, TEST_IMAGES_DIR, out_csv) - - -if __name__ == "__main__": - main() diff --git a/upernet.ipynb b/upernet.ipynb new file mode 100644 index 0000000..e47a758 --- /dev/null +++ b/upernet.ipynb @@ -0,0 +1,511 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 2, + "id": "6aaeb914", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "训练样本数: 2790 | 验证样本数: 1620 | val_wells={6}\n", + "DEVICE: cuda\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Downloading: \"https://download.pytorch.org/models/resnet50-11ad3fa6.pth\" to C:\\Users\\lenovo/.cache\\torch\\hub\\checkpoints\\resnet50-11ad3fa6.pth\n", + "100%|██████████| 97.8M/97.8M [00:09<00:00, 10.3MB/s]\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Epoch 01/20 | train_loss=0.1318 | val_loss=0.1120\n", + " -> 保存最优模型: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_upernet.pth\n", + "Epoch 02/20 | train_loss=0.0745 | val_loss=0.0900\n", + " -> 保存最优模型: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_upernet.pth\n", + "Epoch 03/20 | train_loss=0.0664 | val_loss=0.0940\n", + "Epoch 04/20 | train_loss=0.0616 | val_loss=0.1032\n", + "Epoch 05/20 | train_loss=0.0576 | val_loss=0.1283\n", + "Epoch 06/20 | train_loss=0.0555 | val_loss=0.0903\n", + "Epoch 07/20 | train_loss=0.0538 | val_loss=0.0936\n", + "Epoch 08/20 | train_loss=0.0509 | val_loss=0.1156\n", + "Epoch 09/20 | train_loss=0.0481 | val_loss=0.0915\n", + "Epoch 10/20 | train_loss=0.0446 | val_loss=0.0991\n", + "Epoch 11/20 | train_loss=0.0449 | val_loss=0.1008\n", + "Epoch 12/20 | train_loss=0.0416 | val_loss=0.1002\n", + "Epoch 13/20 | train_loss=0.0402 | val_loss=0.1123\n", + "Epoch 14/20 | train_loss=0.0419 | val_loss=0.1042\n", + "Epoch 15/20 | train_loss=0.0380 | val_loss=0.1052\n", + "Epoch 16/20 | train_loss=0.0379 | val_loss=0.1320\n", + "Epoch 17/20 | train_loss=0.0359 | val_loss=0.1031\n", + "Epoch 18/20 | train_loss=0.0340 | val_loss=0.1163\n", + "Epoch 19/20 | train_loss=0.0341 | val_loss=0.1079\n", + "Epoch 20/20 | train_loss=0.0381 | val_loss=0.1154\n", + "[OK] submission 已保存: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\submission.csv\n" + ] + } + ], + "source": [ + "# -*- coding: utf-8 -*-\n", + "\"\"\"\n", + "UPerNet（ResNet50 backbone，torchvision预训练）完整可跑版本（简单实现）\n", + "- 训练图像目录：X_train_uDRk9z9/images（well1-6）\n", + "- 测试图像目录：X_test_xNbnvIa/images（well7-11）\n", + "- 训练标签：Y_train_T9NrBYo.csv（flatten + -1 padding）\n", + "- 验证：按井划分（例：well6 为 val，其余为 train）\n", + "- 输出：submission.csv（每行一个 patch，flatten，pad 到 160*272 用 -1）\n", + "\n", + "你只需要确认/修改：\n", + "1) DATA_ROOT 路径\n", + "2) EPOCHS / BATCH_SIZE（训练慢就先设小）\n", + "\"\"\"\n", + "\n", + "import re\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "import pandas as pd\n", + "\n", + "import torch\n", + "import torch.nn as nn\n", + "import torch.nn.functional as F\n", + "from torch.utils.data import Dataset, DataLoader, Subset\n", + "\n", + "from torchvision.models import resnet50, ResNet50_Weights\n", + "\n", + "\n", + "# =========================\n", + "# 0. 超参数与路径\n", + "# =========================\n", + "DATA_ROOT = Path(r\"C:\\Users\\lenovo\\Desktop\\deep_datachallenge\") # 改成你的真实路径\n", + "\n", + "TRAIN_IMAGES_DIR = DATA_ROOT / \"X_train_uDRk9z9\" / \"images\"\n", + "TEST_IMAGES_DIR = DATA_ROOT / \"X_test_xNbnvIa\" / \"images\"\n", + "Y_TRAIN_CSV = DATA_ROOT / \"Y_train_T9NrBYo.csv\"\n", + "\n", + "TARGET_H = 160\n", + "TARGET_W = 272\n", + "\n", + "NUM_CLASSES = 3 # CSV 里只有 0/1/2\n", + "IGNORE_INDEX = -1 # CSV padding\n", + "\n", + "BATCH_SIZE = 4 # UPerNet 比较吃显存，4060(8GB)建议先 2~4\n", + "LR = 1e-4 # UPerNet + 预训练 backbone 通常 lr 小一点更稳\n", + "WEIGHT_DECAY = 1e-4\n", + "EPOCHS = 20 # 训练慢就先设 20，跑通后再加\n", + "\n", + "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", + "\n", + "\n", + "# =========================\n", + "# 1. 工具函数\n", + "# =========================\n", + "def parse_well_id(name: str) -> int:\n", + " \"\"\"从 well_1_section_0_patch_0 提取 well id=1\"\"\"\n", + " m = re.search(r\"well_(\\d+)_\", name)\n", + " return int(m.group(1)) if m else -1\n", + "\n", + "\n", + "def minmax_normalize(x: np.ndarray) -> np.ndarray:\n", + " \"\"\"min-max 归一化；NaN/inf 置 0\"\"\"\n", + " x = x.astype(np.float32)\n", + " x = np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)\n", + " x_min = float(x.min())\n", + " x_max = float(x.max())\n", + " if x_max - x_min < 1e-6:\n", + " return np.zeros_like(x, dtype=np.float32)\n", + " return (x - x_min) / (x_max - x_min)\n", + "\n", + "\n", + "def pad_to_160x272(img: np.ndarray, fill_value: float = 0.0) -> np.ndarray:\n", + " \"\"\"把 (160,160) 或 (160,272) pad 到 (160,272)\"\"\"\n", + " h, w = img.shape\n", + " assert h == TARGET_H, f\"期望高度 {TARGET_H}，但拿到 {h}\"\n", + " if w == TARGET_W:\n", + " return img\n", + " if w < TARGET_W:\n", + " out = np.full((TARGET_H, TARGET_W), fill_value, dtype=img.dtype)\n", + " out[:, :w] = img\n", + " return out\n", + " return img[:, :TARGET_W]\n", + "\n", + "\n", + "def decode_mask_from_csv_row(row_values: np.ndarray) -> np.ndarray:\n", + " \"\"\"\n", + " 从 CSV 一行恢复 mask：\n", + " - row_values: flatten + -1 padding\n", + " - 去掉 -1 后 reshape 成 (160, w)\n", + " \"\"\"\n", + " valid = row_values[row_values != IGNORE_INDEX]\n", + " assert len(valid) % TARGET_H == 0, f\"mask 有效长度 {len(valid)} 不能被 160 整除\"\n", + " w = len(valid) // TARGET_H\n", + " return valid.reshape(TARGET_H, w).astype(np.int64)\n", + "\n", + "\n", + "def pad_mask_to_160x272(mask: np.ndarray) -> np.ndarray:\n", + " \"\"\"把 (160,w) pad 到 (160,272)，pad 用 -1（ignore）\"\"\"\n", + " h, w = mask.shape\n", + " assert h == TARGET_H\n", + " if w == TARGET_W:\n", + " return mask\n", + " out = np.full((TARGET_H, TARGET_W), IGNORE_INDEX, dtype=np.int64)\n", + " out[:, :w] = mask\n", + " return out\n", + "\n", + "\n", + "# =========================\n", + "# 2. Dataset（训练/测试共用）\n", + "# =========================\n", + "class WellSegDataset(Dataset):\n", + " def __init__(self, images_dir: Path, y_csv_path: Path = None):\n", + " \"\"\"\n", + " y_csv_path=None 表示无标签（测试）\n", + " \"\"\"\n", + " self.images_dir = images_dir\n", + " self.has_label = y_csv_path is not None\n", + "\n", + " self.image_paths = sorted(images_dir.glob(\"*.npy\"))\n", + " self.names = [p.stem for p in self.image_paths]\n", + "\n", + " if self.has_label:\n", + " self.y_df = pd.read_csv(y_csv_path, index_col=0)\n", + " else:\n", + " self.y_df = None\n", + "\n", + " def __len__(self):\n", + " return len(self.image_paths)\n", + "\n", + " def __getitem__(self, idx: int):\n", + " name = self.names[idx]\n", + " img_path = self.image_paths[idx]\n", + "\n", + " img = np.load(img_path) # (160,160) or (160,272)\n", + " raw_w = img.shape[1] # 推理时裁回原始宽度\n", + " img = minmax_normalize(img)\n", + " img = pad_to_160x272(img, fill_value=0.0)\n", + " img_t = torch.from_numpy(img).unsqueeze(0).float() # (1,160,272)\n", + "\n", + " if not self.has_label:\n", + " return {\"name\": name, \"image\": img_t, \"raw_w\": raw_w}\n", + "\n", + " row = self.y_df.loc[name].values.astype(np.int64)\n", + " mask = decode_mask_from_csv_row(row) # (160,w)\n", + " mask = pad_mask_to_160x272(mask) # (160,272)\n", + " mask_t = torch.from_numpy(mask).long()\n", + "\n", + " return {\"name\": name, \"image\": img_t, \"mask\": mask_t, \"raw_w\": raw_w}\n", + "\n", + "\n", + "# =========================\n", + "# 3. UPerNet Head（简化实现）\n", + "# =========================\n", + "class ConvBNReLU(nn.Module):\n", + " def __init__(self, in_ch, out_ch, k=3, p=1):\n", + " super().__init__()\n", + " self.block = nn.Sequential(\n", + " nn.Conv2d(in_ch, out_ch, kernel_size=k, padding=p, bias=False),\n", + " nn.BatchNorm2d(out_ch),\n", + " nn.ReLU(inplace=True),\n", + " )\n", + "\n", + " def forward(self, x):\n", + " return self.block(x)\n", + "\n", + "\n", + "class PSPModule(nn.Module):\n", + " \"\"\"\n", + " PSP: 对最高层特征做多尺度池化，再concat回去\n", + " \"\"\"\n", + " def __init__(self, in_ch, out_ch=256, pool_sizes=(1, 2, 3, 6)):\n", + " super().__init__()\n", + " self.stages = nn.ModuleList()\n", + " for ps in pool_sizes:\n", + " self.stages.append(nn.Sequential(\n", + " nn.AdaptiveAvgPool2d((ps, ps)),\n", + " nn.Conv2d(in_ch, out_ch, kernel_size=1, bias=False),\n", + " nn.BatchNorm2d(out_ch),\n", + " nn.ReLU(inplace=True),\n", + " ))\n", + " self.bottleneck = ConvBNReLU(in_ch + len(pool_sizes) * out_ch, out_ch, k=3, p=1)\n", + "\n", + " def forward(self, x):\n", + " h, w = x.shape[-2:]\n", + " priors = [x]\n", + " for stage in self.stages:\n", + " y = stage(x)\n", + " y = F.interpolate(y, size=(h, w), mode=\"bilinear\", align_corners=False)\n", + " priors.append(y)\n", + " x = torch.cat(priors, dim=1)\n", + " return self.bottleneck(x) # (B, out_ch, h, w)\n", + "\n", + "\n", + "class UPerHead(nn.Module):\n", + " \"\"\"\n", + " UPerNet Head = PSP + FPN\n", + " 输入特征：c2(1/4), c3(1/8), c4(1/16), c5(1/32)\n", + " \"\"\"\n", + " def __init__(self, in_channels=(256, 512, 1024, 2048), fpn_dim=256, num_classes=3):\n", + " super().__init__()\n", + "\n", + " c2, c3, c4, c5 = in_channels\n", + "\n", + " # PSP 在最高层 c5 上\n", + " self.psp = PSPModule(c5, out_ch=fpn_dim)\n", + "\n", + " # lateral 1x1：把各层通道统一到 fpn_dim\n", + " self.lateral_c2 = nn.Conv2d(c2, fpn_dim, kernel_size=1, bias=False)\n", + " self.lateral_c3 = nn.Conv2d(c3, fpn_dim, kernel_size=1, bias=False)\n", + " self.lateral_c4 = nn.Conv2d(c4, fpn_dim, kernel_size=1, bias=False)\n", + "\n", + " # FPN 3x3：融合后再平滑\n", + " self.fpn_c2 = ConvBNReLU(fpn_dim, fpn_dim, k=3, p=1)\n", + " self.fpn_c3 = ConvBNReLU(fpn_dim, fpn_dim, k=3, p=1)\n", + " self.fpn_c4 = ConvBNReLU(fpn_dim, fpn_dim, k=3, p=1)\n", + " self.fpn_c5 = ConvBNReLU(fpn_dim, fpn_dim, k=3, p=1)\n", + "\n", + " # 融合输出：concat(4层) -> 256\n", + " self.fuse = ConvBNReLU(fpn_dim * 4, fpn_dim, k=3, p=1)\n", + " self.cls = nn.Conv2d(fpn_dim, num_classes, kernel_size=1)\n", + "\n", + " def forward(self, c2, c3, c4, c5):\n", + " # PSP 输出 p5（仍在 1/32 分辨率）\n", + " p5 = self.psp(c5) # (B,256, H/32, W/32)\n", + "\n", + " # lateral\n", + " p4 = self.lateral_c4(c4) # (B,256, H/16, W/16)\n", + " p3 = self.lateral_c3(c3) # (B,256, H/8, W/8 )\n", + " p2 = self.lateral_c2(c2) # (B,256, H/4, W/4 )\n", + "\n", + " # top-down 融合\n", + " p4 = p4 + F.interpolate(p5, size=p4.shape[-2:], mode=\"bilinear\", align_corners=False)\n", + " p3 = p3 + F.interpolate(p4, size=p3.shape[-2:], mode=\"bilinear\", align_corners=False)\n", + " p2 = p2 + F.interpolate(p3, size=p2.shape[-2:], mode=\"bilinear\", align_corners=False)\n", + "\n", + " # 平滑\n", + " p5 = self.fpn_c5(p5)\n", + " p4 = self.fpn_c4(p4)\n", + " p3 = self.fpn_c3(p3)\n", + " p2 = self.fpn_c2(p2)\n", + "\n", + " # 全部上采样到 p2 尺寸，再 concat\n", + " p5_u = F.interpolate(p5, size=p2.shape[-2:], mode=\"bilinear\", align_corners=False)\n", + " p4_u = F.interpolate(p4, size=p2.shape[-2:], mode=\"bilinear\", align_corners=False)\n", + " p3_u = F.interpolate(p3, size=p2.shape[-2:], mode=\"bilinear\", align_corners=False)\n", + "\n", + " x = torch.cat([p2, p3_u, p4_u, p5_u], dim=1)\n", + " x = self.fuse(x)\n", + " logits = self.cls(x) # (B, num_classes, H/4, W/4)\n", + " return logits\n", + "\n", + "\n", + "# =========================\n", + "# 4. UPerNet（ResNet50 backbone）\n", + "# =========================\n", + "class UPerNet(nn.Module):\n", + " def __init__(self, num_classes: int):\n", + " super().__init__()\n", + "\n", + " # ResNet50 预训练\n", + " backbone = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)\n", + "\n", + " # 输入单通道：conv1 从 3 通道改为 1 通道（用均值初始化）\n", + " old_conv1 = backbone.conv1\n", + " new_conv1 = nn.Conv2d(\n", + " in_channels=1,\n", + " out_channels=old_conv1.out_channels,\n", + " kernel_size=old_conv1.kernel_size,\n", + " stride=old_conv1.stride,\n", + " padding=old_conv1.padding,\n", + " bias=False,\n", + " )\n", + " with torch.no_grad():\n", + " new_conv1.weight[:] = old_conv1.weight.mean(dim=1, keepdim=True)\n", + " backbone.conv1 = new_conv1\n", + "\n", + " self.backbone = backbone\n", + "\n", + " # UPerHead：resnet50 的输出通道 (layer1..4) = 256,512,1024,2048\n", + " self.head = UPerHead(in_channels=(256, 512, 1024, 2048), fpn_dim=256, num_classes=num_classes)\n", + "\n", + " def forward(self, x):\n", + " # ResNet 前几层\n", + " x = self.backbone.conv1(x)\n", + " x = self.backbone.bn1(x)\n", + " x = self.backbone.relu(x)\n", + " x = self.backbone.maxpool(x)\n", + "\n", + " # 取各层特征\n", + " c2 = self.backbone.layer1(x) # 1/4\n", + " c3 = self.backbone.layer2(c2) # 1/8\n", + " c4 = self.backbone.layer3(c3) # 1/16\n", + " c5 = self.backbone.layer4(c4) # 1/32\n", + "\n", + " logits_1_4 = self.head(c2, c3, c4, c5)\n", + "\n", + " # 上采样到原图尺寸 (160,272)\n", + " logits = F.interpolate(logits_1_4, size=(TARGET_H, TARGET_W), mode=\"bilinear\", align_corners=False)\n", + " return logits\n", + "\n", + "\n", + "# =========================\n", + "# 5. 训练与验证（最简单）\n", + "# =========================\n", + "def train_one_epoch(model, loader, optimizer):\n", + " model.train()\n", + " total_loss = 0.0\n", + " for batch in loader:\n", + " x = batch[\"image\"].to(DEVICE)\n", + " y = batch[\"mask\"].to(DEVICE)\n", + "\n", + " logits = model(x)\n", + " loss = F.cross_entropy(logits, y, ignore_index=IGNORE_INDEX)\n", + "\n", + " optimizer.zero_grad()\n", + " loss.backward()\n", + " optimizer.step()\n", + "\n", + " total_loss += loss.item() * x.size(0)\n", + "\n", + " return total_loss / len(loader.dataset)\n", + "\n", + "\n", + "@torch.no_grad()\n", + "def valid_one_epoch(model, loader):\n", + " model.eval()\n", + " total_loss = 0.0\n", + " for batch in loader:\n", + " x = batch[\"image\"].to(DEVICE)\n", + " y = batch[\"mask\"].to(DEVICE)\n", + "\n", + " logits = model(x)\n", + " loss = F.cross_entropy(logits, y, ignore_index=IGNORE_INDEX)\n", + " total_loss += loss.item() * x.size(0)\n", + "\n", + " return total_loss / len(loader.dataset)\n", + "\n", + "\n", + "# =========================\n", + "# 6. 推理并生成提交 CSV\n", + "# =========================\n", + "@torch.no_grad()\n", + "def predict_and_make_submission(model, test_images_dir: Path, out_csv_path: Path):\n", + " \"\"\"\n", + " 对 test_images_dir 全部 npy 预测并生成提交 CSV\n", + " - 每行：一个 patch\n", + " - 长度：160*272\n", + " - 如果原始宽度 < 272，剩余用 -1 padding\n", + " \"\"\"\n", + " model.eval()\n", + "\n", + " test_ds = WellSegDataset(test_images_dir, y_csv_path=None)\n", + " test_loader = DataLoader(test_ds, batch_size=1, shuffle=False, num_workers=0)\n", + "\n", + " preds_dict = {}\n", + "\n", + " for batch in test_loader:\n", + " name = batch[\"name\"][0]\n", + " raw_w = int(batch[\"raw_w\"][0])\n", + " x = batch[\"image\"].to(DEVICE)\n", + "\n", + " logits = model(x)\n", + " pred_full = torch.argmax(logits, dim=1).squeeze(0).cpu().numpy().astype(np.int64) # (160,272)\n", + "\n", + " pred = pred_full[:, :raw_w] # 裁回原始宽度\n", + " if raw_w < TARGET_W:\n", + " padded = np.full((TARGET_H * TARGET_W,), IGNORE_INDEX, dtype=np.int64)\n", + " padded[: TARGET_H * raw_w] = pred.flatten()\n", + " preds_dict[name] = padded\n", + " else:\n", + " preds_dict[name] = pred.flatten()\n", + "\n", + " sub = pd.DataFrame(preds_dict, dtype=\"int64\").T\n", + " sub.to_csv(out_csv_path)\n", + " print(f\"[OK] submission 已保存: {out_csv_path}\")\n", + "\n", + "\n", + "# =========================\n", + "# 7. 主函数：训练(井1-5) + 验证(井6) + 预测test(井7-11目录)\n", + "# =========================\n", + "def main():\n", + " # (A) 构建训练集（well1-6）\n", + " train_ds_all = WellSegDataset(TRAIN_IMAGES_DIR, Y_TRAIN_CSV)\n", + "\n", + " # (B) 按井划分 train/val：well6 做验证\n", + " VAL_WELLS = {6}\n", + " train_indices, val_indices = [], []\n", + " for i, name in enumerate(train_ds_all.names):\n", + " w = parse_well_id(name)\n", + " if w in VAL_WELLS:\n", + " val_indices.append(i)\n", + " else:\n", + " train_indices.append(i)\n", + "\n", + " train_ds = Subset(train_ds_all, train_indices) # well1-5\n", + " val_ds = Subset(train_ds_all, val_indices) # well6\n", + "\n", + " train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)\n", + " val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)\n", + "\n", + " print(f\"训练样本数: {len(train_ds)} | 验证样本数: {len(val_ds)} | val_wells={VAL_WELLS}\")\n", + " print(f\"DEVICE: {DEVICE}\")\n", + "\n", + " # (C) 模型与优化器\n", + " model = UPerNet(num_classes=NUM_CLASSES).to(DEVICE)\n", + " optimizer = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)\n", + "\n", + " # (D) 训练\n", + " best_val = 1e9\n", + " best_path = DATA_ROOT / \"best_upernet.pth\"\n", + "\n", + " for epoch in range(1, EPOCHS + 1):\n", + " tr_loss = train_one_epoch(model, train_loader, optimizer)\n", + " va_loss = valid_one_epoch(model, val_loader)\n", + "\n", + " print(f\"Epoch {epoch:02d}/{EPOCHS} | train_loss={tr_loss:.4f} | val_loss={va_loss:.4f}\")\n", + "\n", + " if va_loss < best_val:\n", + " best_val = va_loss\n", + " torch.save(model.state_dict(), best_path)\n", + " print(f\" -> 保存最优模型: {best_path}\")\n", + "\n", + " # (E) 生成提交（测试目录 well7-11）\n", + " out_csv = DATA_ROOT / \"submission.csv\"\n", + " state_dict = torch.load(best_path, map_location=DEVICE, weights_only=True)\n", + " model.load_state_dict(state_dict)\n", + " predict_and_make_submission(model, TEST_IMAGES_DIR, out_csv)\n", + "\n", + "\n", + "if __name__ == \"__main__\":\n", + " main()\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "deep-torch", + "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.10.19" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From b16f94772bba1e5e340293aca03ad6229c3a714c Mon Sep 17 00:00:00 2001 From: HOOK290 Date: Sat, 20 Dec 2025 00:04:34 +0100 Subject: [PATCH 04/11] 3 new model :scores in README Removed suggestions for next steps in the README. --- README.md | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/README.md b/README.md index 0914cf0..ef8fb69 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,8 @@ # datachallenge_deep transformer最强 + + 下面按“类型”把**已经用过**的模型和**还可以用**的模型列出来（只列结构名，不展开细节）。 --- @@ -83,7 +85,3 @@ transformer最强 --- -如果你要我下一步直接落地，我建议你从**“还可用”里最值得做、且代码不会爆炸**的两类选一个： - -* **UNet++**（最贴合你任务的细长结构） -* **DeepLabV3+**（比 V3 更强，改动也不算太大） From ddc1669ed33c372a1c1f9b624167076ec4d40c42 Mon Sep 17 00:00:00 2001 From: Haokun Date: Sun, 21 Dec 2025 00:07:00 +0100 Subject: [PATCH 05/11] 8 scores in README --- DPT.ipynb | 402 ++++++++++++++++++++++ README.md | 78 ++++- UPerNet_Swin_ViT_backbone.ipynb | 330 ++++++++++++++++++ attention_u_net.ipynb | 501 +++++++++++++++++++++++++++ mask2former.ipynb | 592 ++++++++++++++++++++++++++++++++ swin-unet.ipynb | 434 +++++++++++++++++++++++ 6 files changed, 2329 insertions(+), 8 deletions(-) create mode 100644 DPT.ipynb create mode 100644 UPerNet_Swin_ViT_backbone.ipynb create mode 100644 attention_u_net.ipynb create mode 100644 mask2former.ipynb create mode 100644 swin-unet.ipynb diff --git a/DPT.ipynb b/DPT.ipynb new file mode 100644 index 0000000..6adc67a --- /dev/null +++ b/DPT.ipynb @@ -0,0 +1,402 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "d12efdc1", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "c:\\Users\\lenovo\\anaconda3\\envs\\deep-torch\\lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", + " from .autonotebook import tqdm as notebook_tqdm\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "DEVICE: cuda\n", + "Train: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\X_train_uDRk9z9\\images\n", + "Test : C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\X_test_xNbnvIa\\images\n", + "Train samples: 2790 | Val samples: 1620\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "c:\\Users\\lenovo\\anaconda3\\envs\\deep-torch\\lib\\site-packages\\huggingface_hub\\file_download.py:143: UserWarning: `huggingface_hub` cache-system uses symlinks by default to efficiently store duplicated files but your machine does not support them in C:\\Users\\lenovo\\.cache\\huggingface\\hub\\models--timm--vit_base_patch16_224.augreg2_in21k_ft_in1k. Caching files will still work but in a degraded version that might require more space on your disk. This warning can be disabled by setting the `HF_HUB_DISABLE_SYMLINKS_WARNING` environment variable. For more details, see https://huggingface.co/docs/huggingface_hub/how-to-cache#limitations.\n", + "To support symlinks on Windows, you either need to activate Developer Mode or to run Python as an administrator. In order to activate developer mode, see this article: https://docs.microsoft.com/en-us/windows/apps/get-started/enable-your-device-for-development\n", + " warnings.warn(message)\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Epoch 01/20 | train_loss=0.1310 | val_loss=0.0810\n", + " -> best saved: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_dpt.pth\n", + "Epoch 02/20 | train_loss=0.0776 | val_loss=0.0825\n", + "Epoch 03/20 | train_loss=0.0669 | val_loss=0.0937\n", + "Epoch 04/20 | train_loss=0.0643 | val_loss=0.0849\n", + "Epoch 05/20 | train_loss=0.0606 | val_loss=0.1009\n", + "Epoch 06/20 | train_loss=0.0598 | val_loss=0.0824\n", + "Epoch 07/20 | train_loss=0.0532 | val_loss=0.0823\n", + "Epoch 08/20 | train_loss=0.0509 | val_loss=0.0885\n", + "Epoch 09/20 | train_loss=0.0520 | val_loss=0.1033\n", + "Epoch 10/20 | train_loss=0.0472 | val_loss=0.0876\n", + "Epoch 11/20 | train_loss=0.0444 | val_loss=0.0891\n", + "Epoch 12/20 | train_loss=0.0443 | val_loss=0.0897\n", + "Epoch 13/20 | train_loss=0.0406 | val_loss=0.0911\n", + "Epoch 14/20 | train_loss=0.0420 | val_loss=0.0902\n", + "Epoch 15/20 | train_loss=0.0395 | val_loss=0.0965\n", + "Epoch 16/20 | train_loss=0.0346 | val_loss=0.0980\n", + "Epoch 17/20 | train_loss=0.0315 | val_loss=0.1069\n", + "Epoch 18/20 | train_loss=0.0291 | val_loss=0.1004\n", + "Epoch 19/20 | train_loss=0.0279 | val_loss=0.1170\n", + "Epoch 20/20 | train_loss=0.0275 | val_loss=0.1033\n", + "[OK] submission saved to: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\submission.csv\n" + ] + } + ], + "source": [ + "# -*- coding: utf-8 -*-\n", + "\"\"\"\n", + "DPT (Dense Prediction Transformer) - Full runnable code (NO cv2)\n", + "Using segmentation-models-pytorch (smp) built-in DPT.\n", + "\n", + "Why this works:\n", + "- DPT uses a ViT-like backbone, usually requires fixed input size (e.g., 224x224).\n", + "- We follow the same safe pipeline as before:\n", + " (160, w) -> pad to (160,272) -> resize to (224,224) -> model\n", + " model output (224,224) -> resize back to (160,272) -> crop to raw_w\n", + "- submission.csv must be 160*272 flattened with -1 padding\n", + "\n", + "Data:\n", + "- Train images: Desktop/deep_datachallenge/X_train_uDRk9z9/images (well1-6)\n", + "- Train labels: Desktop/deep_datachallenge/Y_train_T9NrBYo.csv\n", + "- Test images: Desktop/deep_datachallenge/X_test_xNbnvIa/images (well7-11)\n", + "\n", + "Split:\n", + "- Train: well1-5\n", + "- Val: well6\n", + "\"\"\"\n", + "\n", + "import re\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "import pandas as pd\n", + "\n", + "import torch\n", + "import torch.nn.functional as F\n", + "from torch.utils.data import Dataset, DataLoader, Subset\n", + "\n", + "import segmentation_models_pytorch as smp\n", + "\n", + "\n", + "# =========================\n", + "# 0) Config\n", + "# =========================\n", + "DATA_ROOT = Path(r\"C:\\Users\\lenovo\\Desktop\\deep_datachallenge\")\n", + "\n", + "TRAIN_IMG_DIR = DATA_ROOT / \"X_train_uDRk9z9\" / \"images\"\n", + "TEST_IMG_DIR = DATA_ROOT / \"X_test_xNbnvIa\" / \"images\"\n", + "Y_TRAIN_CSV = DATA_ROOT / \"Y_train_T9NrBYo.csv\"\n", + "\n", + "# submission resolution (fixed by challenge)\n", + "H_SUB, W_SUB = 160, 272\n", + "\n", + "# model input resolution (ViT/DPT often expects fixed size)\n", + "H_MODEL, W_MODEL = 224, 224\n", + "\n", + "NUM_CLASSES = 3\n", + "IGNORE_INDEX = -1\n", + "\n", + "BATCH_SIZE = 4 # 4060(8GB): start with 2~4\n", + "LR = 1e-4\n", + "WEIGHT_DECAY = 1e-4\n", + "EPOCHS = 20\n", + "\n", + "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", + "\n", + "\n", + "# =========================\n", + "# 1) Utils\n", + "# =========================\n", + "def parse_well_id(name: str) -> int:\n", + " \"\"\"Extract well id from 'well_6_section_...' -> 6\"\"\"\n", + " m = re.search(r\"well_(\\d+)_\", name)\n", + " return int(m.group(1)) if m else -1\n", + "\n", + "\n", + "def minmax_norm(x: np.ndarray) -> np.ndarray:\n", + " \"\"\"Normalize to [0,1], handle NaN/Inf.\"\"\"\n", + " x = x.astype(np.float32)\n", + " x = np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)\n", + " mn, mx = float(x.min()), float(x.max())\n", + " if mx - mn < 1e-6:\n", + " return np.zeros_like(x, dtype=np.float32)\n", + " return (x - mn) / (mx - mn)\n", + "\n", + "\n", + "def pad_to_160x272(img: np.ndarray) -> np.ndarray:\n", + " \"\"\"(160,w)->(160,272), pad right with 0.\"\"\"\n", + " h, w = img.shape\n", + " assert h == H_SUB, f\"Expected height {H_SUB}, got {h}\"\n", + " out = np.zeros((H_SUB, W_SUB), dtype=img.dtype)\n", + " out[:, :min(w, W_SUB)] = img[:, :min(w, W_SUB)]\n", + " return out\n", + "\n", + "\n", + "def decode_mask_from_csv_row(row: np.ndarray) -> np.ndarray:\n", + " \"\"\"\n", + " CSV row is flattened mask with -1 padding.\n", + " Remove -1 then reshape to (160, w).\n", + " \"\"\"\n", + " valid = row[row != IGNORE_INDEX]\n", + " assert len(valid) % H_SUB == 0, f\"Valid mask length {len(valid)} not divisible by {H_SUB}\"\n", + " w = len(valid) // H_SUB\n", + " return valid.reshape(H_SUB, w).astype(np.int64)\n", + "\n", + "\n", + "def pad_mask_to_160x272(mask: np.ndarray) -> np.ndarray:\n", + " \"\"\"(160,w)->(160,272) with -1 padding on the right.\"\"\"\n", + " h, w = mask.shape\n", + " assert h == H_SUB\n", + " out = np.full((H_SUB, W_SUB), IGNORE_INDEX, dtype=np.int64)\n", + " out[:, :min(w, W_SUB)] = mask[:, :min(w, W_SUB)]\n", + " return out\n", + "\n", + "\n", + "def resize_img_np_to_model(img_160x272: np.ndarray) -> np.ndarray:\n", + " \"\"\"(160,272)->(224,224), bilinear\"\"\"\n", + " x = torch.from_numpy(img_160x272).unsqueeze(0).unsqueeze(0).float() # (1,1,160,272)\n", + " x = F.interpolate(x, size=(H_MODEL, W_MODEL), mode=\"bilinear\", align_corners=False)\n", + " return x.squeeze(0).squeeze(0).numpy()\n", + "\n", + "\n", + "def resize_mask_np_to_model(mask_160x272: np.ndarray) -> np.ndarray:\n", + " \"\"\"(160,272)->(224,224), nearest (for labels)\"\"\"\n", + " x = torch.from_numpy(mask_160x272).unsqueeze(0).unsqueeze(0).float() # (1,1,160,272)\n", + " x = F.interpolate(x, size=(H_MODEL, W_MODEL), mode=\"nearest\")\n", + " return x.squeeze(0).squeeze(0).long().numpy()\n", + "\n", + "\n", + "def resize_pred_np_to_sub(pred_224x224: np.ndarray) -> np.ndarray:\n", + " \"\"\"(224,224)->(160,272), nearest\"\"\"\n", + " x = torch.from_numpy(pred_224x224).unsqueeze(0).unsqueeze(0).float()\n", + " x = F.interpolate(x, size=(H_SUB, W_SUB), mode=\"nearest\")\n", + " return x.squeeze(0).squeeze(0).long().numpy()\n", + "\n", + "\n", + "# =========================\n", + "# 2) Dataset\n", + "# =========================\n", + "class WellSegDataset(Dataset):\n", + " def __init__(self, images_dir: Path, y_csv_path: Path = None):\n", + " self.images_dir = images_dir\n", + " self.paths = sorted(images_dir.glob(\"*.npy\"))\n", + " self.names = [p.stem for p in self.paths]\n", + "\n", + " self.has_label = y_csv_path is not None\n", + " self.y_df = pd.read_csv(y_csv_path, index_col=0) if self.has_label else None\n", + "\n", + " def __len__(self):\n", + " return len(self.paths)\n", + "\n", + " def __getitem__(self, idx: int):\n", + " name = self.names[idx]\n", + " img = np.load(self.paths[idx]) # (160,160) or (160,272)\n", + " raw_w = int(img.shape[1])\n", + "\n", + " img = minmax_norm(img)\n", + " img_160x272 = pad_to_160x272(img)\n", + " img_224 = resize_img_np_to_model(img_160x272)\n", + " x = torch.from_numpy(img_224).unsqueeze(0).float() # (1,224,224)\n", + "\n", + " if not self.has_label:\n", + " return {\"name\": name, \"image\": x, \"raw_w\": raw_w}\n", + "\n", + " row = self.y_df.loc[name].values.astype(np.int64)\n", + " mask = decode_mask_from_csv_row(row) # (160,w)\n", + " mask_160x272 = pad_mask_to_160x272(mask) # (160,272)\n", + " mask_224 = resize_mask_np_to_model(mask_160x272) # (224,224)\n", + " y = torch.from_numpy(mask_224).long()\n", + "\n", + " return {\"name\": name, \"image\": x, \"mask\": y, \"raw_w\": raw_w}\n", + "\n", + "\n", + "# =========================\n", + "# 3) Model: DPT (smp)\n", + "# =========================\n", + "def build_dpt(num_classes: int) -> torch.nn.Module:\n", + " \"\"\"\n", + " SMP has DPT model. You choose a ViT / transformer encoder from timm.\n", + "\n", + " Common safe encoders (usually available):\n", + " - \"tu-vit_base_patch16_224\"\n", + " - \"vit_base_patch16_224\"\n", + " If pretrained weights download fails, set encoder_weights=None.\n", + " \"\"\"\n", + " encoder_name = \"tu-vit_base_patch16_224\" # you can also try \"vit_base_patch16_224\"\n", + " encoder_weights = \"imagenet\" # if this fails, set to None\n", + "\n", + " model = smp.DPT(\n", + " encoder_name=encoder_name,\n", + " encoder_weights=encoder_weights,\n", + " in_channels=1,\n", + " classes=num_classes,\n", + " activation=None,\n", + " # If you want to try variable sizes (if supported by encoder):\n", + " # dynamic_img_size=True,\n", + " )\n", + " return model\n", + "\n", + "\n", + "# =========================\n", + "# 4) Train / Val\n", + "# =========================\n", + "def train_one_epoch(model, loader, optimizer):\n", + " model.train()\n", + " total = 0.0\n", + "\n", + " for b in loader:\n", + " x = b[\"image\"].to(DEVICE) # (B,1,224,224)\n", + " y = b[\"mask\"].to(DEVICE) # (B,224,224)\n", + "\n", + " logits = model(x) # (B,C,224,224)\n", + " loss = F.cross_entropy(logits, y, ignore_index=IGNORE_INDEX)\n", + "\n", + " optimizer.zero_grad()\n", + " loss.backward()\n", + " optimizer.step()\n", + "\n", + " total += float(loss.item()) * x.size(0)\n", + "\n", + " return total / len(loader.dataset)\n", + "\n", + "\n", + "@torch.no_grad()\n", + "def eval_one_epoch(model, loader):\n", + " model.eval()\n", + " total = 0.0\n", + "\n", + " for b in loader:\n", + " x = b[\"image\"].to(DEVICE)\n", + " y = b[\"mask\"].to(DEVICE)\n", + " logits = model(x)\n", + " loss = F.cross_entropy(logits, y, ignore_index=IGNORE_INDEX)\n", + " total += float(loss.item()) * x.size(0)\n", + "\n", + " return total / len(loader.dataset)\n", + "\n", + "\n", + "# =========================\n", + "# 5) Inference & submission\n", + "# =========================\n", + "@torch.no_grad()\n", + "def predict_and_make_submission(model, test_images_dir: Path, out_csv: Path):\n", + " model.eval()\n", + "\n", + " test_ds = WellSegDataset(test_images_dir, y_csv_path=None)\n", + " test_loader = DataLoader(test_ds, batch_size=1, shuffle=False, num_workers=0)\n", + "\n", + " preds = {}\n", + "\n", + " for b in test_loader:\n", + " name = b[\"name\"][0]\n", + " raw_w = int(b[\"raw_w\"][0])\n", + " x = b[\"image\"].to(DEVICE) # (1,1,224,224)\n", + "\n", + " logits = model(x) # (1,C,224,224)\n", + " pred224 = torch.argmax(logits, dim=1).squeeze(0).cpu().numpy().astype(np.int64) # (224,224)\n", + "\n", + " pred160x272 = resize_pred_np_to_sub(pred224) # (160,272)\n", + " pred160 = pred160x272[:, :raw_w] # (160,raw_w)\n", + "\n", + " flat = np.full((H_SUB * W_SUB,), IGNORE_INDEX, dtype=np.int64)\n", + " flat[: H_SUB * raw_w] = pred160.flatten()\n", + " preds[name] = flat\n", + "\n", + " pd.DataFrame(preds, dtype=\"int64\").T.to_csv(out_csv)\n", + " print(f\"[OK] submission saved to: {out_csv}\")\n", + "\n", + "\n", + "# =========================\n", + "# 6) Main\n", + "# =========================\n", + "def main():\n", + " print(\"DEVICE:\", DEVICE)\n", + " print(\"Train:\", TRAIN_IMG_DIR)\n", + " print(\"Test :\", TEST_IMG_DIR)\n", + "\n", + " ds_all = WellSegDataset(TRAIN_IMG_DIR, Y_TRAIN_CSV)\n", + "\n", + " train_idx, val_idx = [], []\n", + " for i, n in enumerate(ds_all.names):\n", + " if parse_well_id(n) == 6:\n", + " val_idx.append(i)\n", + " else:\n", + " train_idx.append(i)\n", + "\n", + " train_ds = Subset(ds_all, train_idx) # well1-5\n", + " val_ds = Subset(ds_all, val_idx) # well6\n", + "\n", + " train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)\n", + " val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)\n", + "\n", + " print(f\"Train samples: {len(train_ds)} | Val samples: {len(val_ds)}\")\n", + "\n", + " model = build_dpt(NUM_CLASSES).to(DEVICE)\n", + " optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)\n", + "\n", + " best_val = 1e9\n", + " best_path = DATA_ROOT / \"best_dpt.pth\"\n", + "\n", + " for epoch in range(1, EPOCHS + 1):\n", + " tr = train_one_epoch(model, train_loader, optimizer)\n", + " va = eval_one_epoch(model, val_loader)\n", + " print(f\"Epoch {epoch:02d}/{EPOCHS} | train_loss={tr:.4f} | val_loss={va:.4f}\")\n", + "\n", + " if va < best_val:\n", + " best_val = va\n", + " torch.save(model.state_dict(), best_path)\n", + " print(\" -> best saved:\", best_path)\n", + "\n", + " # test inference\n", + " model.load_state_dict(torch.load(best_path, map_location=DEVICE, weights_only=True))\n", + " out_csv = DATA_ROOT / \"submission.csv\"\n", + " predict_and_make_submission(model, TEST_IMG_DIR, out_csv)\n", + "\n", + "\n", + "if __name__ == \"__main__\":\n", + " main()\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "deep-torch", + "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.10.19" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/README.md b/README.md index ef8fb69..042d30f 100644 --- a/README.md +++ b/README.md @@ -11,14 +11,15 @@ transformer最强 ### 你已用 -* **U-Net（keep272）0.4850** +* **SegNet / segnet_gen 0.62125** +* **Attention-U-Net 0.5137568332816138** * **ResNet34-U-Net 0.4858865707148295** +* **U-Net（keep272）0.4850** ### 仍可用 * **ResNet18/50-U-Net** * **UNet++（Nested U-Net）** -* **Attention U-Net** * **U-Net 3+** * **R2U-Net / R2AttU-Net（递归残差）** * **BCDU-Net** @@ -46,13 +47,13 @@ transformer最强 ### 你已用 * **SegFormer 0.649377** +* **Mask2Former（更强但更复杂） 0.6441724128817745** +* **UPerNet + Swin / ViT backbone 0.6402961815634437** +* **DPT（Dense Prediction Transformer）0.6234547207117457** +* **Swin-UNet / Swin-Transformer segmentation 0.5850994834704769** ### 仍可用 -* **Swin-UNet / Swin-Transformer segmentation** -* **Mask2Former（更强但更复杂）** -* **UPerNet + Swin / ViT backbone** -* **DPT（Dense Prediction Transformer）** --- @@ -66,7 +67,7 @@ transformer最强 * **PSPNet** * **FPN-Seg** -* **HRNet + OCR（高分辨率）** +* **HRNet + OCR（高分辨率）**3 * **PAN（Path Aggregation Network）** --- @@ -81,7 +82,68 @@ transformer最强 * **Boundary Head（分割 + 边界分支）** * **Distance Transform 回归 + 分割** -* **Gated-SCNN（边界引导）** +* **Gated-SCNN（边界引导）**4 + +--- + +### 👉 结论（不拐弯）： + +> **你现在继续“堆更复杂 Transformer 结构”，收益已经开始下降了。** + +原因不是模型不强，而是： + +* 数据量有限（几千 patch） +* 标签噪声 + 细长结构 +* 评价是 **IoU（对边界极其敏感）** + +在这种任务里，**SegFormer 已经吃满了“Transformer 红利”**。 + +--- + +## 🎯 现在最有可能再涨分的 3 条路（按性价比排序） + +### 🥇 路线 A（最推荐）：**SegFormer + 强化训练策略** + +这是**最可能把 0.649 → 0.67+** 的方式。 + +你可以加： + +1. **Dice / Tversky / Lovász-Softmax loss（替代纯 CE）** +2. **TTA（左右翻转 + 轻微 scale）** +3. **连通域 / morphology 后处理（修断裂）** + +👉 不换模型，只“榨干 SegFormer” + +--- + +### 🥈 路线 B：**Mask2Former（现在可跑，但要“比赛版”配置）** + +如果你坚持用它，**一定要注意**： + +* batch size = 1 或 2 +* 冻结 backbone 前几层 +* 更长 warmup +* 不然很容易过拟合 well1–5 + +👉 **上限高，但不稳** + +--- + +### 🥉 路线 C：**UPerNet + Swin + 边界辅助** + +单纯 UPerNet+Swin 已经证明不够，你需要： + +* 加 **Boundary Head** +* 或 Distance Transform loss + +👉 复杂度高，回报不确定 + +--- + +## 🚀 我建议你现在这样走（最理性） + +> **把 SegFormer 当主模型冲榜** +> Mask2Former 作为“你已经会了的备选方案” --- diff --git a/UPerNet_Swin_ViT_backbone.ipynb b/UPerNet_Swin_ViT_backbone.ipynb new file mode 100644 index 0000000..fb66ab5 --- /dev/null +++ b/UPerNet_Swin_ViT_backbone.ipynb @@ -0,0 +1,330 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 4, + "id": "1d87bb20", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "DEVICE: cuda\n", + "Epoch 01/20 | train=0.1124 | val=0.0792\n", + " -> best model saved\n", + "Epoch 02/20 | train=0.0740 | val=0.0748\n", + " -> best model saved\n", + "Epoch 03/20 | train=0.0674 | val=0.0782\n", + "Epoch 04/20 | train=0.0624 | val=0.0876\n", + "Epoch 05/20 | train=0.0597 | val=0.0759\n", + "Epoch 06/20 | train=0.0535 | val=0.0836\n", + "Epoch 07/20 | train=0.0491 | val=0.0806\n", + "Epoch 08/20 | train=0.0463 | val=0.0891\n", + "Epoch 09/20 | train=0.0423 | val=0.0952\n", + "Epoch 10/20 | train=0.0395 | val=0.1013\n", + "Epoch 11/20 | train=0.0362 | val=0.0972\n", + "Epoch 12/20 | train=0.0334 | val=0.1049\n", + "Epoch 13/20 | train=0.0334 | val=0.1016\n", + "Epoch 14/20 | train=0.0295 | val=0.1128\n", + "Epoch 15/20 | train=0.0279 | val=0.1291\n", + "Epoch 16/20 | train=0.0273 | val=0.1160\n", + "Epoch 17/20 | train=0.0250 | val=0.1186\n", + "Epoch 18/20 | train=0.0289 | val=0.1225\n", + "Epoch 19/20 | train=0.0227 | val=0.1314\n", + "Epoch 20/20 | train=0.0234 | val=0.1459\n", + "[OK] submission saved to C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\submission.csv\n" + ] + } + ], + "source": [ + "# -*- coding: utf-8 -*-\n", + "\"\"\"\n", + "UPerNet + Swin Transformer backbone\n", + "FINAL VERSION (NO cv2, PyTorch interpolate only)\n", + "\n", + "Key points:\n", + "- Swin requires 224x224 input\n", + "- Original data: (160,160) or (160,272)\n", + "- Strategy:\n", + " (160,w) -> pad to (160,272) -> resize to (224,224) -> model\n", + " output (224,224) -> resize back to (160,272) -> crop raw_w\n", + "- Submission remains 160*272 with -1 padding\n", + "\"\"\"\n", + "\n", + "import re\n", + "from pathlib import Path\n", + "import numpy as np\n", + "import pandas as pd\n", + "\n", + "import torch\n", + "import torch.nn.functional as F\n", + "from torch.utils.data import Dataset, DataLoader, Subset\n", + "\n", + "import segmentation_models_pytorch as smp\n", + "\n", + "\n", + "# =========================\n", + "# 0) Config\n", + "# =========================\n", + "DATA_ROOT = Path(r\"C:\\Users\\lenovo\\Desktop\\deep_datachallenge\")\n", + "\n", + "TRAIN_IMG_DIR = DATA_ROOT / \"X_train_uDRk9z9\" / \"images\"\n", + "TEST_IMG_DIR = DATA_ROOT / \"X_test_xNbnvIa\" / \"images\"\n", + "Y_TRAIN_CSV = DATA_ROOT / \"Y_train_T9NrBYo.csv\"\n", + "\n", + "# submission size (fixed by challenge)\n", + "H_SUB, W_SUB = 160, 272\n", + "\n", + "# model input size (Swin requirement)\n", + "H_MODEL, W_MODEL = 224, 224\n", + "\n", + "NUM_CLASSES = 3\n", + "IGNORE_INDEX = -1\n", + "\n", + "BATCH_SIZE = 4 # RTX 4060 (8GB): 2~4 recommended\n", + "LR = 1e-4\n", + "WEIGHT_DECAY = 1e-4\n", + "EPOCHS = 20\n", + "\n", + "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", + "\n", + "\n", + "# =========================\n", + "# 1) Utils\n", + "# =========================\n", + "def parse_well_id(name: str) -> int:\n", + " m = re.search(r\"well_(\\d+)_\", name)\n", + " return int(m.group(1)) if m else -1\n", + "\n", + "\n", + "def minmax_norm(x: np.ndarray) -> np.ndarray:\n", + " x = x.astype(np.float32)\n", + " x = np.nan_to_num(x)\n", + " mn, mx = x.min(), x.max()\n", + " if mx - mn < 1e-6:\n", + " return np.zeros_like(x)\n", + " return (x - mn) / (mx - mn)\n", + "\n", + "\n", + "def resize_image_to_224(img: np.ndarray) -> np.ndarray:\n", + " \"\"\"\n", + " img: (160, w)\n", + " -> pad to (160,272)\n", + " -> resize to (224,224) using bilinear\n", + " \"\"\"\n", + " padded = np.zeros((H_SUB, W_SUB), dtype=img.dtype)\n", + " h, w = img.shape\n", + " padded[:, :w] = img\n", + "\n", + " x = torch.from_numpy(padded).unsqueeze(0).unsqueeze(0).float() # (1,1,160,272)\n", + " x = F.interpolate(x, size=(H_MODEL, W_MODEL), mode=\"bilinear\", align_corners=False)\n", + " return x.squeeze(0).squeeze(0).numpy()\n", + "\n", + "\n", + "def decode_mask(row: np.ndarray) -> np.ndarray:\n", + " valid = row[row != IGNORE_INDEX]\n", + " w = len(valid) // H_SUB\n", + " return valid.reshape(H_SUB, w).astype(np.int64)\n", + "\n", + "\n", + "def resize_mask_to_224(mask: np.ndarray) -> np.ndarray:\n", + " \"\"\"\n", + " mask: (160, w)\n", + " -> pad to (160,272) with -1\n", + " -> resize to (224,224) using nearest\n", + " \"\"\"\n", + " padded = np.full((H_SUB, W_SUB), IGNORE_INDEX, dtype=np.int64)\n", + " h, w = mask.shape\n", + " padded[:, :w] = mask\n", + "\n", + " x = torch.from_numpy(padded).unsqueeze(0).unsqueeze(0).float()\n", + " x = F.interpolate(x, size=(H_MODEL, W_MODEL), mode=\"nearest\")\n", + " return x.squeeze(0).squeeze(0).long().numpy()\n", + "\n", + "\n", + "# =========================\n", + "# 2) Dataset\n", + "# =========================\n", + "class WellDataset(Dataset):\n", + " def __init__(self, img_dir: Path, y_csv: Path = None):\n", + " self.paths = sorted(img_dir.glob(\"*.npy\"))\n", + " self.names = [p.stem for p in self.paths]\n", + " self.has_label = y_csv is not None\n", + " self.y_df = pd.read_csv(y_csv, index_col=0) if self.has_label else None\n", + "\n", + " def __len__(self):\n", + " return len(self.paths)\n", + "\n", + " def __getitem__(self, idx):\n", + " name = self.names[idx]\n", + " img = np.load(self.paths[idx]) # (160, w)\n", + " raw_w = img.shape[1]\n", + "\n", + " img = minmax_norm(img)\n", + " img224 = resize_image_to_224(img)\n", + " x = torch.from_numpy(img224).unsqueeze(0).float() # (1,224,224)\n", + "\n", + " if not self.has_label:\n", + " return {\"name\": name, \"image\": x, \"raw_w\": raw_w}\n", + "\n", + " row = self.y_df.loc[name].values.astype(np.int64)\n", + " mask = decode_mask(row)\n", + " mask224 = resize_mask_to_224(mask)\n", + " y = torch.from_numpy(mask224).long() # (224,224)\n", + "\n", + " return {\"name\": name, \"image\": x, \"mask\": y, \"raw_w\": raw_w}\n", + "\n", + "\n", + "# =========================\n", + "# 3) Model\n", + "# =========================\n", + "def build_model():\n", + " return smp.UPerNet(\n", + " encoder_name=\"tu-swin_small_patch4_window7_224\",\n", + " encoder_weights=\"imagenet\",\n", + " in_channels=1,\n", + " classes=NUM_CLASSES,\n", + " activation=None,\n", + " )\n", + "\n", + "\n", + "# =========================\n", + "# 4) Train / Validate\n", + "# =========================\n", + "def train_one_epoch(model, loader, optimizer):\n", + " model.train()\n", + " total = 0.0\n", + "\n", + " for b in loader:\n", + " x = b[\"image\"].to(DEVICE)\n", + " y = b[\"mask\"].to(DEVICE)\n", + "\n", + " logits = model(x)\n", + " loss = F.cross_entropy(logits, y, ignore_index=IGNORE_INDEX)\n", + "\n", + " optimizer.zero_grad()\n", + " loss.backward()\n", + " optimizer.step()\n", + "\n", + " total += loss.item() * x.size(0)\n", + "\n", + " return total / len(loader.dataset)\n", + "\n", + "\n", + "@torch.no_grad()\n", + "def eval_one_epoch(model, loader):\n", + " model.eval()\n", + " total = 0.0\n", + "\n", + " for b in loader:\n", + " x = b[\"image\"].to(DEVICE)\n", + " y = b[\"mask\"].to(DEVICE)\n", + " logits = model(x)\n", + " loss = F.cross_entropy(logits, y, ignore_index=IGNORE_INDEX)\n", + " total += loss.item() * x.size(0)\n", + "\n", + " return total / len(loader.dataset)\n", + "\n", + "\n", + "# =========================\n", + "# 5) Inference & Submission\n", + "# =========================\n", + "@torch.no_grad()\n", + "def predict_and_submit(model, out_csv: Path):\n", + " model.eval()\n", + " test_ds = WellDataset(TEST_IMG_DIR)\n", + " loader = DataLoader(test_ds, batch_size=1, shuffle=False)\n", + "\n", + " preds = {}\n", + " for b in loader:\n", + " name = b[\"name\"][0]\n", + " raw_w = int(b[\"raw_w\"][0])\n", + " x = b[\"image\"].to(DEVICE)\n", + "\n", + " logits = model(x)\n", + " pred224 = torch.argmax(logits, dim=1).squeeze(0).cpu().numpy()\n", + "\n", + " # resize back to (160,272)\n", + " x2 = torch.from_numpy(pred224).unsqueeze(0).unsqueeze(0).float()\n", + " x2 = F.interpolate(x2, size=(H_SUB, W_SUB), mode=\"nearest\")\n", + " pred160 = x2.squeeze(0).squeeze(0).long().numpy()\n", + "\n", + " pred160 = pred160[:, :raw_w]\n", + "\n", + " flat = np.full((H_SUB * W_SUB,), IGNORE_INDEX, dtype=np.int64)\n", + " flat[: H_SUB * raw_w] = pred160.flatten()\n", + " preds[name] = flat\n", + "\n", + " pd.DataFrame(preds, dtype=\"int64\").T.to_csv(out_csv)\n", + " print(f\"[OK] submission saved to {out_csv}\")\n", + "\n", + "\n", + "# =========================\n", + "# 6) Main\n", + "# =========================\n", + "def main():\n", + " print(\"DEVICE:\", DEVICE)\n", + "\n", + " full_ds = WellDataset(TRAIN_IMG_DIR, Y_TRAIN_CSV)\n", + "\n", + " train_idx, val_idx = [], []\n", + " for i, n in enumerate(full_ds.names):\n", + " if parse_well_id(n) == 6:\n", + " val_idx.append(i)\n", + " else:\n", + " train_idx.append(i)\n", + "\n", + " train_ds = Subset(full_ds, train_idx) # well1–5\n", + " val_ds = Subset(full_ds, val_idx) # well6\n", + "\n", + " train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True)\n", + " val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE)\n", + "\n", + " model = build_model().to(DEVICE)\n", + " optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)\n", + "\n", + " best_val = 1e9\n", + " best_path = DATA_ROOT / \"best_upernet_swin.pth\"\n", + "\n", + " for epoch in range(1, EPOCHS + 1):\n", + " tr = train_one_epoch(model, train_loader, optimizer)\n", + " va = eval_one_epoch(model, val_loader)\n", + " print(f\"Epoch {epoch:02d}/{EPOCHS} | train={tr:.4f} | val={va:.4f}\")\n", + "\n", + " if va < best_val:\n", + " best_val = va\n", + " torch.save(model.state_dict(), best_path)\n", + " print(\" -> best model saved\")\n", + "\n", + " model.load_state_dict(torch.load(best_path, map_location=DEVICE, weights_only=True))\n", + " predict_and_submit(model, DATA_ROOT / \"submission.csv\")\n", + "\n", + "\n", + "if __name__ == \"__main__\":\n", + " main()\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "deep-torch", + "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.10.19" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/attention_u_net.ipynb b/attention_u_net.ipynb new file mode 100644 index 0000000..01beaa9 --- /dev/null +++ b/attention_u_net.ipynb @@ -0,0 +1,501 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "9bd2e078", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "DEVICE: cuda\n", + "Train dir: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\X_train_uDRk9z9\\images\n", + "Test dir: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\X_test_xNbnvIa\\images\n", + "Train samples: 2790 | Val samples: 1620 | val_wells={6}\n", + "Epoch 01/20 | train_loss=0.2101 | val_loss=0.1172\n", + " -> Best model saved: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_attention_unet.pth\n", + "Epoch 02/20 | train_loss=0.0944 | val_loss=0.1888\n", + "Epoch 03/20 | train_loss=0.0832 | val_loss=0.1081\n", + " -> Best model saved: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_attention_unet.pth\n", + "Epoch 04/20 | train_loss=0.0848 | val_loss=0.1979\n", + "Epoch 05/20 | train_loss=0.0776 | val_loss=0.0860\n", + " -> Best model saved: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_attention_unet.pth\n", + "Epoch 06/20 | train_loss=0.0754 | val_loss=0.1289\n", + "Epoch 07/20 | train_loss=0.0760 | val_loss=0.0803\n", + " -> Best model saved: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_attention_unet.pth\n", + "Epoch 08/20 | train_loss=0.0732 | val_loss=0.0971\n", + "Epoch 09/20 | train_loss=0.0735 | val_loss=0.1063\n", + "Epoch 10/20 | train_loss=0.0712 | val_loss=0.1057\n", + "Epoch 11/20 | train_loss=0.0703 | val_loss=0.0877\n", + "Epoch 12/20 | train_loss=0.0699 | val_loss=0.0881\n", + "Epoch 13/20 | train_loss=0.0699 | val_loss=0.0905\n", + "Epoch 14/20 | train_loss=0.0680 | val_loss=0.0802\n", + " -> Best model saved: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_attention_unet.pth\n", + "Epoch 15/20 | train_loss=0.0665 | val_loss=0.0923\n", + "Epoch 16/20 | train_loss=0.0673 | val_loss=0.0847\n", + "Epoch 17/20 | train_loss=0.0670 | val_loss=0.0874\n", + "Epoch 18/20 | train_loss=0.0666 | val_loss=0.0949\n", + "Epoch 19/20 | train_loss=0.0644 | val_loss=0.1550\n", + "Epoch 20/20 | train_loss=0.0645 | val_loss=0.1053\n", + "[OK] submission saved to: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\submission.csv\n" + ] + } + ], + "source": [ + "# -*- coding: utf-8 -*-\n", + "\"\"\"\n", + "Attention U-Net - Full runnable semantic segmentation code (no extra libs)\n", + "\n", + "Dataset:\n", + "- Train images: X_train_uDRk9z9/images (well1-6)\n", + "- Train labels: Y_train_T9NrBYo.csv (flatten + -1 padding)\n", + "- Test images: X_test_xNbnvIa/images (well7-11)\n", + "\n", + "Split:\n", + "- Train: well1-5\n", + "- Val: well6\n", + "\n", + "Output:\n", + "- submission.csv (each row = one patch name, flattened, padded to 160*272 with -1)\n", + "\n", + "Notes:\n", + "- Input patches are (160,160) or (160,272). We pad to (160,272).\n", + "- Class ids: 0/1/2; padding is -1 (ignored in loss).\n", + "\"\"\"\n", + "\n", + "import re\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "import pandas as pd\n", + "\n", + "import torch\n", + "import torch.nn as nn\n", + "import torch.nn.functional as F\n", + "from torch.utils.data import Dataset, DataLoader, Subset\n", + "\n", + "\n", + "# =========================\n", + "# 0) Paths & Hyperparameters\n", + "# =========================\n", + "DATA_ROOT = Path(r\"C:\\Users\\lenovo\\Desktop\\deep_datachallenge\") # change to your path\n", + "\n", + "TRAIN_IMAGES_DIR = DATA_ROOT / \"X_train_uDRk9z9\" / \"images\"\n", + "TEST_IMAGES_DIR = DATA_ROOT / \"X_test_xNbnvIa\" / \"images\"\n", + "Y_TRAIN_CSV = DATA_ROOT / \"Y_train_T9NrBYo.csv\"\n", + "\n", + "H = 160\n", + "W = 272\n", + "NUM_CLASSES = 3\n", + "IGNORE_INDEX = -1\n", + "\n", + "BATCH_SIZE = 8 # Attention U-Net is lighter than Swin/Mask2Former; 8 is often OK on 4060(8GB)\n", + "LR = 1e-3\n", + "WEIGHT_DECAY = 1e-4\n", + "EPOCHS = 20\n", + "\n", + "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", + "\n", + "\n", + "# =========================\n", + "# 1) Utils\n", + "# =========================\n", + "def parse_well_id(name: str) -> int:\n", + " \"\"\"Extract well id from: well_1_section_0_patch_0 -> 1\"\"\"\n", + " m = re.search(r\"well_(\\d+)_\", name)\n", + " return int(m.group(1)) if m else -1\n", + "\n", + "\n", + "def minmax_normalize(x: np.ndarray) -> np.ndarray:\n", + " \"\"\"Min-max normalize; replace NaN/inf with 0.\"\"\"\n", + " x = x.astype(np.float32)\n", + " x = np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)\n", + " x_min = float(x.min())\n", + " x_max = float(x.max())\n", + " if x_max - x_min < 1e-6:\n", + " return np.zeros_like(x, dtype=np.float32)\n", + " return (x - x_min) / (x_max - x_min)\n", + "\n", + "\n", + "def pad_to_160x272(img: np.ndarray, fill_value: float = 0.0) -> np.ndarray:\n", + " \"\"\"Pad (160,160) or (160,272) to (160,272).\"\"\"\n", + " h, w = img.shape\n", + " assert h == H, f\"Expected height {H}, got {h}\"\n", + " if w == W:\n", + " return img\n", + " if w < W:\n", + " out = np.full((H, W), fill_value, dtype=img.dtype)\n", + " out[:, :w] = img\n", + " return out\n", + " return img[:, :W]\n", + "\n", + "\n", + "def decode_mask_from_csv_row(row_values: np.ndarray) -> np.ndarray:\n", + " \"\"\"\n", + " Decode one CSV row -> (160,w) semantic mask\n", + " - row_values: flattened mask with -1 padding\n", + " \"\"\"\n", + " valid = row_values[row_values != IGNORE_INDEX]\n", + " assert len(valid) % H == 0, f\"Valid mask length {len(valid)} not divisible by 160\"\n", + " ww = len(valid) // H\n", + " return valid.reshape(H, ww).astype(np.int64)\n", + "\n", + "\n", + "def pad_mask_to_160x272(mask: np.ndarray) -> np.ndarray:\n", + " \"\"\"Pad (160,w) -> (160,272) using -1 for padding.\"\"\"\n", + " h, w = mask.shape\n", + " assert h == H\n", + " if w == W:\n", + " return mask\n", + " out = np.full((H, W), IGNORE_INDEX, dtype=np.int64)\n", + " out[:, :w] = mask\n", + " return out\n", + "\n", + "\n", + "# =========================\n", + "# 2) Dataset\n", + "# =========================\n", + "class WellSegDataset(Dataset):\n", + " def __init__(self, images_dir: Path, y_csv_path: Path = None, wells=None):\n", + " \"\"\"\n", + " y_csv_path=None => test mode.\n", + " wells: optional set of well ids to filter.\n", + " \"\"\"\n", + " self.images_dir = images_dir\n", + " self.has_label = y_csv_path is not None\n", + " self.wells = set(wells) if wells is not None else None\n", + "\n", + " all_paths = sorted(images_dir.glob(\"*.npy\"))\n", + " all_names = [p.stem for p in all_paths]\n", + "\n", + " if self.wells is None:\n", + " self.image_paths = all_paths\n", + " self.names = all_names\n", + " else:\n", + " keep = []\n", + " keep_names = []\n", + " for p, n in zip(all_paths, all_names):\n", + " if parse_well_id(n) in self.wells:\n", + " keep.append(p)\n", + " keep_names.append(n)\n", + " self.image_paths = keep\n", + " self.names = keep_names\n", + "\n", + " self.y_df = pd.read_csv(y_csv_path, index_col=0) if self.has_label else None\n", + "\n", + " def __len__(self):\n", + " return len(self.image_paths)\n", + "\n", + " def __getitem__(self, idx: int):\n", + " name = self.names[idx]\n", + " img_path = self.image_paths[idx]\n", + "\n", + " img = np.load(img_path) # (160,160) or (160,272)\n", + " raw_w = int(img.shape[1])\n", + " img = minmax_normalize(img)\n", + " img = pad_to_160x272(img, fill_value=0.0)\n", + " x = torch.from_numpy(img).unsqueeze(0).float() # (1,160,272)\n", + "\n", + " if not self.has_label:\n", + " return {\"name\": name, \"image\": x, \"raw_w\": raw_w}\n", + "\n", + " row = self.y_df.loc[name].values.astype(np.int64)\n", + " mask = decode_mask_from_csv_row(row) # (160,w)\n", + " mask = pad_mask_to_160x272(mask) # (160,272)\n", + " y = torch.from_numpy(mask).long() # (160,272)\n", + "\n", + " return {\"name\": name, \"image\": x, \"mask\": y, \"raw_w\": raw_w}\n", + "\n", + "\n", + "# =========================\n", + "# 3) Attention U-Net building blocks\n", + "# =========================\n", + "class DoubleConv(nn.Module):\n", + " \"\"\"(Conv -> BN -> ReLU) * 2\"\"\"\n", + " def __init__(self, in_ch, out_ch):\n", + " super().__init__()\n", + " self.net = nn.Sequential(\n", + " nn.Conv2d(in_ch, out_ch, kernel_size=3, padding=1, bias=False),\n", + " nn.BatchNorm2d(out_ch),\n", + " nn.ReLU(inplace=True),\n", + " nn.Conv2d(out_ch, out_ch, kernel_size=3, padding=1, bias=False),\n", + " nn.BatchNorm2d(out_ch),\n", + " nn.ReLU(inplace=True),\n", + " )\n", + "\n", + " def forward(self, x):\n", + " return self.net(x)\n", + "\n", + "\n", + "class Down(nn.Module):\n", + " \"\"\"Downscaling with maxpool then double conv\"\"\"\n", + " def __init__(self, in_ch, out_ch):\n", + " super().__init__()\n", + " self.net = nn.Sequential(\n", + " nn.MaxPool2d(2),\n", + " DoubleConv(in_ch, out_ch)\n", + " )\n", + "\n", + " def forward(self, x):\n", + " return self.net(x)\n", + "\n", + "\n", + "class Up(nn.Module):\n", + " \"\"\"Upscaling then double conv (we use bilinear upsample)\"\"\"\n", + " def __init__(self, in_ch, out_ch):\n", + " super().__init__()\n", + " self.conv = DoubleConv(in_ch, out_ch)\n", + "\n", + " def forward(self, x, skip):\n", + " x = F.interpolate(x, size=skip.shape[-2:], mode=\"bilinear\", align_corners=False)\n", + " x = torch.cat([skip, x], dim=1)\n", + " return self.conv(x)\n", + "\n", + "\n", + "class AttentionGate(nn.Module):\n", + " \"\"\"\n", + " Attention Gate (AG) from Attention U-Net:\n", + " - g: gating signal from decoder\n", + " - x: skip connection from encoder\n", + " output: attended skip features\n", + " \"\"\"\n", + " def __init__(self, F_g, F_l, F_int):\n", + " super().__init__()\n", + " self.W_g = nn.Sequential(\n", + " nn.Conv2d(F_g, F_int, kernel_size=1, bias=False),\n", + " nn.BatchNorm2d(F_int)\n", + " )\n", + " self.W_x = nn.Sequential(\n", + " nn.Conv2d(F_l, F_int, kernel_size=1, bias=False),\n", + " nn.BatchNorm2d(F_int)\n", + " )\n", + " self.psi = nn.Sequential(\n", + " nn.Conv2d(F_int, 1, kernel_size=1, bias=False),\n", + " nn.BatchNorm2d(1),\n", + " nn.Sigmoid()\n", + " )\n", + " self.relu = nn.ReLU(inplace=True)\n", + "\n", + " def forward(self, g, x):\n", + " # g and x may have different spatial sizes; resize g to x size\n", + " g1 = self.W_g(g)\n", + " x1 = self.W_x(x)\n", + " if g1.shape[-2:] != x1.shape[-2:]:\n", + " g1 = F.interpolate(g1, size=x1.shape[-2:], mode=\"bilinear\", align_corners=False)\n", + "\n", + " psi = self.relu(g1 + x1)\n", + " psi = self.psi(psi) # (B,1,H,W) attention mask in [0,1]\n", + " return x * psi\n", + "\n", + "\n", + "class AttentionUNet(nn.Module):\n", + " def __init__(self, in_channels=1, num_classes=3, base=32):\n", + " super().__init__()\n", + " # Encoder\n", + " self.inc = DoubleConv(in_channels, base) # 32\n", + " self.down1 = Down(base, base*2) # 64\n", + " self.down2 = Down(base*2, base*4) # 128\n", + " self.down3 = Down(base*4, base*8) # 256\n", + " self.down4 = Down(base*8, base*16) # 512\n", + "\n", + " # Decoder + Attention gates\n", + " self.att4 = AttentionGate(F_g=base*16, F_l=base*8, F_int=base*4)\n", + " self.up4 = Up(in_ch=base*16 + base*8, out_ch=base*8)\n", + "\n", + " self.att3 = AttentionGate(F_g=base*8, F_l=base*4, F_int=base*2)\n", + " self.up3 = Up(in_ch=base*8 + base*4, out_ch=base*4)\n", + "\n", + " self.att2 = AttentionGate(F_g=base*4, F_l=base*2, F_int=base)\n", + " self.up2 = Up(in_ch=base*4 + base*2, out_ch=base*2)\n", + "\n", + " self.att1 = AttentionGate(F_g=base*2, F_l=base, F_int=base//2)\n", + " self.up1 = Up(in_ch=base*2 + base, out_ch=base)\n", + "\n", + " self.outc = nn.Conv2d(base, num_classes, kernel_size=1)\n", + "\n", + " def forward(self, x):\n", + " # Encoder\n", + " x1 = self.inc(x) # (B,base, H, W)\n", + " x2 = self.down1(x1) # (B,base*2, H/2, W/2)\n", + " x3 = self.down2(x2) # (B,base*4, H/4, W/4)\n", + " x4 = self.down3(x3) # (B,base*8, H/8, W/8)\n", + " x5 = self.down4(x4) # (B,base*16,H/16,W/16)\n", + "\n", + " # Decoder with attention on skip connections\n", + " s4 = self.att4(g=x5, x=x4)\n", + " d4 = self.up4(x5, s4)\n", + "\n", + " s3 = self.att3(g=d4, x=x3)\n", + " d3 = self.up3(d4, s3)\n", + "\n", + " s2 = self.att2(g=d3, x=x2)\n", + " d2 = self.up2(d3, s2)\n", + "\n", + " s1 = self.att1(g=d2, x=x1)\n", + " d1 = self.up1(d2, s1)\n", + "\n", + " logits = self.outc(d1) # (B,C,H,W)\n", + " return logits\n", + "\n", + "\n", + "# =========================\n", + "# 4) Train / Validate\n", + "# =========================\n", + "def train_one_epoch(model, loader, optimizer):\n", + " model.train()\n", + " total_loss = 0.0\n", + "\n", + " for batch in loader:\n", + " x = batch[\"image\"].to(DEVICE) # (B,1,160,272)\n", + " y = batch[\"mask\"].to(DEVICE) # (B,160,272)\n", + "\n", + " logits = model(x) # (B,C,160,272)\n", + " loss = F.cross_entropy(logits, y, ignore_index=IGNORE_INDEX)\n", + "\n", + " optimizer.zero_grad()\n", + " loss.backward()\n", + " optimizer.step()\n", + "\n", + " total_loss += float(loss.item()) * x.size(0)\n", + "\n", + " return total_loss / len(loader.dataset)\n", + "\n", + "\n", + "@torch.no_grad()\n", + "def valid_one_epoch(model, loader):\n", + " model.eval()\n", + " total_loss = 0.0\n", + "\n", + " for batch in loader:\n", + " x = batch[\"image\"].to(DEVICE)\n", + " y = batch[\"mask\"].to(DEVICE)\n", + "\n", + " logits = model(x)\n", + " loss = F.cross_entropy(logits, y, ignore_index=IGNORE_INDEX)\n", + " total_loss += float(loss.item()) * x.size(0)\n", + "\n", + " return total_loss / len(loader.dataset)\n", + "\n", + "\n", + "# =========================\n", + "# 5) Inference & submission\n", + "# =========================\n", + "@torch.no_grad()\n", + "def predict_and_make_submission(model, test_images_dir: Path, out_csv_path: Path):\n", + " \"\"\"\n", + " Predict all test patches and write submission.csv:\n", + " - argmax over classes\n", + " - crop to raw width\n", + " - pad to 160*272 using -1\n", + " \"\"\"\n", + " model.eval()\n", + "\n", + " test_ds = WellSegDataset(test_images_dir, y_csv_path=None)\n", + " test_loader = DataLoader(test_ds, batch_size=1, shuffle=False, num_workers=0)\n", + "\n", + " preds_dict = {}\n", + "\n", + " for batch in test_loader:\n", + " name = batch[\"name\"][0]\n", + " raw_w = int(batch[\"raw_w\"][0])\n", + " x = batch[\"image\"].to(DEVICE)\n", + "\n", + " logits = model(x)\n", + " pred_full = torch.argmax(logits, dim=1).squeeze(0).cpu().numpy().astype(np.int64) # (160,272)\n", + "\n", + " pred = pred_full[:, :raw_w] # crop to original width\n", + "\n", + " if raw_w < W:\n", + " padded = np.full((H * W,), IGNORE_INDEX, dtype=np.int64)\n", + " padded[: H * raw_w] = pred.flatten()\n", + " preds_dict[name] = padded\n", + " else:\n", + " preds_dict[name] = pred.flatten()\n", + "\n", + " sub = pd.DataFrame(preds_dict, dtype=\"int64\").T\n", + " sub.to_csv(out_csv_path)\n", + " print(f\"[OK] submission saved to: {out_csv_path}\")\n", + "\n", + "\n", + "# =========================\n", + "# 6) Main\n", + "# =========================\n", + "def main():\n", + " print(f\"DEVICE: {DEVICE}\")\n", + " print(f\"Train dir: {TRAIN_IMAGES_DIR}\")\n", + " print(f\"Test dir: {TEST_IMAGES_DIR}\")\n", + "\n", + " # Load all train data (well1-6)\n", + " train_ds_all = WellSegDataset(TRAIN_IMAGES_DIR, Y_TRAIN_CSV)\n", + "\n", + " # Split by well: well6 as validation\n", + " VAL_WELLS = {6}\n", + " train_indices, val_indices = [], []\n", + " for i, name in enumerate(train_ds_all.names):\n", + " w = parse_well_id(name)\n", + " if w in VAL_WELLS:\n", + " val_indices.append(i)\n", + " else:\n", + " train_indices.append(i)\n", + "\n", + " train_ds = Subset(train_ds_all, train_indices) # well1-5\n", + " val_ds = Subset(train_ds_all, val_indices) # well6\n", + "\n", + " train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)\n", + " val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)\n", + "\n", + " print(f\"Train samples: {len(train_ds)} | Val samples: {len(val_ds)} | val_wells={VAL_WELLS}\")\n", + "\n", + " model = AttentionUNet(in_channels=1, num_classes=NUM_CLASSES, base=32).to(DEVICE)\n", + " optimizer = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)\n", + "\n", + " best_val = 1e9\n", + " best_path = DATA_ROOT / \"best_attention_unet.pth\"\n", + "\n", + " for epoch in range(1, EPOCHS + 1):\n", + " tr_loss = train_one_epoch(model, train_loader, optimizer)\n", + " va_loss = valid_one_epoch(model, val_loader)\n", + "\n", + " print(f\"Epoch {epoch:02d}/{EPOCHS} | train_loss={tr_loss:.4f} | val_loss={va_loss:.4f}\")\n", + "\n", + " if va_loss < best_val:\n", + " best_val = va_loss\n", + " torch.save(model.state_dict(), best_path)\n", + " print(f\" -> Best model saved: {best_path}\")\n", + "\n", + " # Predict test and write submission\n", + " out_csv = DATA_ROOT / \"submission.csv\"\n", + " state_dict = torch.load(best_path, map_location=DEVICE, weights_only=True)\n", + " model.load_state_dict(state_dict)\n", + " predict_and_make_submission(model, TEST_IMAGES_DIR, out_csv)\n", + "\n", + "\n", + "if __name__ == \"__main__\":\n", + " main()\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "deep-torch", + "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.10.19" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/mask2former.ipynb b/mask2former.ipynb new file mode 100644 index 0000000..98659bd --- /dev/null +++ b/mask2former.ipynb @@ -0,0 +1,592 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "e5988541", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "c:\\Users\\lenovo\\anaconda3\\envs\\deep-torch\\lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", + " from .autonotebook import tqdm as notebook_tqdm\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "DEVICE: cuda\n", + "Train dir: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\X_train_uDRk9z9\\images\n", + "Test dir: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\X_test_xNbnvIa\\images\n", + "Model input size: 224x224\n", + "Pretrained: facebook/mask2former-swin-tiny-ade-semantic\n", + "Train samples: 2790 | Val samples: 1620 | val_wells={6}\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Using a slow image processor as `use_fast` is unset and a slow processor was saved with this model. `use_fast=True` will be the default behavior in v4.52, even if the model was saved with a slow processor. This will result in minor differences in outputs. You'll still be able to use a slow processor with `use_fast=False`.\n", + "c:\\Users\\lenovo\\anaconda3\\envs\\deep-torch\\lib\\site-packages\\transformers\\image_processing_base.py:417: UserWarning: The following named arguments are not valid for `Mask2FormerImageProcessor.__init__` and were ignored: '_max_size', 'reduce_labels'\n", + " image_processor = cls(**image_processor_dict)\n", + "Some weights of Mask2FormerForUniversalSegmentation were not initialized from the model checkpoint at facebook/mask2former-swin-tiny-ade-semantic and are newly initialized because the shapes did not match:\n", + "- class_predictor.weight: found shape torch.Size([151, 256]) in the checkpoint and torch.Size([4, 256]) in the model instantiated\n", + "- class_predictor.bias: found shape torch.Size([151]) in the checkpoint and torch.Size([4]) in the model instantiated\n", + "- criterion.empty_weight: found shape torch.Size([151]) in the checkpoint and torch.Size([4]) in the model instantiated\n", + "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Epoch 01/20 | train_loss=17.3520 | val_loss=15.0105\n", + " -> Best model saved: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_mask2former.pth\n", + "Epoch 02/20 | train_loss=13.7775 | val_loss=15.2349\n", + "Epoch 03/20 | train_loss=12.6401 | val_loss=14.5714\n", + " -> Best model saved: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_mask2former.pth\n", + "Epoch 04/20 | train_loss=12.0788 | val_loss=14.2691\n", + " -> Best model saved: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_mask2former.pth\n", + "Epoch 05/20 | train_loss=11.6676 | val_loss=15.0381\n", + "Epoch 06/20 | train_loss=11.4770 | val_loss=14.3574\n", + "Epoch 07/20 | train_loss=10.9957 | val_loss=14.1694\n", + " -> Best model saved: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_mask2former.pth\n", + "Epoch 08/20 | train_loss=10.6843 | val_loss=14.9301\n", + "Epoch 09/20 | train_loss=10.3577 | val_loss=15.2518\n", + "Epoch 10/20 | train_loss=10.2286 | val_loss=15.0672\n", + "Epoch 11/20 | train_loss=9.6430 | val_loss=15.5822\n", + "Epoch 12/20 | train_loss=9.4476 | val_loss=14.6725\n", + "Epoch 13/20 | train_loss=9.0929 | val_loss=15.5858\n", + "Epoch 14/20 | train_loss=9.0179 | val_loss=15.7629\n", + "Epoch 15/20 | train_loss=8.4409 | val_loss=16.1311\n", + "Epoch 16/20 | train_loss=8.3673 | val_loss=16.3036\n", + "Epoch 17/20 | train_loss=8.0378 | val_loss=17.1427\n", + "Epoch 18/20 | train_loss=7.8582 | val_loss=17.3488\n", + "Epoch 19/20 | train_loss=7.5451 | val_loss=17.6415\n", + "Epoch 20/20 | train_loss=7.1668 | val_loss=16.5966\n", + "[OK] submission saved to: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\submission.csv\n" + ] + } + ], + "source": [ + "# -*- coding: utf-8 -*-\n", + "\"\"\"\n", + "Mask2Former (HuggingFace Transformers) - Full runnable semantic segmentation code\n", + "\n", + "What this script does:\n", + "- Reads .npy ultrasound patches\n", + "- Trains on well1-5, validates on well6 (from X_train_uDRk9z9/images)\n", + "- Predicts on X_test_xNbnvIa/images (well7-11)\n", + "- Writes submission.csv with the SAME format as before:\n", + " - each row = one patch name\n", + " - flattened mask\n", + " - padded to 160*272 with -1\n", + "\n", + "Why we resize:\n", + "- Mask2Former backbones are usually trained on larger resolutions.\n", + "- To keep it simple and fit RTX 4060 (8GB), we resize inputs to 224x224 during training/inference,\n", + " then upsample predictions back to (160,272) for submission.\n", + "\n", + "Install (in your CUDA environment):\n", + " pip install transformers accelerate\n", + "\n", + "Notes:\n", + "- If your machine cannot download pretrained weights (no internet), set PRETRAINED=None and it will start from scratch.\n", + "- Mask2Former expects instance-style labels: a set of binary masks + class ids per image.\n", + " We convert your (H,W) semantic mask into that format automatically.\n", + "\"\"\"\n", + "\n", + "import re\n", + "from pathlib import Path\n", + "from typing import Dict, List, Tuple\n", + "\n", + "import numpy as np\n", + "import pandas as pd\n", + "\n", + "import torch\n", + "import torch.nn.functional as F\n", + "from torch.utils.data import Dataset, DataLoader, Subset\n", + "\n", + "from transformers import (\n", + " AutoImageProcessor,\n", + " Mask2FormerForUniversalSegmentation,\n", + ")\n", + "\n", + "\n", + "# =========================\n", + "# 0. Paths & Hyperparameters\n", + "# =========================\n", + "DATA_ROOT = Path(r\"C:\\Users\\lenovo\\Desktop\\deep_datachallenge\") # change to your path\n", + "\n", + "TRAIN_IMAGES_DIR = DATA_ROOT / \"X_train_uDRk9z9\" / \"images\"\n", + "TEST_IMAGES_DIR = DATA_ROOT / \"X_test_xNbnvIa\" / \"images\"\n", + "Y_TRAIN_CSV = DATA_ROOT / \"Y_train_T9NrBYo.csv\"\n", + "\n", + "# Original submission size\n", + "TARGET_H = 160\n", + "TARGET_W = 272\n", + "\n", + "# Model input size (keep small for 4060)\n", + "MODEL_H = 224\n", + "MODEL_W = 224\n", + "\n", + "NUM_CLASSES = 3\n", + "IGNORE_INDEX = -1\n", + "\n", + "BATCH_SIZE = 2 # Mask2Former is heavy; start with 1~2 on 4060 8GB\n", + "LR = 5e-5\n", + "WEIGHT_DECAY = 1e-4\n", + "EPOCHS = 20\n", + "\n", + "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", + "\n", + "# A strong semantic pretrained checkpoint (Swin-T backbone)\n", + "# If you have no internet, set PRETRAINED = None\n", + "PRETRAINED = \"facebook/mask2former-swin-tiny-ade-semantic\"\n", + "\n", + "\n", + "# =========================\n", + "# 1. Utils\n", + "# =========================\n", + "def parse_well_id(name: str) -> int:\n", + " \"\"\"Extract well id from: well_1_section_0_patch_0 -> 1\"\"\"\n", + " m = re.search(r\"well_(\\d+)_\", name)\n", + " return int(m.group(1)) if m else -1\n", + "\n", + "\n", + "def minmax_normalize(x: np.ndarray) -> np.ndarray:\n", + " \"\"\"Min-max normalize; replace NaN/inf with 0.\"\"\"\n", + " x = x.astype(np.float32)\n", + " x = np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)\n", + " x_min = float(x.min())\n", + " x_max = float(x.max())\n", + " if x_max - x_min < 1e-6:\n", + " return np.zeros_like(x, dtype=np.float32)\n", + " return (x - x_min) / (x_max - x_min)\n", + "\n", + "\n", + "def pad_to_160x272(img: np.ndarray, fill_value: float = 0.0) -> np.ndarray:\n", + " \"\"\"Pad (160,160) or (160,272) to (160,272).\"\"\"\n", + " h, w = img.shape\n", + " assert h == TARGET_H, f\"Expected height {TARGET_H}, got {h}\"\n", + " if w == TARGET_W:\n", + " return img\n", + " if w < TARGET_W:\n", + " out = np.full((TARGET_H, TARGET_W), fill_value, dtype=img.dtype)\n", + " out[:, :w] = img\n", + " return out\n", + " return img[:, :TARGET_W]\n", + "\n", + "\n", + "def decode_mask_from_csv_row(row_values: np.ndarray) -> np.ndarray:\n", + " \"\"\"\n", + " Decode one CSV row -> (160,w) semantic mask\n", + " - row_values: flattened mask with -1 padding\n", + " \"\"\"\n", + " valid = row_values[row_values != IGNORE_INDEX]\n", + " assert len(valid) % TARGET_H == 0, f\"Valid mask length {len(valid)} not divisible by 160\"\n", + " w = len(valid) // TARGET_H\n", + " return valid.reshape(TARGET_H, w).astype(np.int64)\n", + "\n", + "\n", + "def pad_mask_to_160x272(mask: np.ndarray) -> np.ndarray:\n", + " \"\"\"Pad (160,w) -> (160,272) using -1 for padding.\"\"\"\n", + " h, w = mask.shape\n", + " assert h == TARGET_H\n", + " if w == TARGET_W:\n", + " return mask\n", + " out = np.full((TARGET_H, TARGET_W), IGNORE_INDEX, dtype=np.int64)\n", + " out[:, :w] = mask\n", + " return out\n", + "\n", + "\n", + "def resize_image_torch(img_1hw: torch.Tensor, h: int, w: int) -> torch.Tensor:\n", + " \"\"\"Resize image tensor (1,H,W) -> (1,h,w) (bilinear).\"\"\"\n", + " x = img_1hw.unsqueeze(0) # (1,1,H,W)\n", + " x = F.interpolate(x, size=(h, w), mode=\"bilinear\", align_corners=False)\n", + " return x.squeeze(0) # (1,h,w)\n", + "\n", + "\n", + "def resize_mask_torch(mask_hw: torch.Tensor, h: int, w: int) -> torch.Tensor:\n", + " \"\"\"Resize mask tensor (H,W) -> (h,w) (nearest).\"\"\"\n", + " y = mask_hw.unsqueeze(0).unsqueeze(0).float() # (1,1,H,W)\n", + " y = F.interpolate(y, size=(h, w), mode=\"nearest\")\n", + " return y.squeeze(0).squeeze(0).long()\n", + "\n", + "\n", + "def semantic_to_mask2former_targets(\n", + " semantic_mask: torch.Tensor,\n", + " num_classes: int,\n", + " ignore_index: int = -1,\n", + ") -> Tuple[torch.Tensor, torch.Tensor]:\n", + " \"\"\"\n", + " Convert a semantic mask (H,W) into Mask2Former targets:\n", + " - class_labels: (N,) long\n", + " - mask_labels: (N,H,W) float (0/1)\n", + "\n", + " We create one binary mask per class present in the image (excluding ignore_index).\n", + " \"\"\"\n", + " # semantic_mask: (H,W)\n", + " valid = semantic_mask != ignore_index\n", + " if valid.sum() == 0:\n", + " # If everything is ignore, create a dummy empty target (rare).\n", + " # Use background class 0 with an all-zero mask.\n", + " class_labels = torch.tensor([0], dtype=torch.long)\n", + " mask_labels = torch.zeros((1, semantic_mask.shape[0], semantic_mask.shape[1]), dtype=torch.float32)\n", + " return class_labels, mask_labels\n", + "\n", + " present_classes = torch.unique(semantic_mask[valid]).tolist()\n", + " present_classes = [int(c) for c in present_classes if 0 <= int(c) < num_classes]\n", + "\n", + " if len(present_classes) == 0:\n", + " class_labels = torch.tensor([0], dtype=torch.long)\n", + " mask_labels = torch.zeros((1, semantic_mask.shape[0], semantic_mask.shape[1]), dtype=torch.float32)\n", + " return class_labels, mask_labels\n", + "\n", + " masks = []\n", + " classes = []\n", + " for c in present_classes:\n", + " m = (semantic_mask == c) & valid\n", + " if m.sum() == 0:\n", + " continue\n", + " masks.append(m.float())\n", + " classes.append(c)\n", + "\n", + " if len(classes) == 0:\n", + " class_labels = torch.tensor([0], dtype=torch.long)\n", + " mask_labels = torch.zeros((1, semantic_mask.shape[0], semantic_mask.shape[1]), dtype=torch.float32)\n", + " return class_labels, mask_labels\n", + "\n", + " class_labels = torch.tensor(classes, dtype=torch.long)\n", + " mask_labels = torch.stack(masks, dim=0).float() # (N,H,W)\n", + " return class_labels, mask_labels\n", + "\n", + "\n", + "# =========================\n", + "# 2. Dataset\n", + "# =========================\n", + "class WellSegDataset(Dataset):\n", + " def __init__(self, images_dir: Path, y_csv_path: Path = None):\n", + " self.images_dir = images_dir\n", + " self.has_label = y_csv_path is not None\n", + "\n", + " self.image_paths = sorted(images_dir.glob(\"*.npy\"))\n", + " self.names = [p.stem for p in self.image_paths]\n", + "\n", + " self.y_df = pd.read_csv(y_csv_path, index_col=0) if self.has_label else None\n", + "\n", + " def __len__(self):\n", + " return len(self.image_paths)\n", + "\n", + " def __getitem__(self, idx: int):\n", + " name = self.names[idx]\n", + " img_path = self.image_paths[idx]\n", + "\n", + " img = np.load(img_path) # (160,160) or (160,272)\n", + " raw_w = int(img.shape[1]) # used to crop back for submission\n", + " img = minmax_normalize(img)\n", + " img = pad_to_160x272(img, fill_value=0.0)\n", + "\n", + " img_t = torch.from_numpy(img).unsqueeze(0).float() # (1,160,272)\n", + " img_t = resize_image_torch(img_t, MODEL_H, MODEL_W) # (1,224,224)\n", + "\n", + " if not self.has_label:\n", + " return {\"name\": name, \"image\": img_t, \"raw_w\": raw_w}\n", + "\n", + " row = self.y_df.loc[name].values.astype(np.int64)\n", + " mask = decode_mask_from_csv_row(row) # (160,w)\n", + " mask = pad_mask_to_160x272(mask) # (160,272)\n", + " mask_t = torch.from_numpy(mask).long() # (160,272)\n", + " mask_t = resize_mask_torch(mask_t, MODEL_H, MODEL_W) # (224,224)\n", + "\n", + " return {\"name\": name, \"image\": img_t, \"mask\": mask_t, \"raw_w\": raw_w}\n", + "\n", + "\n", + "# =========================\n", + "# 3. Collate for Mask2Former\n", + "# =========================\n", + "def collate_mask2former(batch: List[Dict]) -> Dict:\n", + " \"\"\"\n", + " Build a batch dict for Mask2Former:\n", + " - pixel_values: (B,3,224,224) float\n", + " - pixel_mask: (B,224,224) bool/long (1=valid)\n", + " - mask_labels: list of (Ni,224,224) float\n", + " - class_labels: list of (Ni,) long\n", + " \"\"\"\n", + " names = [b[\"name\"] for b in batch]\n", + " raw_ws = torch.tensor([b[\"raw_w\"] for b in batch], dtype=torch.long)\n", + "\n", + " # image: (1,224,224) -> (3,224,224) by repeating channel\n", + " imgs_1 = torch.stack([b[\"image\"] for b in batch], dim=0) # (B,1,224,224)\n", + " pixel_values = imgs_1.repeat(1, 3, 1, 1) # (B,3,224,224)\n", + "\n", + " pixel_mask = torch.ones((pixel_values.shape[0], MODEL_H, MODEL_W), dtype=torch.long)\n", + "\n", + " out = {\n", + " \"names\": names,\n", + " \"raw_ws\": raw_ws,\n", + " \"pixel_values\": pixel_values,\n", + " \"pixel_mask\": pixel_mask,\n", + " }\n", + "\n", + " if \"mask\" in batch[0]:\n", + " class_labels_list = []\n", + " mask_labels_list = []\n", + " for b in batch:\n", + " y = b[\"mask\"] # (224,224)\n", + " cls, msk = semantic_to_mask2former_targets(y, NUM_CLASSES, IGNORE_INDEX)\n", + " class_labels_list.append(cls)\n", + " mask_labels_list.append(msk)\n", + "\n", + " out[\"class_labels\"] = class_labels_list\n", + " out[\"mask_labels\"] = mask_labels_list\n", + "\n", + " return out\n", + "\n", + "\n", + "# =========================\n", + "# 4. Model builder\n", + "# =========================\n", + "def build_model_and_processor(num_classes: int):\n", + " id2label = {0: \"class0\", 1: \"class1\", 2: \"class2\"}\n", + " label2id = {v: k for k, v in id2label.items()}\n", + "\n", + " if PRETRAINED is None:\n", + " # Train from scratch\n", + " processor = AutoImageProcessor.from_pretrained(\"facebook/mask2former-swin-tiny-ade-semantic\")\n", + " model = Mask2FormerForUniversalSegmentation.from_pretrained(\n", + " \"facebook/mask2former-swin-tiny-ade-semantic\",\n", + " ignore_mismatched_sizes=True,\n", + " id2label=id2label,\n", + " label2id=label2id,\n", + " num_labels=num_classes,\n", + " use_safetensors=True,\n", + " )\n", + " else:\n", + " processor = AutoImageProcessor.from_pretrained(PRETRAINED)\n", + " model = Mask2FormerForUniversalSegmentation.from_pretrained(\n", + " PRETRAINED,\n", + " ignore_mismatched_sizes=True, # allow changing num_labels\n", + " id2label=id2label,\n", + " label2id=label2id,\n", + " num_labels=num_classes,\n", + " )\n", + "\n", + " return model, processor\n", + "\n", + "\n", + "# =========================\n", + "# 5. Train / Validate\n", + "# =========================\n", + "def train_one_epoch(model, loader, optimizer):\n", + " model.train()\n", + " total_loss = 0.0\n", + "\n", + " for batch in loader:\n", + " pixel_values = batch[\"pixel_values\"].to(DEVICE) # (B,3,224,224)\n", + " pixel_mask = batch[\"pixel_mask\"].to(DEVICE) # (B,224,224)\n", + "\n", + " # Mask2Former expects lists for labels (length B)\n", + " class_labels = [x.to(DEVICE) for x in batch[\"class_labels\"]]\n", + " mask_labels = [x.to(DEVICE) for x in batch[\"mask_labels\"]]\n", + "\n", + " outputs = model(\n", + " pixel_values=pixel_values,\n", + " pixel_mask=pixel_mask,\n", + " class_labels=class_labels,\n", + " mask_labels=mask_labels,\n", + " )\n", + "\n", + " loss = outputs.loss\n", + "\n", + " optimizer.zero_grad()\n", + " loss.backward()\n", + " optimizer.step()\n", + "\n", + " total_loss += float(loss.item()) * pixel_values.size(0)\n", + "\n", + " return total_loss / len(loader.dataset)\n", + "\n", + "\n", + "@torch.no_grad()\n", + "def valid_one_epoch(model, loader):\n", + " model.eval()\n", + " total_loss = 0.0\n", + "\n", + " for batch in loader:\n", + " pixel_values = batch[\"pixel_values\"].to(DEVICE)\n", + " pixel_mask = batch[\"pixel_mask\"].to(DEVICE)\n", + " class_labels = [x.to(DEVICE) for x in batch[\"class_labels\"]]\n", + " mask_labels = [x.to(DEVICE) for x in batch[\"mask_labels\"]]\n", + "\n", + " outputs = model(\n", + " pixel_values=pixel_values,\n", + " pixel_mask=pixel_mask,\n", + " class_labels=class_labels,\n", + " mask_labels=mask_labels,\n", + " )\n", + " loss = outputs.loss\n", + " total_loss += float(loss.item()) * pixel_values.size(0)\n", + "\n", + " return total_loss / len(loader.dataset)\n", + "\n", + "\n", + "# =========================\n", + "# 6. Inference & submission\n", + "# =========================\n", + "@torch.no_grad()\n", + "def predict_and_make_submission(model, processor, test_images_dir: Path, out_csv_path: Path):\n", + " \"\"\"\n", + " Predict all test patches and write submission.csv.\n", + " Steps:\n", + " - model predicts at 224x224\n", + " - we use processor.post_process_semantic_segmentation to get semantic map\n", + " - upsample semantic map to (160,272)\n", + " - crop to raw width and pad to 160*272 with -1\n", + " \"\"\"\n", + " model.eval()\n", + "\n", + " test_ds = WellSegDataset(test_images_dir, y_csv_path=None)\n", + " test_loader = DataLoader(\n", + " test_ds,\n", + " batch_size=1,\n", + " shuffle=False,\n", + " num_workers=0,\n", + " collate_fn=collate_mask2former,\n", + " )\n", + "\n", + " preds_dict = {}\n", + "\n", + " for batch in test_loader:\n", + " name = batch[\"names\"][0]\n", + " raw_w = int(batch[\"raw_ws\"][0].item())\n", + "\n", + " pixel_values = batch[\"pixel_values\"].to(DEVICE) # (1,3,224,224)\n", + " pixel_mask = batch[\"pixel_mask\"].to(DEVICE)\n", + "\n", + " outputs = model(pixel_values=pixel_values, pixel_mask=pixel_mask)\n", + "\n", + " # Post-process semantic segmentation\n", + " target_sizes = [(MODEL_H, MODEL_W)]\n", + " seg_list = processor.post_process_semantic_segmentation(outputs, target_sizes=target_sizes)\n", + " seg_224 = seg_list[0].to(torch.int64) # (224,224)\n", + "\n", + " # Upsample to (160,272) using nearest\n", + " seg_224 = seg_224.unsqueeze(0).unsqueeze(0).float() # (1,1,224,224)\n", + " seg_160_272 = F.interpolate(seg_224, size=(TARGET_H, TARGET_W), mode=\"nearest\").squeeze(0).squeeze(0)\n", + " seg_160_272 = seg_160_272.cpu().numpy().astype(np.int64) # (160,272)\n", + "\n", + " # Crop back to original width\n", + " pred = seg_160_272[:, :raw_w]\n", + "\n", + " if raw_w < TARGET_W:\n", + " padded = np.full((TARGET_H * TARGET_W,), IGNORE_INDEX, dtype=np.int64)\n", + " padded[: TARGET_H * raw_w] = pred.flatten()\n", + " preds_dict[name] = padded\n", + " else:\n", + " preds_dict[name] = pred.flatten()\n", + "\n", + " sub = pd.DataFrame(preds_dict, dtype=\"int64\").T\n", + " sub.to_csv(out_csv_path)\n", + " print(f\"[OK] submission saved to: {out_csv_path}\")\n", + "\n", + "\n", + "# =========================\n", + "# 7. Main\n", + "# =========================\n", + "def main():\n", + " print(f\"DEVICE: {DEVICE}\")\n", + " print(f\"Train dir: {TRAIN_IMAGES_DIR}\")\n", + " print(f\"Test dir: {TEST_IMAGES_DIR}\")\n", + " print(f\"Model input size: {MODEL_H}x{MODEL_W}\")\n", + " print(f\"Pretrained: {PRETRAINED}\")\n", + "\n", + " # Load all train data (well1-6)\n", + " train_ds_all = WellSegDataset(TRAIN_IMAGES_DIR, Y_TRAIN_CSV)\n", + "\n", + " # Split by well: well6 as validation\n", + " VAL_WELLS = {6}\n", + " train_indices, val_indices = [], []\n", + " for i, name in enumerate(train_ds_all.names):\n", + " w = parse_well_id(name)\n", + " if w in VAL_WELLS:\n", + " val_indices.append(i)\n", + " else:\n", + " train_indices.append(i)\n", + "\n", + " train_ds = Subset(train_ds_all, train_indices) # well1-5\n", + " val_ds = Subset(train_ds_all, val_indices) # well6\n", + "\n", + " train_loader = DataLoader(\n", + " train_ds,\n", + " batch_size=BATCH_SIZE,\n", + " shuffle=True,\n", + " num_workers=0,\n", + " collate_fn=collate_mask2former,\n", + " )\n", + " val_loader = DataLoader(\n", + " val_ds,\n", + " batch_size=BATCH_SIZE,\n", + " shuffle=False,\n", + " num_workers=0,\n", + " collate_fn=collate_mask2former,\n", + " )\n", + "\n", + " print(f\"Train samples: {len(train_ds)} | Val samples: {len(val_ds)} | val_wells={VAL_WELLS}\")\n", + "\n", + " model, processor = build_model_and_processor(NUM_CLASSES)\n", + " model = model.to(DEVICE)\n", + "\n", + " optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)\n", + "\n", + " best_val = 1e9\n", + " best_path = DATA_ROOT / \"best_mask2former.pth\"\n", + "\n", + " for epoch in range(1, EPOCHS + 1):\n", + " tr_loss = train_one_epoch(model, train_loader, optimizer)\n", + " va_loss = valid_one_epoch(model, val_loader)\n", + " print(f\"Epoch {epoch:02d}/{EPOCHS} | train_loss={tr_loss:.4f} | val_loss={va_loss:.4f}\")\n", + "\n", + " if va_loss < best_val:\n", + " best_val = va_loss\n", + " torch.save(model.state_dict(), best_path)\n", + " print(f\" -> Best model saved: {best_path}\")\n", + "\n", + " # Predict test and write submission\n", + " out_csv = DATA_ROOT / \"submission.csv\"\n", + " state_dict = torch.load(best_path, map_location=DEVICE, weights_only=True)\n", + " model.load_state_dict(state_dict)\n", + " predict_and_make_submission(model, processor, TEST_IMAGES_DIR, out_csv)\n", + "\n", + "\n", + "if __name__ == \"__main__\":\n", + " main()\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "deep-torch", + "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.10.19" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/swin-unet.ipynb b/swin-unet.ipynb new file mode 100644 index 0000000..3074856 --- /dev/null +++ b/swin-unet.ipynb @@ -0,0 +1,434 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 2, + "id": "e433fe32", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "DEVICE: cuda\n", + "Train dir: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\X_train_uDRk9z9\\images\n", + "Test dir: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\X_test_xNbnvIa\\images\n", + "Model input size: 224x224\n", + "Train samples: 2790 | Val samples: 1620 | val_wells={6}\n", + "Epoch 01/20 | train_loss=0.2381 | val_loss=0.0994\n", + " -> Best model saved: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_swin_unet.pth\n", + "Epoch 02/20 | train_loss=0.0835 | val_loss=0.0767\n", + " -> Best model saved: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_swin_unet.pth\n", + "Epoch 03/20 | train_loss=0.0714 | val_loss=0.0793\n", + "Epoch 04/20 | train_loss=0.0675 | val_loss=0.0746\n", + " -> Best model saved: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_swin_unet.pth\n", + "Epoch 05/20 | train_loss=0.0666 | val_loss=0.0761\n", + "Epoch 06/20 | train_loss=0.0645 | val_loss=0.0755\n", + "Epoch 07/20 | train_loss=0.0617 | val_loss=0.0773\n", + "Epoch 08/20 | train_loss=0.0612 | val_loss=0.0849\n", + "Epoch 09/20 | train_loss=0.0611 | val_loss=0.0784\n", + "Epoch 10/20 | train_loss=0.0592 | val_loss=0.0773\n", + "Epoch 11/20 | train_loss=0.0579 | val_loss=0.0828\n", + "Epoch 12/20 | train_loss=0.0592 | val_loss=0.0735\n", + " -> Best model saved: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_swin_unet.pth\n", + "Epoch 13/20 | train_loss=0.0567 | val_loss=0.0888\n", + "Epoch 14/20 | train_loss=0.0567 | val_loss=0.0778\n", + "Epoch 15/20 | train_loss=0.0560 | val_loss=0.0914\n", + "Epoch 16/20 | train_loss=0.0553 | val_loss=0.0889\n", + "Epoch 17/20 | train_loss=0.0540 | val_loss=0.0812\n", + "Epoch 18/20 | train_loss=0.0536 | val_loss=0.0970\n", + "Epoch 19/20 | train_loss=0.0535 | val_loss=0.0803\n", + "Epoch 20/20 | train_loss=0.0532 | val_loss=0.0901\n", + "[OK] submission saved to: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\submission.csv\n" + ] + } + ], + "source": [ + "# -*- coding: utf-8 -*-\n", + "\"\"\"\n", + "Swin-UNet (Swin Transformer encoder + UNet decoder) - Full runnable code (224x224 input for Swin)\n", + "\n", + "Pipeline:\n", + "- Train images: X_train_uDRk9z9/images (well1-6)\n", + "- Test images: X_test_xNbnvIa/images (well7-11)\n", + "- Train labels: Y_train_T9NrBYo.csv (flatten + -1 padding)\n", + "- Validation split: well6 as val, well1-5 as train\n", + "- Output: submission.csv (each row = one patch, flattened, padded to 160*272 with -1)\n", + "\n", + "Key point:\n", + "- Swin encoders in timm/SMP often expect 224x224 inputs.\n", + "- We resize image/mask to 224x224 for training and inference.\n", + "- During inference we upsample logits back to (160,272), then crop to raw width and pad with -1 to match submission format.\n", + "\n", + "Dependencies:\n", + " pip install timm segmentation-models-pytorch\n", + "\"\"\"\n", + "\n", + "import re\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "import pandas as pd\n", + "\n", + "import torch\n", + "import torch.nn.functional as F\n", + "from torch.utils.data import Dataset, DataLoader, Subset\n", + "\n", + "import timm\n", + "import segmentation_models_pytorch as smp\n", + "\n", + "\n", + "# =========================\n", + "# 0. Paths & Hyperparameters\n", + "# =========================\n", + "DATA_ROOT = Path(r\"C:\\Users\\lenovo\\Desktop\\deep_datachallenge\") # change to your path\n", + "\n", + "TRAIN_IMAGES_DIR = DATA_ROOT / \"X_train_uDRk9z9\" / \"images\"\n", + "TEST_IMAGES_DIR = DATA_ROOT / \"X_test_xNbnvIa\" / \"images\"\n", + "Y_TRAIN_CSV = DATA_ROOT / \"Y_train_T9NrBYo.csv\"\n", + "\n", + "# Original target size used by dataset/submission format\n", + "TARGET_H = 160\n", + "TARGET_W = 272\n", + "\n", + "# Model input size for Swin\n", + "MODEL_H = 224\n", + "MODEL_W = 224\n", + "\n", + "NUM_CLASSES = 3\n", + "IGNORE_INDEX = -1\n", + "\n", + "BATCH_SIZE = 4 # for RTX 4060(8GB), start with 2~4\n", + "LR = 1e-4\n", + "WEIGHT_DECAY = 1e-4\n", + "EPOCHS = 20\n", + "\n", + "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", + "\n", + "\n", + "# =========================\n", + "# 1. Utils\n", + "# =========================\n", + "def parse_well_id(name: str) -> int:\n", + " \"\"\"Extract well id from: well_1_section_0_patch_0 -> 1\"\"\"\n", + " m = re.search(r\"well_(\\d+)_\", name)\n", + " return int(m.group(1)) if m else -1\n", + "\n", + "\n", + "def minmax_normalize(x: np.ndarray) -> np.ndarray:\n", + " \"\"\"Min-max normalize; replace NaN/inf with 0.\"\"\"\n", + " x = x.astype(np.float32)\n", + " x = np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)\n", + " x_min = float(x.min())\n", + " x_max = float(x.max())\n", + " if x_max - x_min < 1e-6:\n", + " return np.zeros_like(x, dtype=np.float32)\n", + " return (x - x_min) / (x_max - x_min)\n", + "\n", + "\n", + "def pad_to_160x272(img: np.ndarray, fill_value: float = 0.0) -> np.ndarray:\n", + " \"\"\"Pad (160,160) or (160,272) to (160,272).\"\"\"\n", + " h, w = img.shape\n", + " assert h == TARGET_H, f\"Expected height {TARGET_H}, got {h}\"\n", + " if w == TARGET_W:\n", + " return img\n", + " if w < TARGET_W:\n", + " out = np.full((TARGET_H, TARGET_W), fill_value, dtype=img.dtype)\n", + " out[:, :w] = img\n", + " return out\n", + " return img[:, :TARGET_W]\n", + "\n", + "\n", + "def decode_mask_from_csv_row(row_values: np.ndarray) -> np.ndarray:\n", + " \"\"\"\n", + " Decode a mask from one CSV row:\n", + " - row_values: flattened mask with -1 padding\n", + " - remove -1 then reshape to (160, w)\n", + " \"\"\"\n", + " valid = row_values[row_values != IGNORE_INDEX]\n", + " assert len(valid) % TARGET_H == 0, f\"Valid mask length {len(valid)} not divisible by 160\"\n", + " w = len(valid) // TARGET_H\n", + " return valid.reshape(TARGET_H, w).astype(np.int64)\n", + "\n", + "\n", + "def pad_mask_to_160x272(mask: np.ndarray) -> np.ndarray:\n", + " \"\"\"Pad (160,w) to (160,272) using -1 for padding (ignore_index).\"\"\"\n", + " h, w = mask.shape\n", + " assert h == TARGET_H\n", + " if w == TARGET_W:\n", + " return mask\n", + " out = np.full((TARGET_H, TARGET_W), IGNORE_INDEX, dtype=np.int64)\n", + " out[:, :w] = mask\n", + " return out\n", + "\n", + "\n", + "def resize_image_torch(img_t_1hw: torch.Tensor, h: int, w: int) -> torch.Tensor:\n", + " \"\"\"\n", + " Resize an image tensor (1,H,W) -> (1,h,w) using bilinear.\n", + " \"\"\"\n", + " x = img_t_1hw.unsqueeze(0) # (1,1,H,W)\n", + " x = F.interpolate(x, size=(h, w), mode=\"bilinear\", align_corners=False)\n", + " return x.squeeze(0) # (1,h,w)\n", + "\n", + "\n", + "def resize_mask_torch(mask_t_hw: torch.Tensor, h: int, w: int) -> torch.Tensor:\n", + " \"\"\"\n", + " Resize a mask tensor (H,W) -> (h,w) using nearest.\n", + " IGNORE_INDEX is kept as-is via nearest interpolation.\n", + " \"\"\"\n", + " y = mask_t_hw.unsqueeze(0).unsqueeze(0).float() # (1,1,H,W)\n", + " y = F.interpolate(y, size=(h, w), mode=\"nearest\")\n", + " return y.squeeze(0).squeeze(0).long() # (h,w)\n", + "\n", + "\n", + "# =========================\n", + "# 2. Dataset\n", + "# =========================\n", + "class WellSegDataset(Dataset):\n", + " def __init__(self, images_dir: Path, y_csv_path: Path = None):\n", + " \"\"\"\n", + " If y_csv_path is None => test mode (no labels).\n", + " \"\"\"\n", + " self.images_dir = images_dir\n", + " self.has_label = y_csv_path is not None\n", + "\n", + " self.image_paths = sorted(images_dir.glob(\"*.npy\"))\n", + " self.names = [p.stem for p in self.image_paths]\n", + "\n", + " if self.has_label:\n", + " self.y_df = pd.read_csv(y_csv_path, index_col=0)\n", + " else:\n", + " self.y_df = None\n", + "\n", + " def __len__(self):\n", + " return len(self.image_paths)\n", + "\n", + " def __getitem__(self, idx: int):\n", + " name = self.names[idx]\n", + " img_path = self.image_paths[idx]\n", + "\n", + " img = np.load(img_path) # (160,160) or (160,272)\n", + " raw_w = img.shape[1] # used to crop back at submission\n", + " img = minmax_normalize(img)\n", + " img = pad_to_160x272(img, fill_value=0.0)\n", + "\n", + " # image: (1,160,272) -> resize to (1,224,224) for Swin\n", + " img_t = torch.from_numpy(img).unsqueeze(0).float() # (1,160,272)\n", + " img_t_224 = resize_image_torch(img_t, MODEL_H, MODEL_W) # (1,224,224)\n", + "\n", + " if not self.has_label:\n", + " return {\"name\": name, \"image\": img_t_224, \"raw_w\": raw_w}\n", + "\n", + " # mask: decode -> pad to (160,272) -> resize to (224,224) for training\n", + " row = self.y_df.loc[name].values.astype(np.int64)\n", + " mask = decode_mask_from_csv_row(row) # (160,w)\n", + " mask = pad_mask_to_160x272(mask) # (160,272)\n", + " mask_t = torch.from_numpy(mask).long() # (160,272)\n", + " mask_t_224 = resize_mask_torch(mask_t, MODEL_H, MODEL_W) # (224,224)\n", + "\n", + " return {\"name\": name, \"image\": img_t_224, \"mask\": mask_t_224, \"raw_w\": raw_w}\n", + "\n", + "\n", + "# =========================\n", + "# 3. Swin-UNet model (SMP + timm)\n", + "# =========================\n", + "def choose_swin_encoder_name() -> str:\n", + " \"\"\"\n", + " Pick a Swin encoder name that exists in timm.\n", + " SMP uses timm models (sometimes with 'tu-' prefix).\n", + " \"\"\"\n", + " candidates = [\n", + " \"tu-swin_tiny_patch4_window7_224\",\n", + " \"tu-swin_small_patch4_window7_224\",\n", + " \"tu-swin_base_patch4_window7_224\",\n", + " \"swin_tiny_patch4_window7_224\",\n", + " \"swin_small_patch4_window7_224\",\n", + " \"swin_base_patch4_window7_224\",\n", + " ]\n", + "\n", + " timm_models = set(timm.list_models())\n", + " for name in candidates:\n", + " raw = name.replace(\"tu-\", \"\")\n", + " if raw in timm_models:\n", + " return name\n", + "\n", + " return \"tu-swin_tiny_patch4_window7_224\"\n", + "\n", + "\n", + "def build_swin_unet(num_classes: int) -> torch.nn.Module:\n", + " encoder_name = choose_swin_encoder_name()\n", + "\n", + " # If your machine cannot download weights, set encoder_weights=None\n", + " model = smp.Unet(\n", + " encoder_name=encoder_name,\n", + " encoder_weights=\"imagenet\",\n", + " in_channels=1,\n", + " classes=num_classes,\n", + " activation=None,\n", + " )\n", + " return model\n", + "\n", + "\n", + "# =========================\n", + "# 4. Train / Validate\n", + "# =========================\n", + "def train_one_epoch(model, loader, optimizer):\n", + " model.train()\n", + " total_loss = 0.0\n", + "\n", + " for batch in loader:\n", + " x = batch[\"image\"].to(DEVICE) # (B,1,224,224)\n", + " y = batch[\"mask\"].to(DEVICE) # (B,224,224)\n", + "\n", + " logits = model(x) # (B,C,224,224)\n", + " loss = F.cross_entropy(logits, y, ignore_index=IGNORE_INDEX)\n", + "\n", + " optimizer.zero_grad()\n", + " loss.backward()\n", + " optimizer.step()\n", + "\n", + " total_loss += loss.item() * x.size(0)\n", + "\n", + " return total_loss / len(loader.dataset)\n", + "\n", + "\n", + "@torch.no_grad()\n", + "def valid_one_epoch(model, loader):\n", + " model.eval()\n", + " total_loss = 0.0\n", + "\n", + " for batch in loader:\n", + " x = batch[\"image\"].to(DEVICE)\n", + " y = batch[\"mask\"].to(DEVICE)\n", + "\n", + " logits = model(x)\n", + " loss = F.cross_entropy(logits, y, ignore_index=IGNORE_INDEX)\n", + " total_loss += loss.item() * x.size(0)\n", + "\n", + " return total_loss / len(loader.dataset)\n", + "\n", + "\n", + "# =========================\n", + "# 5. Inference & submission\n", + "# =========================\n", + "@torch.no_grad()\n", + "def predict_and_make_submission(model, test_images_dir: Path, out_csv_path: Path):\n", + " \"\"\"\n", + " Predict all .npy in test_images_dir and write submission.csv\n", + " - Model runs on 224x224 input\n", + " - Logits are upsampled back to (160,272)\n", + " - Then we crop to raw width and pad with -1 to match submission format\n", + " \"\"\"\n", + " model.eval()\n", + "\n", + " test_ds = WellSegDataset(test_images_dir, y_csv_path=None)\n", + " test_loader = DataLoader(test_ds, batch_size=1, shuffle=False, num_workers=0)\n", + "\n", + " preds_dict = {}\n", + "\n", + " for batch in test_loader:\n", + " name = batch[\"name\"][0]\n", + " raw_w = int(batch[\"raw_w\"][0])\n", + " x = batch[\"image\"].to(DEVICE) # (1,1,224,224)\n", + "\n", + " logits_224 = model(x) # (1,C,224,224)\n", + " logits = F.interpolate(logits_224, size=(TARGET_H, TARGET_W), mode=\"bilinear\", align_corners=False)\n", + " pred_full = torch.argmax(logits, dim=1).squeeze(0).cpu().numpy().astype(np.int64) # (160,272)\n", + "\n", + " pred = pred_full[:, :raw_w] # crop back to original width\n", + "\n", + " if raw_w < TARGET_W:\n", + " padded = np.full((TARGET_H * TARGET_W,), IGNORE_INDEX, dtype=np.int64)\n", + " padded[: TARGET_H * raw_w] = pred.flatten()\n", + " preds_dict[name] = padded\n", + " else:\n", + " preds_dict[name] = pred.flatten()\n", + "\n", + " sub = pd.DataFrame(preds_dict, dtype=\"int64\").T\n", + " sub.to_csv(out_csv_path)\n", + " print(f\"[OK] submission saved to: {out_csv_path}\")\n", + "\n", + "\n", + "# =========================\n", + "# 6. Main\n", + "# =========================\n", + "def main():\n", + " print(f\"DEVICE: {DEVICE}\")\n", + " print(f\"Train dir: {TRAIN_IMAGES_DIR}\")\n", + " print(f\"Test dir: {TEST_IMAGES_DIR}\")\n", + " print(f\"Model input size: {MODEL_H}x{MODEL_W}\")\n", + "\n", + " # (A) Load all train data (well1-6)\n", + " train_ds_all = WellSegDataset(TRAIN_IMAGES_DIR, Y_TRAIN_CSV)\n", + "\n", + " # (B) Split by well: use well6 as validation\n", + " VAL_WELLS = {6}\n", + " train_indices, val_indices = [], []\n", + " for i, name in enumerate(train_ds_all.names):\n", + " w = parse_well_id(name)\n", + " if w in VAL_WELLS:\n", + " val_indices.append(i)\n", + " else:\n", + " train_indices.append(i)\n", + "\n", + " train_ds = Subset(train_ds_all, train_indices) # well1-5\n", + " val_ds = Subset(train_ds_all, val_indices) # well6\n", + "\n", + " train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)\n", + " val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)\n", + "\n", + " print(f\"Train samples: {len(train_ds)} | Val samples: {len(val_ds)} | val_wells={VAL_WELLS}\")\n", + "\n", + " # (C) Model & optimizer\n", + " model = build_swin_unet(num_classes=NUM_CLASSES).to(DEVICE)\n", + " optimizer = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)\n", + "\n", + " # (D) Train\n", + " best_val = 1e9\n", + " best_path = DATA_ROOT / \"best_swin_unet.pth\"\n", + "\n", + " for epoch in range(1, EPOCHS + 1):\n", + " tr_loss = train_one_epoch(model, train_loader, optimizer)\n", + " va_loss = valid_one_epoch(model, val_loader)\n", + "\n", + " print(f\"Epoch {epoch:02d}/{EPOCHS} | train_loss={tr_loss:.4f} | val_loss={va_loss:.4f}\")\n", + "\n", + " if va_loss < best_val:\n", + " best_val = va_loss\n", + " torch.save(model.state_dict(), best_path)\n", + " print(f\" -> Best model saved: {best_path}\")\n", + "\n", + " # (E) Predict test and write submission\n", + " out_csv = DATA_ROOT / \"submission.csv\"\n", + " state_dict = torch.load(best_path, map_location=DEVICE, weights_only=True)\n", + " model.load_state_dict(state_dict)\n", + " predict_and_make_submission(model, TEST_IMAGES_DIR, out_csv)\n", + "\n", + "\n", + "if __name__ == \"__main__\":\n", + " main()\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "deep-torch", + "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.10.19" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From 1ae43122b645365f923d70c49fd0275620622dcc Mon Sep 17 00:00:00 2001 From: Haokun Date: Sat, 10 Jan 2026 15:20:28 +0100 Subject: [PATCH 06/11] Create mask2former_semi_supervised_pseudo_label.py --- mask2former_semi_supervised_pseudo_label.py | 585 ++++++++++++++++++++ 1 file changed, 585 insertions(+) create mode 100644 mask2former_semi_supervised_pseudo_label.py diff --git a/mask2former_semi_supervised_pseudo_label.py b/mask2former_semi_supervised_pseudo_label.py new file mode 100644 index 0000000..db2514c --- /dev/null +++ b/mask2former_semi_supervised_pseudo_label.py @@ -0,0 +1,585 @@ +# -*- coding: utf-8 -*- +""" +Mask2Former + Semi-Supervised Pseudo Label (Full runnable) + +Data: +- Labeled train images: X_train_uDRk9z9/images (well1-6) +- Labeled train labels: Y_train_T9NrBYo.csv (flatten + -1 padding) +- Unlabeled images: X_unlabeled_mtkxUlo/images (well12-14) +- Test images: X_test_xNbnvIa/images (well7-11) + +Split: +- Train labeled: well1-5 +- Val labeled: well6 +- Unlabeled: well12-14 (no labels) + +Output: +- submission.csv, each row = one patch +- flattened mask, padded to 160*272 with -1 + +Install: + pip install transformers accelerate +""" + +import re +from pathlib import Path +from typing import Dict, List, Tuple, Optional + +import numpy as np +import pandas as pd + +import torch +import torch.nn.functional as F +from torch.utils.data import Dataset, DataLoader, Subset + +from transformers import ( + AutoImageProcessor, + Mask2FormerForUniversalSegmentation, +) + +# ========================= +# 0) Paths & Hyperparameters +# ========================= +DATA_ROOT = Path(r"C:\Users\lenovo\Desktop\deep_datachallenge") + +TRAIN_IMAGES_DIR = DATA_ROOT / "X_train_uDRk9z9" / "images" +TEST_IMAGES_DIR = DATA_ROOT / "X_test_xNbnvIa" / "images" +UNLABELED_DIR = DATA_ROOT / "X_unlabeled_mtkxUlo" / "images" +Y_TRAIN_CSV = DATA_ROOT / "Y_train_T9NrBYo.csv" + +# submission size +TARGET_H = 160 +TARGET_W = 272 + +# model size +MODEL_H = 224 +MODEL_W = 224 + +NUM_CLASSES = 3 +IGNORE_INDEX = -1 + +BATCH_SIZE_L = 2 # labeled batch +BATCH_SIZE_U = 2 # unlabeled batch +LR = 5e-5 +WEIGHT_DECAY = 1e-4 +EPOCHS = 10 + +# semi-supervised hyperparams +PSEUDO_TH = 0.85 # 伪标签置信度阈值(越高越保守) +LAMBDA_U = 0.5 # 无标签loss权重(0.2~1.0可调) + +DEVICE = "cuda" if torch.cuda.is_available() else "cpu" + +PRETRAINED = "facebook/mask2former-swin-tiny-ade-semantic" + + +# ========================= +# 1) Utils +# ========================= +def parse_well_id(name: str) -> int: + """well_12_section_0_patch_0 -> 12""" + m = re.search(r"well_(\d+)_", name) + return int(m.group(1)) if m else -1 + + +def minmax_normalize(x: np.ndarray) -> np.ndarray: + x = x.astype(np.float32) + x = np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0) + mn = float(x.min()) + mx = float(x.max()) + if mx - mn < 1e-6: + return np.zeros_like(x, dtype=np.float32) + return (x - mn) / (mx - mn) + + +def pad_to_160x272(img: np.ndarray, fill_value: float = 0.0) -> np.ndarray: + h, w = img.shape + assert h == TARGET_H, f"Expected height {TARGET_H}, got {h}" + if w == TARGET_W: + return img + if w < TARGET_W: + out = np.full((TARGET_H, TARGET_W), fill_value, dtype=img.dtype) + out[:, :w] = img + return out + return img[:, :TARGET_W] + + +def decode_mask_from_csv_row(row_values: np.ndarray) -> np.ndarray: + valid = row_values[row_values != IGNORE_INDEX] + assert len(valid) % TARGET_H == 0, f"Valid mask length {len(valid)} not divisible by 160" + w = len(valid) // TARGET_H + return valid.reshape(TARGET_H, w).astype(np.int64) + + +def pad_mask_to_160x272(mask: np.ndarray) -> np.ndarray: + h, w = mask.shape + assert h == TARGET_H + if w == TARGET_W: + return mask + out = np.full((TARGET_H, TARGET_W), IGNORE_INDEX, dtype=np.int64) + out[:, :w] = mask + return out + + +def resize_image_torch(img_1hw: torch.Tensor, h: int, w: int) -> torch.Tensor: + """(1,H,W)->(1,h,w) bilinear""" + x = img_1hw.unsqueeze(0) # (1,1,H,W) + x = F.interpolate(x, size=(h, w), mode="bilinear", align_corners=False) + return x.squeeze(0) + + +def resize_mask_torch(mask_hw: torch.Tensor, h: int, w: int) -> torch.Tensor: + """(H,W)->(h,w) nearest""" + y = mask_hw.unsqueeze(0).unsqueeze(0).float() + y = F.interpolate(y, size=(h, w), mode="nearest") + return y.squeeze(0).squeeze(0).long() + + +def semantic_to_mask2former_targets( + semantic_mask: torch.Tensor, + num_classes: int, + ignore_index: int = -1, +) -> Tuple[torch.Tensor, torch.Tensor]: + """ + semantic_mask: (H,W) with ignore_index + return: + class_labels: (N,) + mask_labels: (N,H,W) float(0/1) + """ + valid = semantic_mask != ignore_index + if valid.sum() == 0: + class_labels = torch.tensor([0], dtype=torch.long) + mask_labels = torch.zeros((1, semantic_mask.shape[0], semantic_mask.shape[1]), dtype=torch.float32) + return class_labels, mask_labels + + present = torch.unique(semantic_mask[valid]).tolist() + present = [int(c) for c in present if 0 <= int(c) < num_classes] + if len(present) == 0: + class_labels = torch.tensor([0], dtype=torch.long) + mask_labels = torch.zeros((1, semantic_mask.shape[0], semantic_mask.shape[1]), dtype=torch.float32) + return class_labels, mask_labels + + masks, classes = [], [] + for c in present: + m = (semantic_mask == c) & valid + if m.sum() == 0: + continue + masks.append(m.float()) + classes.append(c) + + if len(classes) == 0: + class_labels = torch.tensor([0], dtype=torch.long) + mask_labels = torch.zeros((1, semantic_mask.shape[0], semantic_mask.shape[1]), dtype=torch.float32) + return class_labels, mask_labels + + class_labels = torch.tensor(classes, dtype=torch.long) + mask_labels = torch.stack(masks, dim=0).float() + return class_labels, mask_labels + + +# ========================= +# 2) 简单增强（无标签用） +# ========================= +def aug_weak(x: torch.Tensor) -> torch.Tensor: + """弱增强：随机左右翻转 + 轻噪声""" + # x: (1,224,224) + if torch.rand(1).item() < 0.5: + x = torch.flip(x, dims=[2]) + noise = 0.02 * torch.randn_like(x) + return torch.clamp(x + noise, 0.0, 1.0) + + +def aug_strong(x: torch.Tensor) -> torch.Tensor: + """强增强：随机翻转 + 更强噪声 + 亮度对比度扰动""" + if torch.rand(1).item() < 0.5: + x = torch.flip(x, dims=[2]) + # brightness/contrast + contrast = 0.8 + 0.4 * torch.rand(1).item() # [0.8,1.2] + brightness = -0.1 + 0.2 * torch.rand(1).item() # [-0.1,0.1] + x = x * contrast + brightness + # noise + noise = 0.05 * torch.randn_like(x) + x = x + noise + return torch.clamp(x, 0.0, 1.0) + + +# ========================= +# 3) Dataset +# ========================= +class LabeledDataset(Dataset): + def __init__(self, images_dir: Path, y_csv_path: Path): + self.image_paths = sorted(images_dir.glob("*.npy")) + self.names = [p.stem for p in self.image_paths] + self.y_df = pd.read_csv(y_csv_path, index_col=0) + + def __len__(self): + return len(self.image_paths) + + def __getitem__(self, idx: int): + name = self.names[idx] + img = np.load(self.image_paths[idx]) + raw_w = int(img.shape[1]) + + img = minmax_normalize(img) + img = pad_to_160x272(img, fill_value=0.0) + img_t = torch.from_numpy(img).unsqueeze(0).float() # (1,160,272) + img_t = resize_image_torch(img_t, MODEL_H, MODEL_W) # (1,224,224) + + row = self.y_df.loc[name].values.astype(np.int64) + mask = decode_mask_from_csv_row(row) # (160,w) + mask = pad_mask_to_160x272(mask) # (160,272) + mask_t = torch.from_numpy(mask).long() # (160,272) + mask_t = resize_mask_torch(mask_t, MODEL_H, MODEL_W) # (224,224) + + return {"name": name, "image": img_t, "mask": mask_t, "raw_w": raw_w} + + +class UnlabeledDataset(Dataset): + def __init__(self, images_dir: Path): + self.image_paths = sorted(images_dir.glob("*.npy")) + self.names = [p.stem for p in self.image_paths] + + def __len__(self): + return len(self.image_paths) + + def __getitem__(self, idx: int): + name = self.names[idx] + img = np.load(self.image_paths[idx]) + raw_w = int(img.shape[1]) + + img = minmax_normalize(img) + img = pad_to_160x272(img, fill_value=0.0) + img_t = torch.from_numpy(img).unsqueeze(0).float() # (1,160,272) + img_t = resize_image_torch(img_t, MODEL_H, MODEL_W) # (1,224,224) + + # 返回 base 图像（增强在 collate 做） + return {"name": name, "image": img_t, "raw_w": raw_w} + + +# ========================= +# 4) Collate +# ========================= +def collate_labeled(batch: List[Dict]) -> Dict: + names = [b["name"] for b in batch] + raw_ws = torch.tensor([b["raw_w"] for b in batch], dtype=torch.long) + + imgs_1 = torch.stack([b["image"] for b in batch], dim=0) # (B,1,224,224) + pixel_values = imgs_1.repeat(1, 3, 1, 1) # (B,3,224,224) + pixel_mask = torch.ones((pixel_values.shape[0], MODEL_H, MODEL_W), dtype=torch.long) + + class_labels_list, mask_labels_list = [], [] + for b in batch: + y = b["mask"] # (224,224) + cls, msk = semantic_to_mask2former_targets(y, NUM_CLASSES, IGNORE_INDEX) + class_labels_list.append(cls) + mask_labels_list.append(msk) + + return { + "names": names, + "raw_ws": raw_ws, + "pixel_values": pixel_values, + "pixel_mask": pixel_mask, + "class_labels": class_labels_list, + "mask_labels": mask_labels_list, + } + + +def collate_unlabeled(batch: List[Dict]) -> Dict: + names = [b["name"] for b in batch] + + imgs = [b["image"] for b in batch] # list of (1,224,224) + + # weak / strong augmentation + imgs_w = torch.stack([aug_weak(x.clone()) for x in imgs], dim=0) # (B,1,224,224) + imgs_s = torch.stack([aug_strong(x.clone()) for x in imgs], dim=0) # (B,1,224,224) + + pixel_values_w = imgs_w.repeat(1, 3, 1, 1) # (B,3,224,224) + pixel_values_s = imgs_s.repeat(1, 3, 1, 1) + + pixel_mask = torch.ones((pixel_values_w.shape[0], MODEL_H, MODEL_W), dtype=torch.long) + + return { + "names": names, + "pixel_values_w": pixel_values_w, + "pixel_values_s": pixel_values_s, + "pixel_mask": pixel_mask, + } + + +# ========================= +# 5) Model builder +# ========================= +def build_model(num_classes: int): + id2label = {0: "class0", 1: "class1", 2: "class2"} + label2id = {v: k for k, v in id2label.items()} + + processor = AutoImageProcessor.from_pretrained(PRETRAINED) + model = Mask2FormerForUniversalSegmentation.from_pretrained( + PRETRAINED, + ignore_mismatched_sizes=True, + id2label=id2label, + label2id=label2id, + num_labels=num_classes, + ) + return model, processor + + +# ========================= +# 6) Pseudo label from Mask2Former outputs +# ========================= +@torch.no_grad() +def pseudo_from_outputs(outputs, num_classes: int) -> Tuple[torch.Tensor, torch.Tensor]: + """ + 从 Mask2Former 输出构造像素级类别概率： + class_probs (softmax) * mask_probs (sigmoid) -> per-pixel scores + + outputs.class_queries_logits: (B, Q, C+1) (最后一个通常是 no-object) + outputs.masks_queries_logits: (B, Q, H, W) + + return: + pseudo: (B,H,W) long (0..C-1) + conf: (B,H,W) float (max score) + """ + class_logits = outputs.class_queries_logits # (B,Q,C+1) + mask_logits = outputs.masks_queries_logits # (B,Q,H,W) + + class_prob = class_logits.softmax(dim=-1)[..., :num_classes] # (B,Q,C) + mask_prob = mask_logits.sigmoid() # (B,Q,H,W) + + # (B,C,H,W) einsum: sum_q class_prob[b,q,c] * mask_prob[b,q,h,w] + score = torch.einsum("bqc,bqhw->bchw", class_prob, mask_prob) + conf, pseudo = torch.max(score, dim=1) # (B,H,W) + return pseudo.long(), conf.float() + + +# ========================= +# 7) Train / Validate (Semi-Supervised) +# ========================= +def train_one_epoch_semi(model, labeled_loader, unlabeled_loader, optimizer): + model.train() + + total_l, total_u = 0.0, 0.0 + n_l, n_u = 0, 0 + + unlabeled_iter = iter(unlabeled_loader) + + for batch_l in labeled_loader: + # ---- labeled step ---- + pixel_values = batch_l["pixel_values"].to(DEVICE) + pixel_mask = batch_l["pixel_mask"].to(DEVICE) + class_labels = [x.to(DEVICE) for x in batch_l["class_labels"]] + mask_labels = [x.to(DEVICE) for x in batch_l["mask_labels"]] + + out_l = model( + pixel_values=pixel_values, + pixel_mask=pixel_mask, + class_labels=class_labels, + mask_labels=mask_labels, + ) + loss_l = out_l.loss + + # ---- unlabeled step (pseudo-label) ---- + try: + batch_u = next(unlabeled_iter) + except StopIteration: + unlabeled_iter = iter(unlabeled_loader) + batch_u = next(unlabeled_iter) + + pv_w = batch_u["pixel_values_w"].to(DEVICE) # weak + pv_s = batch_u["pixel_values_s"].to(DEVICE) # strong + pm_u = batch_u["pixel_mask"].to(DEVICE) + + # teacher prediction on weak + model.eval() + out_u_teacher = model(pixel_values=pv_w, pixel_mask=pm_u) + pseudo, conf = pseudo_from_outputs(out_u_teacher, NUM_CLASSES) # (B,224,224) + + # 置信度过滤：低于阈值的像素设为 IGNORE + pseudo = pseudo.clone() + pseudo[conf < PSEUDO_TH] = IGNORE_INDEX + + # 将 pseudo semantic mask -> mask2former targets(list) + class_labels_u, mask_labels_u = [], [] + for i in range(pseudo.shape[0]): + cls_i, msk_i = semantic_to_mask2former_targets(pseudo[i], NUM_CLASSES, IGNORE_INDEX) + class_labels_u.append(cls_i.to(DEVICE)) + mask_labels_u.append(msk_i.to(DEVICE)) + + model.train() + out_u_student = model( + pixel_values=pv_s, + pixel_mask=pm_u, + class_labels=class_labels_u, + mask_labels=mask_labels_u, + ) + loss_u = out_u_student.loss + + # ---- total loss ---- + loss = loss_l + LAMBDA_U * loss_u + + optimizer.zero_grad() + loss.backward() + optimizer.step() + + total_l += float(loss_l.item()) * pixel_values.size(0) + total_u += float(loss_u.item()) * pv_s.size(0) + n_l += pixel_values.size(0) + n_u += pv_s.size(0) + + return total_l / max(1, n_l), total_u / max(1, n_u) + + +@torch.no_grad() +def valid_one_epoch(model, loader): + model.eval() + total = 0.0 + n = 0 + + for batch in loader: + pixel_values = batch["pixel_values"].to(DEVICE) + pixel_mask = batch["pixel_mask"].to(DEVICE) + class_labels = [x.to(DEVICE) for x in batch["class_labels"]] + mask_labels = [x.to(DEVICE) for x in batch["mask_labels"]] + + out = model( + pixel_values=pixel_values, + pixel_mask=pixel_mask, + class_labels=class_labels, + mask_labels=mask_labels, + ) + total += float(out.loss.item()) * pixel_values.size(0) + n += pixel_values.size(0) + + return total / max(1, n) + + +# ========================= +# 8) Inference & submission +# ========================= +@torch.no_grad() +def predict_and_make_submission(model, processor, test_images_dir: Path, out_csv_path: Path): + model.eval() + + # 这里复用 UnlabeledDataset 结构（只有图像，无标签） + test_ds = UnlabeledDataset(test_images_dir) + + def collate_test(batch: List[Dict]) -> Dict: + names = [b["name"] for b in batch] + raw_ws = torch.tensor([b["raw_w"] for b in batch], dtype=torch.long) + imgs_1 = torch.stack([b["image"] for b in batch], dim=0) # (B,1,224,224) + pixel_values = imgs_1.repeat(1, 3, 1, 1) + pixel_mask = torch.ones((pixel_values.shape[0], MODEL_H, MODEL_W), dtype=torch.long) + return {"names": names, "raw_ws": raw_ws, "pixel_values": pixel_values, "pixel_mask": pixel_mask} + + test_loader = DataLoader(test_ds, batch_size=1, shuffle=False, num_workers=0, collate_fn=collate_test) + + preds_dict = {} + + for batch in test_loader: + name = batch["names"][0] + raw_w = int(batch["raw_ws"][0].item()) + + pixel_values = batch["pixel_values"].to(DEVICE) + pixel_mask = batch["pixel_mask"].to(DEVICE) + + outputs = model(pixel_values=pixel_values, pixel_mask=pixel_mask) + + # 使用 processor 的语义后处理（稳定） + seg_list = processor.post_process_semantic_segmentation(outputs, target_sizes=[(MODEL_H, MODEL_W)]) + seg_224 = seg_list[0].to(torch.int64) # (224,224) + + seg_224 = seg_224.unsqueeze(0).unsqueeze(0).float() + seg_160_272 = F.interpolate(seg_224, size=(TARGET_H, TARGET_W), mode="nearest").squeeze(0).squeeze(0) + seg_160_272 = seg_160_272.cpu().numpy().astype(np.int64) + + pred = seg_160_272[:, :raw_w] + + padded = np.full((TARGET_H * TARGET_W,), IGNORE_INDEX, dtype=np.int64) + padded[: TARGET_H * raw_w] = pred.flatten() + preds_dict[name] = padded + + sub = pd.DataFrame(preds_dict, dtype="int64").T + sub.to_csv(out_csv_path) + print(f"[OK] submission saved to: {out_csv_path}") + + +# ========================= +# 9) Main +# ========================= +def main(): + print(f"DEVICE: {DEVICE}") + print(f"Labeled train dir: {TRAIN_IMAGES_DIR}") + print(f"Unlabeled dir: {UNLABELED_DIR}") + print(f"Test dir: {TEST_IMAGES_DIR}") + print(f"Pretrained: {PRETRAINED}") + print(f"Pseudo TH={PSEUDO_TH}, lambda_u={LAMBDA_U}") + + # ---- labeled dataset (well1-6) ---- + ds_all = LabeledDataset(TRAIN_IMAGES_DIR, Y_TRAIN_CSV) + + # split by well (val=6) + train_idx, val_idx = [], [] + for i, name in enumerate(ds_all.names): + w = parse_well_id(name) + if w == 6: + val_idx.append(i) + else: + train_idx.append(i) + + train_ds = Subset(ds_all, train_idx) # well1-5 + val_ds = Subset(ds_all, val_idx) # well6 + + train_loader = DataLoader( + train_ds, + batch_size=BATCH_SIZE_L, + shuffle=True, + num_workers=0, + collate_fn=collate_labeled, + ) + val_loader = DataLoader( + val_ds, + batch_size=BATCH_SIZE_L, + shuffle=False, + num_workers=0, + collate_fn=collate_labeled, + ) + + # ---- unlabeled dataset (well12-14) ---- + unlab_ds = UnlabeledDataset(UNLABELED_DIR) + unlab_loader = DataLoader( + unlab_ds, + batch_size=BATCH_SIZE_U, + shuffle=True, + num_workers=0, + collate_fn=collate_unlabeled, + ) + + print(f"Labeled train: {len(train_ds)} | Val: {len(val_ds)} | Unlabeled: {len(unlab_ds)}") + + model, processor = build_model(NUM_CLASSES) + model = model.to(DEVICE) + + optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY) + + best_val = 1e9 + best_path = DATA_ROOT / "best_mask2former_semi.pth" + + for epoch in range(1, EPOCHS + 1): + tr_l, tr_u = train_one_epoch_semi(model, train_loader, unlab_loader, optimizer) + va = valid_one_epoch(model, val_loader) + + print(f"Epoch {epoch:02d}/{EPOCHS} | train_l={tr_l:.4f} | train_u={tr_u:.4f} | val={va:.4f}") + + if va < best_val: + best_val = va + torch.save(model.state_dict(), best_path) + print(f" -> Best saved: {best_path}") + + # ---- inference ---- + out_csv = DATA_ROOT / "submission.csv" + model.load_state_dict(torch.load(best_path, map_location=DEVICE, weights_only=True)) + predict_and_make_submission(model, processor, TEST_IMAGES_DIR, out_csv) + + +if __name__ == "__main__": + main() From f675e214f18a13d36f7ed70060d01d19c341121f Mon Sep 17 00:00:00 2001 From: Haokun Date: Sat, 10 Jan 2026 15:23:17 +0100 Subject: [PATCH 07/11] ... --- mask2former_semi_supervised_pseudo_label.py => semi.py | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename mask2former_semi_supervised_pseudo_label.py => semi.py (100%) diff --git a/mask2former_semi_supervised_pseudo_label.py b/semi.py similarity index 100% rename from mask2former_semi_supervised_pseudo_label.py rename to semi.py From 600b8f43b7aa3bdb31db48351026e53cc368e87e Mon Sep 17 00:00:00 2001 From: Haokun Date: Sun, 11 Jan 2026 10:34:30 +0100 Subject: [PATCH 08/11] increase mask2former_semi_supervise --- README.md | 1 + "mask2_former_semi-supervis\303\251.ipynb" | 665 +++++++++++++++++++++ semi.py | 585 ------------------ 3 files changed, 666 insertions(+), 585 deletions(-) create mode 100644 "mask2_former_semi-supervis\303\251.ipynb" delete mode 100644 semi.py diff --git a/README.md b/README.md index 042d30f..31645e3 100644 --- a/README.md +++ b/README.md @@ -46,6 +46,7 @@ transformer最强 ### 你已用 +* **Mask2Former + Semi-Supervised Pseudo Label (Full runnable)0.6663556049927326 * **SegFormer 0.649377** * **Mask2Former（更强但更复杂） 0.6441724128817745** * **UPerNet + Swin / ViT backbone 0.6402961815634437** diff --git "a/mask2_former_semi-supervis\303\251.ipynb" "b/mask2_former_semi-supervis\303\251.ipynb" new file mode 100644 index 0000000..141e7f0 --- /dev/null +++ "b/mask2_former_semi-supervis\303\251.ipynb" @@ -0,0 +1,665 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "id": "faa83615", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "DEVICE: cuda\n", + "Labeled train dir: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\X_train_uDRk9z9\\images\n", + "Unlabeled dir: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\X_unlabeled_mtkxUlo\\images\n", + "Test dir: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\X_test_xNbnvIa\\images\n", + "Pretrained: facebook/mask2former-swin-tiny-ade-semantic\n", + "Pseudo TH=0.85, lambda_u=0.5\n", + "Labeled train: 2790 | Val: 1620 | Unlabeled: 1980\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "c:\\Users\\lenovo\\anaconda3\\envs\\deep-torch\\lib\\site-packages\\transformers\\image_processing_base.py:417: UserWarning: The following named arguments are not valid for `Mask2FormerImageProcessor.__init__` and were ignored: '_max_size', 'reduce_labels'\n", + " image_processor = cls(**image_processor_dict)\n", + "Some weights of Mask2FormerForUniversalSegmentation were not initialized from the model checkpoint at facebook/mask2former-swin-tiny-ade-semantic and are newly initialized because the shapes did not match:\n", + "- class_predictor.weight: found shape torch.Size([151, 256]) in the checkpoint and torch.Size([4, 256]) in the model instantiated\n", + "- class_predictor.bias: found shape torch.Size([151]) in the checkpoint and torch.Size([4]) in the model instantiated\n", + "- criterion.empty_weight: found shape torch.Size([151]) in the checkpoint and torch.Size([4]) in the model instantiated\n", + "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Epoch 01/10 | train_l=18.6849 | train_u=41.4849 | val=15.4774\n", + " -> Best saved: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_mask2former_semi.pth\n", + "Epoch 02/10 | train_l=14.1083 | train_u=30.6973 | val=14.3693\n", + " -> Best saved: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_mask2former_semi.pth\n", + "Epoch 03/10 | train_l=13.2919 | train_u=31.3983 | val=14.7217\n", + "Epoch 04/10 | train_l=12.4842 | train_u=23.7818 | val=15.1125\n", + "Epoch 05/10 | train_l=11.9995 | train_u=22.2180 | val=15.1217\n", + "Epoch 06/10 | train_l=11.6591 | train_u=20.8557 | val=14.9213\n", + "Epoch 07/10 | train_l=11.2987 | train_u=18.4966 | val=15.2570\n", + "Epoch 08/10 | train_l=10.9792 | train_u=17.3398 | val=14.8278\n", + "Epoch 09/10 | train_l=10.6290 | train_u=17.2760 | val=14.8505\n", + "Epoch 10/10 | train_l=10.3889 | train_u=17.5382 | val=14.5758\n", + "[OK] submission saved to: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\submission.csv\n" + ] + } + ], + "source": [ + "\n", + "\"\"\"\n", + "Mask2Former + Semi-Supervised Pseudo Label (Full runnable)0.6663556049927326\n", + "\n", + "Data:\n", + "- Labeled train images: X_train_uDRk9z9/images (well1-6)\n", + "- Labeled train labels: Y_train_T9NrBYo.csv (flatten + -1 padding)\n", + "- Unlabeled images: X_unlabeled_mtkxUlo/images (well12-14)\n", + "- Test images: X_test_xNbnvIa/images (well7-11)\n", + "\n", + "Split:\n", + "- Train labeled: well1-5\n", + "- Val labeled: well6\n", + "- Unlabeled: well12-14 (no labels)\n", + "\n", + "Output:\n", + "- submission.csv, each row = one patch\n", + "- flattened mask, padded to 160*272 with -1\n", + "\n", + "Install:\n", + " pip install transformers accelerate\n", + "\"\"\"\n", + "\n", + "import re\n", + "from pathlib import Path\n", + "from typing import Dict, List, Tuple, Optional\n", + "\n", + "import numpy as np\n", + "import pandas as pd\n", + "\n", + "import torch\n", + "import torch.nn.functional as F\n", + "from torch.utils.data import Dataset, DataLoader, Subset\n", + "\n", + "from transformers import (\n", + " AutoImageProcessor,\n", + " Mask2FormerForUniversalSegmentation,\n", + ")\n", + "\n", + "# =========================\n", + "# 0) Paths & Hyperparameters\n", + "# =========================\n", + "DATA_ROOT = Path(r\"C:\\Users\\lenovo\\Desktop\\deep_datachallenge\")\n", + "\n", + "TRAIN_IMAGES_DIR = DATA_ROOT / \"X_train_uDRk9z9\" / \"images\"\n", + "TEST_IMAGES_DIR = DATA_ROOT / \"X_test_xNbnvIa\" / \"images\"\n", + "UNLABELED_DIR = DATA_ROOT / \"X_unlabeled_mtkxUlo\" / \"images\"\n", + "Y_TRAIN_CSV = DATA_ROOT / \"Y_train_T9NrBYo.csv\"\n", + "\n", + "# submission size\n", + "TARGET_H = 160\n", + "TARGET_W = 272\n", + "\n", + "# model size\n", + "MODEL_H = 224\n", + "MODEL_W = 224\n", + "\n", + "NUM_CLASSES = 3\n", + "IGNORE_INDEX = -1\n", + "\n", + "BATCH_SIZE_L = 2 # labeled batch\n", + "BATCH_SIZE_U = 2 # unlabeled batch\n", + "LR = 5e-5\n", + "WEIGHT_DECAY = 1e-4\n", + "EPOCHS = 10\n", + "\n", + "# semi-supervised hyperparams\n", + "PSEUDO_TH = 0.85 # 伪标签置信度阈值(越高越保守)\n", + "LAMBDA_U = 0.5 # 无标签loss权重(0.2~1.0可调)\n", + "\n", + "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", + "\n", + "PRETRAINED = \"facebook/mask2former-swin-tiny-ade-semantic\"\n", + "\n", + "\n", + "# =========================\n", + "# 1) Utils\n", + "# =========================\n", + "def parse_well_id(name: str) -> int:\n", + " \"\"\"well_12_section_0_patch_0 -> 12\"\"\"\n", + " m = re.search(r\"well_(\\d+)_\", name)\n", + " return int(m.group(1)) if m else -1\n", + "\n", + "\n", + "def minmax_normalize(x: np.ndarray) -> np.ndarray:\n", + " x = x.astype(np.float32)\n", + " x = np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)\n", + " mn = float(x.min())\n", + " mx = float(x.max())\n", + " if mx - mn < 1e-6:\n", + " return np.zeros_like(x, dtype=np.float32)\n", + " return (x - mn) / (mx - mn)\n", + "\n", + "\n", + "def pad_to_160x272(img: np.ndarray, fill_value: float = 0.0) -> np.ndarray:\n", + " h, w = img.shape\n", + " assert h == TARGET_H, f\"Expected height {TARGET_H}, got {h}\"\n", + " if w == TARGET_W:\n", + " return img\n", + " if w < TARGET_W:\n", + " out = np.full((TARGET_H, TARGET_W), fill_value, dtype=img.dtype)\n", + " out[:, :w] = img\n", + " return out\n", + " return img[:, :TARGET_W]\n", + "\n", + "\n", + "def decode_mask_from_csv_row(row_values: np.ndarray) -> np.ndarray:\n", + " valid = row_values[row_values != IGNORE_INDEX]\n", + " assert len(valid) % TARGET_H == 0, f\"Valid mask length {len(valid)} not divisible by 160\"\n", + " w = len(valid) // TARGET_H\n", + " return valid.reshape(TARGET_H, w).astype(np.int64)\n", + "\n", + "\n", + "def pad_mask_to_160x272(mask: np.ndarray) -> np.ndarray:\n", + " h, w = mask.shape\n", + " assert h == TARGET_H\n", + " if w == TARGET_W:\n", + " return mask\n", + " out = np.full((TARGET_H, TARGET_W), IGNORE_INDEX, dtype=np.int64)\n", + " out[:, :w] = mask\n", + " return out\n", + "\n", + "\n", + "def resize_image_torch(img_1hw: torch.Tensor, h: int, w: int) -> torch.Tensor:\n", + " \"\"\"(1,H,W)->(1,h,w) bilinear\"\"\"\n", + " x = img_1hw.unsqueeze(0) # (1,1,H,W)\n", + " x = F.interpolate(x, size=(h, w), mode=\"bilinear\", align_corners=False)\n", + " return x.squeeze(0)\n", + "\n", + "\n", + "def resize_mask_torch(mask_hw: torch.Tensor, h: int, w: int) -> torch.Tensor:\n", + " \"\"\"(H,W)->(h,w) nearest\"\"\"\n", + " y = mask_hw.unsqueeze(0).unsqueeze(0).float()\n", + " y = F.interpolate(y, size=(h, w), mode=\"nearest\")\n", + " return y.squeeze(0).squeeze(0).long()\n", + "\n", + "\n", + "def semantic_to_mask2former_targets(\n", + " semantic_mask: torch.Tensor,\n", + " num_classes: int,\n", + " ignore_index: int = -1,\n", + ") -> Tuple[torch.Tensor, torch.Tensor]:\n", + " \"\"\"\n", + " semantic_mask: (H,W) with ignore_index\n", + " return:\n", + " class_labels: (N,)\n", + " mask_labels: (N,H,W) float(0/1)\n", + " \"\"\"\n", + " valid = semantic_mask != ignore_index\n", + " if valid.sum() == 0:\n", + " class_labels = torch.tensor([0], dtype=torch.long)\n", + " mask_labels = torch.zeros((1, semantic_mask.shape[0], semantic_mask.shape[1]), dtype=torch.float32)\n", + " return class_labels, mask_labels\n", + "\n", + " present = torch.unique(semantic_mask[valid]).tolist()\n", + " present = [int(c) for c in present if 0 <= int(c) < num_classes]\n", + " if len(present) == 0:\n", + " class_labels = torch.tensor([0], dtype=torch.long)\n", + " mask_labels = torch.zeros((1, semantic_mask.shape[0], semantic_mask.shape[1]), dtype=torch.float32)\n", + " return class_labels, mask_labels\n", + "\n", + " masks, classes = [], []\n", + " for c in present:\n", + " m = (semantic_mask == c) & valid\n", + " if m.sum() == 0:\n", + " continue\n", + " masks.append(m.float())\n", + " classes.append(c)\n", + "\n", + " if len(classes) == 0:\n", + " class_labels = torch.tensor([0], dtype=torch.long)\n", + " mask_labels = torch.zeros((1, semantic_mask.shape[0], semantic_mask.shape[1]), dtype=torch.float32)\n", + " return class_labels, mask_labels\n", + "\n", + " class_labels = torch.tensor(classes, dtype=torch.long)\n", + " mask_labels = torch.stack(masks, dim=0).float()\n", + " return class_labels, mask_labels\n", + "\n", + "\n", + "# =========================\n", + "# 2) 简单增强（无标签用）\n", + "# =========================\n", + "def aug_weak(x: torch.Tensor) -> torch.Tensor:\n", + " \"\"\"弱增强：随机左右翻转 + 轻噪声\"\"\"\n", + " # x: (1,224,224)\n", + " if torch.rand(1).item() < 0.5:\n", + " x = torch.flip(x, dims=[2])\n", + " noise = 0.02 * torch.randn_like(x)\n", + " return torch.clamp(x + noise, 0.0, 1.0)\n", + "\n", + "\n", + "def aug_strong(x: torch.Tensor) -> torch.Tensor:\n", + " \"\"\"强增强：随机翻转 + 更强噪声 + 亮度对比度扰动\"\"\"\n", + " if torch.rand(1).item() < 0.5:\n", + " x = torch.flip(x, dims=[2])\n", + " # brightness/contrast\n", + " contrast = 0.8 + 0.4 * torch.rand(1).item() # [0.8,1.2]\n", + " brightness = -0.1 + 0.2 * torch.rand(1).item() # [-0.1,0.1]\n", + " x = x * contrast + brightness\n", + " # noise\n", + " noise = 0.05 * torch.randn_like(x)\n", + " x = x + noise\n", + " return torch.clamp(x, 0.0, 1.0)\n", + "\n", + "\n", + "# =========================\n", + "# 3) Dataset\n", + "# =========================\n", + "class LabeledDataset(Dataset):\n", + " def __init__(self, images_dir: Path, y_csv_path: Path):\n", + " self.image_paths = sorted(images_dir.glob(\"*.npy\"))\n", + " self.names = [p.stem for p in self.image_paths]\n", + " self.y_df = pd.read_csv(y_csv_path, index_col=0)\n", + "\n", + " def __len__(self):\n", + " return len(self.image_paths)\n", + "\n", + " def __getitem__(self, idx: int):\n", + " name = self.names[idx]\n", + " img = np.load(self.image_paths[idx])\n", + " raw_w = int(img.shape[1])\n", + "\n", + " img = minmax_normalize(img)\n", + " img = pad_to_160x272(img, fill_value=0.0)\n", + " img_t = torch.from_numpy(img).unsqueeze(0).float() # (1,160,272)\n", + " img_t = resize_image_torch(img_t, MODEL_H, MODEL_W) # (1,224,224)\n", + "\n", + " row = self.y_df.loc[name].values.astype(np.int64)\n", + " mask = decode_mask_from_csv_row(row) # (160,w)\n", + " mask = pad_mask_to_160x272(mask) # (160,272)\n", + " mask_t = torch.from_numpy(mask).long() # (160,272)\n", + " mask_t = resize_mask_torch(mask_t, MODEL_H, MODEL_W) # (224,224)\n", + "\n", + " return {\"name\": name, \"image\": img_t, \"mask\": mask_t, \"raw_w\": raw_w}\n", + "\n", + "\n", + "class UnlabeledDataset(Dataset):\n", + " def __init__(self, images_dir: Path):\n", + " self.image_paths = sorted(images_dir.glob(\"*.npy\"))\n", + " self.names = [p.stem for p in self.image_paths]\n", + "\n", + " def __len__(self):\n", + " return len(self.image_paths)\n", + "\n", + " def __getitem__(self, idx: int):\n", + " name = self.names[idx]\n", + " img = np.load(self.image_paths[idx])\n", + " raw_w = int(img.shape[1])\n", + "\n", + " img = minmax_normalize(img)\n", + " img = pad_to_160x272(img, fill_value=0.0)\n", + " img_t = torch.from_numpy(img).unsqueeze(0).float() # (1,160,272)\n", + " img_t = resize_image_torch(img_t, MODEL_H, MODEL_W) # (1,224,224)\n", + "\n", + " # 返回 base 图像（增强在 collate 做）\n", + " return {\"name\": name, \"image\": img_t, \"raw_w\": raw_w}\n", + "\n", + "\n", + "# =========================\n", + "# 4) Collate\n", + "# =========================\n", + "def collate_labeled(batch: List[Dict]) -> Dict:\n", + " names = [b[\"name\"] for b in batch]\n", + " raw_ws = torch.tensor([b[\"raw_w\"] for b in batch], dtype=torch.long)\n", + "\n", + " imgs_1 = torch.stack([b[\"image\"] for b in batch], dim=0) # (B,1,224,224)\n", + " pixel_values = imgs_1.repeat(1, 3, 1, 1) # (B,3,224,224)\n", + " pixel_mask = torch.ones((pixel_values.shape[0], MODEL_H, MODEL_W), dtype=torch.long)\n", + "\n", + " class_labels_list, mask_labels_list = [], []\n", + " for b in batch:\n", + " y = b[\"mask\"] # (224,224)\n", + " cls, msk = semantic_to_mask2former_targets(y, NUM_CLASSES, IGNORE_INDEX)\n", + " class_labels_list.append(cls)\n", + " mask_labels_list.append(msk)\n", + "\n", + " return {\n", + " \"names\": names,\n", + " \"raw_ws\": raw_ws,\n", + " \"pixel_values\": pixel_values,\n", + " \"pixel_mask\": pixel_mask,\n", + " \"class_labels\": class_labels_list,\n", + " \"mask_labels\": mask_labels_list,\n", + " }\n", + "\n", + "\n", + "def collate_unlabeled(batch: List[Dict]) -> Dict:\n", + " names = [b[\"name\"] for b in batch]\n", + "\n", + " imgs = [b[\"image\"] for b in batch] # list of (1,224,224)\n", + "\n", + " # weak / strong augmentation\n", + " imgs_w = torch.stack([aug_weak(x.clone()) for x in imgs], dim=0) # (B,1,224,224)\n", + " imgs_s = torch.stack([aug_strong(x.clone()) for x in imgs], dim=0) # (B,1,224,224)\n", + "\n", + " pixel_values_w = imgs_w.repeat(1, 3, 1, 1) # (B,3,224,224)\n", + " pixel_values_s = imgs_s.repeat(1, 3, 1, 1)\n", + "\n", + " pixel_mask = torch.ones((pixel_values_w.shape[0], MODEL_H, MODEL_W), dtype=torch.long)\n", + "\n", + " return {\n", + " \"names\": names,\n", + " \"pixel_values_w\": pixel_values_w,\n", + " \"pixel_values_s\": pixel_values_s,\n", + " \"pixel_mask\": pixel_mask,\n", + " }\n", + "\n", + "\n", + "# =========================\n", + "# 5) Model builder\n", + "# =========================\n", + "def build_model(num_classes: int):\n", + " id2label = {0: \"class0\", 1: \"class1\", 2: \"class2\"}\n", + " label2id = {v: k for k, v in id2label.items()}\n", + "\n", + " processor = AutoImageProcessor.from_pretrained(PRETRAINED)\n", + " model = Mask2FormerForUniversalSegmentation.from_pretrained(\n", + " PRETRAINED,\n", + " ignore_mismatched_sizes=True,\n", + " id2label=id2label,\n", + " label2id=label2id,\n", + " num_labels=num_classes,\n", + " )\n", + " return model, processor\n", + "\n", + "\n", + "# =========================\n", + "# 6) Pseudo label from Mask2Former outputs\n", + "# =========================\n", + "@torch.no_grad()\n", + "def pseudo_from_outputs(outputs, num_classes: int) -> Tuple[torch.Tensor, torch.Tensor]:\n", + " \"\"\"\n", + " 从 Mask2Former 输出构造像素级类别概率：\n", + " class_probs (softmax) * mask_probs (sigmoid) -> per-pixel scores\n", + "\n", + " outputs.class_queries_logits: (B, Q, C+1) (最后一个通常是 no-object)\n", + " outputs.masks_queries_logits: (B, Q, H, W)\n", + "\n", + " return:\n", + " pseudo: (B,H,W) long (0..C-1)\n", + " conf: (B,H,W) float (max score)\n", + " \"\"\"\n", + " class_logits = outputs.class_queries_logits # (B,Q,C+1)\n", + " mask_logits = outputs.masks_queries_logits # (B,Q,H,W)\n", + "\n", + " class_prob = class_logits.softmax(dim=-1)[..., :num_classes] # (B,Q,C)\n", + " mask_prob = mask_logits.sigmoid() # (B,Q,H,W)\n", + "\n", + " # (B,C,H,W) einsum: sum_q class_prob[b,q,c] * mask_prob[b,q,h,w]\n", + " score = torch.einsum(\"bqc,bqhw->bchw\", class_prob, mask_prob)\n", + " conf, pseudo = torch.max(score, dim=1) # (B,H,W)\n", + " return pseudo.long(), conf.float()\n", + "\n", + "\n", + "# =========================\n", + "# 7) Train / Validate (Semi-Supervised)\n", + "# =========================\n", + "def train_one_epoch_semi(model, labeled_loader, unlabeled_loader, optimizer):\n", + " model.train()\n", + "\n", + " total_l, total_u = 0.0, 0.0\n", + " n_l, n_u = 0, 0\n", + "\n", + " unlabeled_iter = iter(unlabeled_loader)\n", + "\n", + " for batch_l in labeled_loader:\n", + " # ---- labeled step ----\n", + " pixel_values = batch_l[\"pixel_values\"].to(DEVICE)\n", + " pixel_mask = batch_l[\"pixel_mask\"].to(DEVICE)\n", + " class_labels = [x.to(DEVICE) for x in batch_l[\"class_labels\"]]\n", + " mask_labels = [x.to(DEVICE) for x in batch_l[\"mask_labels\"]]\n", + "\n", + " out_l = model(\n", + " pixel_values=pixel_values,\n", + " pixel_mask=pixel_mask,\n", + " class_labels=class_labels,\n", + " mask_labels=mask_labels,\n", + " )\n", + " loss_l = out_l.loss\n", + "\n", + " # ---- unlabeled step (pseudo-label) ----\n", + " try:\n", + " batch_u = next(unlabeled_iter)\n", + " except StopIteration:\n", + " unlabeled_iter = iter(unlabeled_loader)\n", + " batch_u = next(unlabeled_iter)\n", + "\n", + " pv_w = batch_u[\"pixel_values_w\"].to(DEVICE) # weak\n", + " pv_s = batch_u[\"pixel_values_s\"].to(DEVICE) # strong\n", + " pm_u = batch_u[\"pixel_mask\"].to(DEVICE)\n", + "\n", + " # teacher prediction on weak\n", + " model.eval()\n", + " out_u_teacher = model(pixel_values=pv_w, pixel_mask=pm_u)\n", + " pseudo, conf = pseudo_from_outputs(out_u_teacher, NUM_CLASSES) # (B,224,224)\n", + "\n", + " # 置信度过滤：低于阈值的像素设为 IGNORE\n", + " pseudo = pseudo.clone()\n", + " pseudo[conf < PSEUDO_TH] = IGNORE_INDEX\n", + "\n", + " # 将 pseudo semantic mask -> mask2former targets(list)\n", + " class_labels_u, mask_labels_u = [], []\n", + " for i in range(pseudo.shape[0]):\n", + " cls_i, msk_i = semantic_to_mask2former_targets(pseudo[i], NUM_CLASSES, IGNORE_INDEX)\n", + " class_labels_u.append(cls_i.to(DEVICE))\n", + " mask_labels_u.append(msk_i.to(DEVICE))\n", + "\n", + " model.train()\n", + " out_u_student = model(\n", + " pixel_values=pv_s,\n", + " pixel_mask=pm_u,\n", + " class_labels=class_labels_u,\n", + " mask_labels=mask_labels_u,\n", + " )\n", + " loss_u = out_u_student.loss\n", + "\n", + " # ---- total loss ----\n", + " loss = loss_l + LAMBDA_U * loss_u\n", + "\n", + " optimizer.zero_grad()\n", + " loss.backward()\n", + " optimizer.step()\n", + "\n", + " total_l += float(loss_l.item()) * pixel_values.size(0)\n", + " total_u += float(loss_u.item()) * pv_s.size(0)\n", + " n_l += pixel_values.size(0)\n", + " n_u += pv_s.size(0)\n", + "\n", + " return total_l / max(1, n_l), total_u / max(1, n_u)\n", + "\n", + "\n", + "@torch.no_grad()\n", + "def valid_one_epoch(model, loader):\n", + " model.eval()\n", + " total = 0.0\n", + " n = 0\n", + "\n", + " for batch in loader:\n", + " pixel_values = batch[\"pixel_values\"].to(DEVICE)\n", + " pixel_mask = batch[\"pixel_mask\"].to(DEVICE)\n", + " class_labels = [x.to(DEVICE) for x in batch[\"class_labels\"]]\n", + " mask_labels = [x.to(DEVICE) for x in batch[\"mask_labels\"]]\n", + "\n", + " out = model(\n", + " pixel_values=pixel_values,\n", + " pixel_mask=pixel_mask,\n", + " class_labels=class_labels,\n", + " mask_labels=mask_labels,\n", + " )\n", + " total += float(out.loss.item()) * pixel_values.size(0)\n", + " n += pixel_values.size(0)\n", + "\n", + " return total / max(1, n)\n", + "\n", + "\n", + "# =========================\n", + "# 8) Inference & submission\n", + "# =========================\n", + "@torch.no_grad()\n", + "def predict_and_make_submission(model, processor, test_images_dir: Path, out_csv_path: Path):\n", + " model.eval()\n", + "\n", + " # 这里复用 UnlabeledDataset 结构（只有图像，无标签）\n", + " test_ds = UnlabeledDataset(test_images_dir)\n", + "\n", + " def collate_test(batch: List[Dict]) -> Dict:\n", + " names = [b[\"name\"] for b in batch]\n", + " raw_ws = torch.tensor([b[\"raw_w\"] for b in batch], dtype=torch.long)\n", + " imgs_1 = torch.stack([b[\"image\"] for b in batch], dim=0) # (B,1,224,224)\n", + " pixel_values = imgs_1.repeat(1, 3, 1, 1)\n", + " pixel_mask = torch.ones((pixel_values.shape[0], MODEL_H, MODEL_W), dtype=torch.long)\n", + " return {\"names\": names, \"raw_ws\": raw_ws, \"pixel_values\": pixel_values, \"pixel_mask\": pixel_mask}\n", + "\n", + " test_loader = DataLoader(test_ds, batch_size=1, shuffle=False, num_workers=0, collate_fn=collate_test)\n", + "\n", + " preds_dict = {}\n", + "\n", + " for batch in test_loader:\n", + " name = batch[\"names\"][0]\n", + " raw_w = int(batch[\"raw_ws\"][0].item())\n", + "\n", + " pixel_values = batch[\"pixel_values\"].to(DEVICE)\n", + " pixel_mask = batch[\"pixel_mask\"].to(DEVICE)\n", + "\n", + " outputs = model(pixel_values=pixel_values, pixel_mask=pixel_mask)\n", + "\n", + " # 使用 processor 的语义后处理（稳定）\n", + " seg_list = processor.post_process_semantic_segmentation(outputs, target_sizes=[(MODEL_H, MODEL_W)])\n", + " seg_224 = seg_list[0].to(torch.int64) # (224,224)\n", + "\n", + " seg_224 = seg_224.unsqueeze(0).unsqueeze(0).float()\n", + " seg_160_272 = F.interpolate(seg_224, size=(TARGET_H, TARGET_W), mode=\"nearest\").squeeze(0).squeeze(0)\n", + " seg_160_272 = seg_160_272.cpu().numpy().astype(np.int64)\n", + "\n", + " pred = seg_160_272[:, :raw_w]\n", + "\n", + " padded = np.full((TARGET_H * TARGET_W,), IGNORE_INDEX, dtype=np.int64)\n", + " padded[: TARGET_H * raw_w] = pred.flatten()\n", + " preds_dict[name] = padded\n", + "\n", + " sub = pd.DataFrame(preds_dict, dtype=\"int64\").T\n", + " sub.to_csv(out_csv_path)\n", + " print(f\"[OK] submission saved to: {out_csv_path}\")\n", + "\n", + "\n", + "# =========================\n", + "# 9) Main\n", + "# =========================\n", + "def main():\n", + " print(f\"DEVICE: {DEVICE}\")\n", + " print(f\"Labeled train dir: {TRAIN_IMAGES_DIR}\")\n", + " print(f\"Unlabeled dir: {UNLABELED_DIR}\")\n", + " print(f\"Test dir: {TEST_IMAGES_DIR}\")\n", + " print(f\"Pretrained: {PRETRAINED}\")\n", + " print(f\"Pseudo TH={PSEUDO_TH}, lambda_u={LAMBDA_U}\")\n", + "\n", + " # ---- labeled dataset (well1-6) ----\n", + " ds_all = LabeledDataset(TRAIN_IMAGES_DIR, Y_TRAIN_CSV)\n", + "\n", + " # split by well (val=6)\n", + " train_idx, val_idx = [], []\n", + " for i, name in enumerate(ds_all.names):\n", + " w = parse_well_id(name)\n", + " if w == 6:\n", + " val_idx.append(i)\n", + " else:\n", + " train_idx.append(i)\n", + "\n", + " train_ds = Subset(ds_all, train_idx) # well1-5\n", + " val_ds = Subset(ds_all, val_idx) # well6\n", + "\n", + " train_loader = DataLoader(\n", + " train_ds,\n", + " batch_size=BATCH_SIZE_L,\n", + " shuffle=True,\n", + " num_workers=0,\n", + " collate_fn=collate_labeled,\n", + " )\n", + " val_loader = DataLoader(\n", + " val_ds,\n", + " batch_size=BATCH_SIZE_L,\n", + " shuffle=False,\n", + " num_workers=0,\n", + " collate_fn=collate_labeled,\n", + " )\n", + "\n", + " # ---- unlabeled dataset (well12-14) ----\n", + " unlab_ds = UnlabeledDataset(UNLABELED_DIR)\n", + " unlab_loader = DataLoader(\n", + " unlab_ds,\n", + " batch_size=BATCH_SIZE_U,\n", + " shuffle=True,\n", + " num_workers=0,\n", + " collate_fn=collate_unlabeled,\n", + " )\n", + "\n", + " print(f\"Labeled train: {len(train_ds)} | Val: {len(val_ds)} | Unlabeled: {len(unlab_ds)}\")\n", + "\n", + " model, processor = build_model(NUM_CLASSES)\n", + " model = model.to(DEVICE)\n", + "\n", + " optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)\n", + "\n", + " best_val = 1e9\n", + " best_path = DATA_ROOT / \"best_mask2former_semi.pth\"\n", + "\n", + " for epoch in range(1, EPOCHS + 1):\n", + " tr_l, tr_u = train_one_epoch_semi(model, train_loader, unlab_loader, optimizer)\n", + " va = valid_one_epoch(model, val_loader)\n", + "\n", + " print(f\"Epoch {epoch:02d}/{EPOCHS} | train_l={tr_l:.4f} | train_u={tr_u:.4f} | val={va:.4f}\")\n", + "\n", + " if va < best_val:\n", + " best_val = va\n", + " torch.save(model.state_dict(), best_path)\n", + " print(f\" -> Best saved: {best_path}\")\n", + "\n", + " # ---- inference ----\n", + " out_csv = DATA_ROOT / \"submission.csv\"\n", + " model.load_state_dict(torch.load(best_path, map_location=DEVICE, weights_only=True))\n", + " predict_and_make_submission(model, processor, TEST_IMAGES_DIR, out_csv)\n", + "\n", + "\n", + "if __name__ == \"__main__\":\n", + " main()\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "deep-torch", + "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.10.19" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/semi.py b/semi.py deleted file mode 100644 index db2514c..0000000 --- a/semi.py +++ /dev/null @@ -1,585 +0,0 @@ -# -*- coding: utf-8 -*- -""" -Mask2Former + Semi-Supervised Pseudo Label (Full runnable) - -Data: -- Labeled train images: X_train_uDRk9z9/images (well1-6) -- Labeled train labels: Y_train_T9NrBYo.csv (flatten + -1 padding) -- Unlabeled images: X_unlabeled_mtkxUlo/images (well12-14) -- Test images: X_test_xNbnvIa/images (well7-11) - -Split: -- Train labeled: well1-5 -- Val labeled: well6 -- Unlabeled: well12-14 (no labels) - -Output: -- submission.csv, each row = one patch -- flattened mask, padded to 160*272 with -1 - -Install: - pip install transformers accelerate -""" - -import re -from pathlib import Path -from typing import Dict, List, Tuple, Optional - -import numpy as np -import pandas as pd - -import torch -import torch.nn.functional as F -from torch.utils.data import Dataset, DataLoader, Subset - -from transformers import ( - AutoImageProcessor, - Mask2FormerForUniversalSegmentation, -) - -# ========================= -# 0) Paths & Hyperparameters -# ========================= -DATA_ROOT = Path(r"C:\Users\lenovo\Desktop\deep_datachallenge") - -TRAIN_IMAGES_DIR = DATA_ROOT / "X_train_uDRk9z9" / "images" -TEST_IMAGES_DIR = DATA_ROOT / "X_test_xNbnvIa" / "images" -UNLABELED_DIR = DATA_ROOT / "X_unlabeled_mtkxUlo" / "images" -Y_TRAIN_CSV = DATA_ROOT / "Y_train_T9NrBYo.csv" - -# submission size -TARGET_H = 160 -TARGET_W = 272 - -# model size -MODEL_H = 224 -MODEL_W = 224 - -NUM_CLASSES = 3 -IGNORE_INDEX = -1 - -BATCH_SIZE_L = 2 # labeled batch -BATCH_SIZE_U = 2 # unlabeled batch -LR = 5e-5 -WEIGHT_DECAY = 1e-4 -EPOCHS = 10 - -# semi-supervised hyperparams -PSEUDO_TH = 0.85 # 伪标签置信度阈值(越高越保守) -LAMBDA_U = 0.5 # 无标签loss权重(0.2~1.0可调) - -DEVICE = "cuda" if torch.cuda.is_available() else "cpu" - -PRETRAINED = "facebook/mask2former-swin-tiny-ade-semantic" - - -# ========================= -# 1) Utils -# ========================= -def parse_well_id(name: str) -> int: - """well_12_section_0_patch_0 -> 12""" - m = re.search(r"well_(\d+)_", name) - return int(m.group(1)) if m else -1 - - -def minmax_normalize(x: np.ndarray) -> np.ndarray: - x = x.astype(np.float32) - x = np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0) - mn = float(x.min()) - mx = float(x.max()) - if mx - mn < 1e-6: - return np.zeros_like(x, dtype=np.float32) - return (x - mn) / (mx - mn) - - -def pad_to_160x272(img: np.ndarray, fill_value: float = 0.0) -> np.ndarray: - h, w = img.shape - assert h == TARGET_H, f"Expected height {TARGET_H}, got {h}" - if w == TARGET_W: - return img - if w < TARGET_W: - out = np.full((TARGET_H, TARGET_W), fill_value, dtype=img.dtype) - out[:, :w] = img - return out - return img[:, :TARGET_W] - - -def decode_mask_from_csv_row(row_values: np.ndarray) -> np.ndarray: - valid = row_values[row_values != IGNORE_INDEX] - assert len(valid) % TARGET_H == 0, f"Valid mask length {len(valid)} not divisible by 160" - w = len(valid) // TARGET_H - return valid.reshape(TARGET_H, w).astype(np.int64) - - -def pad_mask_to_160x272(mask: np.ndarray) -> np.ndarray: - h, w = mask.shape - assert h == TARGET_H - if w == TARGET_W: - return mask - out = np.full((TARGET_H, TARGET_W), IGNORE_INDEX, dtype=np.int64) - out[:, :w] = mask - return out - - -def resize_image_torch(img_1hw: torch.Tensor, h: int, w: int) -> torch.Tensor: - """(1,H,W)->(1,h,w) bilinear""" - x = img_1hw.unsqueeze(0) # (1,1,H,W) - x = F.interpolate(x, size=(h, w), mode="bilinear", align_corners=False) - return x.squeeze(0) - - -def resize_mask_torch(mask_hw: torch.Tensor, h: int, w: int) -> torch.Tensor: - """(H,W)->(h,w) nearest""" - y = mask_hw.unsqueeze(0).unsqueeze(0).float() - y = F.interpolate(y, size=(h, w), mode="nearest") - return y.squeeze(0).squeeze(0).long() - - -def semantic_to_mask2former_targets( - semantic_mask: torch.Tensor, - num_classes: int, - ignore_index: int = -1, -) -> Tuple[torch.Tensor, torch.Tensor]: - """ - semantic_mask: (H,W) with ignore_index - return: - class_labels: (N,) - mask_labels: (N,H,W) float(0/1) - """ - valid = semantic_mask != ignore_index - if valid.sum() == 0: - class_labels = torch.tensor([0], dtype=torch.long) - mask_labels = torch.zeros((1, semantic_mask.shape[0], semantic_mask.shape[1]), dtype=torch.float32) - return class_labels, mask_labels - - present = torch.unique(semantic_mask[valid]).tolist() - present = [int(c) for c in present if 0 <= int(c) < num_classes] - if len(present) == 0: - class_labels = torch.tensor([0], dtype=torch.long) - mask_labels = torch.zeros((1, semantic_mask.shape[0], semantic_mask.shape[1]), dtype=torch.float32) - return class_labels, mask_labels - - masks, classes = [], [] - for c in present: - m = (semantic_mask == c) & valid - if m.sum() == 0: - continue - masks.append(m.float()) - classes.append(c) - - if len(classes) == 0: - class_labels = torch.tensor([0], dtype=torch.long) - mask_labels = torch.zeros((1, semantic_mask.shape[0], semantic_mask.shape[1]), dtype=torch.float32) - return class_labels, mask_labels - - class_labels = torch.tensor(classes, dtype=torch.long) - mask_labels = torch.stack(masks, dim=0).float() - return class_labels, mask_labels - - -# ========================= -# 2) 简单增强（无标签用） -# ========================= -def aug_weak(x: torch.Tensor) -> torch.Tensor: - """弱增强：随机左右翻转 + 轻噪声""" - # x: (1,224,224) - if torch.rand(1).item() < 0.5: - x = torch.flip(x, dims=[2]) - noise = 0.02 * torch.randn_like(x) - return torch.clamp(x + noise, 0.0, 1.0) - - -def aug_strong(x: torch.Tensor) -> torch.Tensor: - """强增强：随机翻转 + 更强噪声 + 亮度对比度扰动""" - if torch.rand(1).item() < 0.5: - x = torch.flip(x, dims=[2]) - # brightness/contrast - contrast = 0.8 + 0.4 * torch.rand(1).item() # [0.8,1.2] - brightness = -0.1 + 0.2 * torch.rand(1).item() # [-0.1,0.1] - x = x * contrast + brightness - # noise - noise = 0.05 * torch.randn_like(x) - x = x + noise - return torch.clamp(x, 0.0, 1.0) - - -# ========================= -# 3) Dataset -# ========================= -class LabeledDataset(Dataset): - def __init__(self, images_dir: Path, y_csv_path: Path): - self.image_paths = sorted(images_dir.glob("*.npy")) - self.names = [p.stem for p in self.image_paths] - self.y_df = pd.read_csv(y_csv_path, index_col=0) - - def __len__(self): - return len(self.image_paths) - - def __getitem__(self, idx: int): - name = self.names[idx] - img = np.load(self.image_paths[idx]) - raw_w = int(img.shape[1]) - - img = minmax_normalize(img) - img = pad_to_160x272(img, fill_value=0.0) - img_t = torch.from_numpy(img).unsqueeze(0).float() # (1,160,272) - img_t = resize_image_torch(img_t, MODEL_H, MODEL_W) # (1,224,224) - - row = self.y_df.loc[name].values.astype(np.int64) - mask = decode_mask_from_csv_row(row) # (160,w) - mask = pad_mask_to_160x272(mask) # (160,272) - mask_t = torch.from_numpy(mask).long() # (160,272) - mask_t = resize_mask_torch(mask_t, MODEL_H, MODEL_W) # (224,224) - - return {"name": name, "image": img_t, "mask": mask_t, "raw_w": raw_w} - - -class UnlabeledDataset(Dataset): - def __init__(self, images_dir: Path): - self.image_paths = sorted(images_dir.glob("*.npy")) - self.names = [p.stem for p in self.image_paths] - - def __len__(self): - return len(self.image_paths) - - def __getitem__(self, idx: int): - name = self.names[idx] - img = np.load(self.image_paths[idx]) - raw_w = int(img.shape[1]) - - img = minmax_normalize(img) - img = pad_to_160x272(img, fill_value=0.0) - img_t = torch.from_numpy(img).unsqueeze(0).float() # (1,160,272) - img_t = resize_image_torch(img_t, MODEL_H, MODEL_W) # (1,224,224) - - # 返回 base 图像（增强在 collate 做） - return {"name": name, "image": img_t, "raw_w": raw_w} - - -# ========================= -# 4) Collate -# ========================= -def collate_labeled(batch: List[Dict]) -> Dict: - names = [b["name"] for b in batch] - raw_ws = torch.tensor([b["raw_w"] for b in batch], dtype=torch.long) - - imgs_1 = torch.stack([b["image"] for b in batch], dim=0) # (B,1,224,224) - pixel_values = imgs_1.repeat(1, 3, 1, 1) # (B,3,224,224) - pixel_mask = torch.ones((pixel_values.shape[0], MODEL_H, MODEL_W), dtype=torch.long) - - class_labels_list, mask_labels_list = [], [] - for b in batch: - y = b["mask"] # (224,224) - cls, msk = semantic_to_mask2former_targets(y, NUM_CLASSES, IGNORE_INDEX) - class_labels_list.append(cls) - mask_labels_list.append(msk) - - return { - "names": names, - "raw_ws": raw_ws, - "pixel_values": pixel_values, - "pixel_mask": pixel_mask, - "class_labels": class_labels_list, - "mask_labels": mask_labels_list, - } - - -def collate_unlabeled(batch: List[Dict]) -> Dict: - names = [b["name"] for b in batch] - - imgs = [b["image"] for b in batch] # list of (1,224,224) - - # weak / strong augmentation - imgs_w = torch.stack([aug_weak(x.clone()) for x in imgs], dim=0) # (B,1,224,224) - imgs_s = torch.stack([aug_strong(x.clone()) for x in imgs], dim=0) # (B,1,224,224) - - pixel_values_w = imgs_w.repeat(1, 3, 1, 1) # (B,3,224,224) - pixel_values_s = imgs_s.repeat(1, 3, 1, 1) - - pixel_mask = torch.ones((pixel_values_w.shape[0], MODEL_H, MODEL_W), dtype=torch.long) - - return { - "names": names, - "pixel_values_w": pixel_values_w, - "pixel_values_s": pixel_values_s, - "pixel_mask": pixel_mask, - } - - -# ========================= -# 5) Model builder -# ========================= -def build_model(num_classes: int): - id2label = {0: "class0", 1: "class1", 2: "class2"} - label2id = {v: k for k, v in id2label.items()} - - processor = AutoImageProcessor.from_pretrained(PRETRAINED) - model = Mask2FormerForUniversalSegmentation.from_pretrained( - PRETRAINED, - ignore_mismatched_sizes=True, - id2label=id2label, - label2id=label2id, - num_labels=num_classes, - ) - return model, processor - - -# ========================= -# 6) Pseudo label from Mask2Former outputs -# ========================= -@torch.no_grad() -def pseudo_from_outputs(outputs, num_classes: int) -> Tuple[torch.Tensor, torch.Tensor]: - """ - 从 Mask2Former 输出构造像素级类别概率： - class_probs (softmax) * mask_probs (sigmoid) -> per-pixel scores - - outputs.class_queries_logits: (B, Q, C+1) (最后一个通常是 no-object) - outputs.masks_queries_logits: (B, Q, H, W) - - return: - pseudo: (B,H,W) long (0..C-1) - conf: (B,H,W) float (max score) - """ - class_logits = outputs.class_queries_logits # (B,Q,C+1) - mask_logits = outputs.masks_queries_logits # (B,Q,H,W) - - class_prob = class_logits.softmax(dim=-1)[..., :num_classes] # (B,Q,C) - mask_prob = mask_logits.sigmoid() # (B,Q,H,W) - - # (B,C,H,W) einsum: sum_q class_prob[b,q,c] * mask_prob[b,q,h,w] - score = torch.einsum("bqc,bqhw->bchw", class_prob, mask_prob) - conf, pseudo = torch.max(score, dim=1) # (B,H,W) - return pseudo.long(), conf.float() - - -# ========================= -# 7) Train / Validate (Semi-Supervised) -# ========================= -def train_one_epoch_semi(model, labeled_loader, unlabeled_loader, optimizer): - model.train() - - total_l, total_u = 0.0, 0.0 - n_l, n_u = 0, 0 - - unlabeled_iter = iter(unlabeled_loader) - - for batch_l in labeled_loader: - # ---- labeled step ---- - pixel_values = batch_l["pixel_values"].to(DEVICE) - pixel_mask = batch_l["pixel_mask"].to(DEVICE) - class_labels = [x.to(DEVICE) for x in batch_l["class_labels"]] - mask_labels = [x.to(DEVICE) for x in batch_l["mask_labels"]] - - out_l = model( - pixel_values=pixel_values, - pixel_mask=pixel_mask, - class_labels=class_labels, - mask_labels=mask_labels, - ) - loss_l = out_l.loss - - # ---- unlabeled step (pseudo-label) ---- - try: - batch_u = next(unlabeled_iter) - except StopIteration: - unlabeled_iter = iter(unlabeled_loader) - batch_u = next(unlabeled_iter) - - pv_w = batch_u["pixel_values_w"].to(DEVICE) # weak - pv_s = batch_u["pixel_values_s"].to(DEVICE) # strong - pm_u = batch_u["pixel_mask"].to(DEVICE) - - # teacher prediction on weak - model.eval() - out_u_teacher = model(pixel_values=pv_w, pixel_mask=pm_u) - pseudo, conf = pseudo_from_outputs(out_u_teacher, NUM_CLASSES) # (B,224,224) - - # 置信度过滤：低于阈值的像素设为 IGNORE - pseudo = pseudo.clone() - pseudo[conf < PSEUDO_TH] = IGNORE_INDEX - - # 将 pseudo semantic mask -> mask2former targets(list) - class_labels_u, mask_labels_u = [], [] - for i in range(pseudo.shape[0]): - cls_i, msk_i = semantic_to_mask2former_targets(pseudo[i], NUM_CLASSES, IGNORE_INDEX) - class_labels_u.append(cls_i.to(DEVICE)) - mask_labels_u.append(msk_i.to(DEVICE)) - - model.train() - out_u_student = model( - pixel_values=pv_s, - pixel_mask=pm_u, - class_labels=class_labels_u, - mask_labels=mask_labels_u, - ) - loss_u = out_u_student.loss - - # ---- total loss ---- - loss = loss_l + LAMBDA_U * loss_u - - optimizer.zero_grad() - loss.backward() - optimizer.step() - - total_l += float(loss_l.item()) * pixel_values.size(0) - total_u += float(loss_u.item()) * pv_s.size(0) - n_l += pixel_values.size(0) - n_u += pv_s.size(0) - - return total_l / max(1, n_l), total_u / max(1, n_u) - - -@torch.no_grad() -def valid_one_epoch(model, loader): - model.eval() - total = 0.0 - n = 0 - - for batch in loader: - pixel_values = batch["pixel_values"].to(DEVICE) - pixel_mask = batch["pixel_mask"].to(DEVICE) - class_labels = [x.to(DEVICE) for x in batch["class_labels"]] - mask_labels = [x.to(DEVICE) for x in batch["mask_labels"]] - - out = model( - pixel_values=pixel_values, - pixel_mask=pixel_mask, - class_labels=class_labels, - mask_labels=mask_labels, - ) - total += float(out.loss.item()) * pixel_values.size(0) - n += pixel_values.size(0) - - return total / max(1, n) - - -# ========================= -# 8) Inference & submission -# ========================= -@torch.no_grad() -def predict_and_make_submission(model, processor, test_images_dir: Path, out_csv_path: Path): - model.eval() - - # 这里复用 UnlabeledDataset 结构（只有图像，无标签） - test_ds = UnlabeledDataset(test_images_dir) - - def collate_test(batch: List[Dict]) -> Dict: - names = [b["name"] for b in batch] - raw_ws = torch.tensor([b["raw_w"] for b in batch], dtype=torch.long) - imgs_1 = torch.stack([b["image"] for b in batch], dim=0) # (B,1,224,224) - pixel_values = imgs_1.repeat(1, 3, 1, 1) - pixel_mask = torch.ones((pixel_values.shape[0], MODEL_H, MODEL_W), dtype=torch.long) - return {"names": names, "raw_ws": raw_ws, "pixel_values": pixel_values, "pixel_mask": pixel_mask} - - test_loader = DataLoader(test_ds, batch_size=1, shuffle=False, num_workers=0, collate_fn=collate_test) - - preds_dict = {} - - for batch in test_loader: - name = batch["names"][0] - raw_w = int(batch["raw_ws"][0].item()) - - pixel_values = batch["pixel_values"].to(DEVICE) - pixel_mask = batch["pixel_mask"].to(DEVICE) - - outputs = model(pixel_values=pixel_values, pixel_mask=pixel_mask) - - # 使用 processor 的语义后处理（稳定） - seg_list = processor.post_process_semantic_segmentation(outputs, target_sizes=[(MODEL_H, MODEL_W)]) - seg_224 = seg_list[0].to(torch.int64) # (224,224) - - seg_224 = seg_224.unsqueeze(0).unsqueeze(0).float() - seg_160_272 = F.interpolate(seg_224, size=(TARGET_H, TARGET_W), mode="nearest").squeeze(0).squeeze(0) - seg_160_272 = seg_160_272.cpu().numpy().astype(np.int64) - - pred = seg_160_272[:, :raw_w] - - padded = np.full((TARGET_H * TARGET_W,), IGNORE_INDEX, dtype=np.int64) - padded[: TARGET_H * raw_w] = pred.flatten() - preds_dict[name] = padded - - sub = pd.DataFrame(preds_dict, dtype="int64").T - sub.to_csv(out_csv_path) - print(f"[OK] submission saved to: {out_csv_path}") - - -# ========================= -# 9) Main -# ========================= -def main(): - print(f"DEVICE: {DEVICE}") - print(f"Labeled train dir: {TRAIN_IMAGES_DIR}") - print(f"Unlabeled dir: {UNLABELED_DIR}") - print(f"Test dir: {TEST_IMAGES_DIR}") - print(f"Pretrained: {PRETRAINED}") - print(f"Pseudo TH={PSEUDO_TH}, lambda_u={LAMBDA_U}") - - # ---- labeled dataset (well1-6) ---- - ds_all = LabeledDataset(TRAIN_IMAGES_DIR, Y_TRAIN_CSV) - - # split by well (val=6) - train_idx, val_idx = [], [] - for i, name in enumerate(ds_all.names): - w = parse_well_id(name) - if w == 6: - val_idx.append(i) - else: - train_idx.append(i) - - train_ds = Subset(ds_all, train_idx) # well1-5 - val_ds = Subset(ds_all, val_idx) # well6 - - train_loader = DataLoader( - train_ds, - batch_size=BATCH_SIZE_L, - shuffle=True, - num_workers=0, - collate_fn=collate_labeled, - ) - val_loader = DataLoader( - val_ds, - batch_size=BATCH_SIZE_L, - shuffle=False, - num_workers=0, - collate_fn=collate_labeled, - ) - - # ---- unlabeled dataset (well12-14) ---- - unlab_ds = UnlabeledDataset(UNLABELED_DIR) - unlab_loader = DataLoader( - unlab_ds, - batch_size=BATCH_SIZE_U, - shuffle=True, - num_workers=0, - collate_fn=collate_unlabeled, - ) - - print(f"Labeled train: {len(train_ds)} | Val: {len(val_ds)} | Unlabeled: {len(unlab_ds)}") - - model, processor = build_model(NUM_CLASSES) - model = model.to(DEVICE) - - optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY) - - best_val = 1e9 - best_path = DATA_ROOT / "best_mask2former_semi.pth" - - for epoch in range(1, EPOCHS + 1): - tr_l, tr_u = train_one_epoch_semi(model, train_loader, unlab_loader, optimizer) - va = valid_one_epoch(model, val_loader) - - print(f"Epoch {epoch:02d}/{EPOCHS} | train_l={tr_l:.4f} | train_u={tr_u:.4f} | val={va:.4f}") - - if va < best_val: - best_val = va - torch.save(model.state_dict(), best_path) - print(f" -> Best saved: {best_path}") - - # ---- inference ---- - out_csv = DATA_ROOT / "submission.csv" - model.load_state_dict(torch.load(best_path, map_location=DEVICE, weights_only=True)) - predict_and_make_submission(model, processor, TEST_IMAGES_DIR, out_csv) - - -if __name__ == "__main__": - main() From 17838433ad8fe0090e1448c5fe124fe07cd6bc59 Mon Sep 17 00:00:00 2001 From: Haokun Date: Sun, 11 Jan 2026 15:55:24 +0100 Subject: [PATCH 09/11] Create Mask2Former_Semi_Supervised2.ipynb --- Mask2Former_Semi_Supervised2.ipynb | 0 1 file changed, 0 insertions(+), 0 deletions(-) create mode 100644 Mask2Former_Semi_Supervised2.ipynb diff --git a/Mask2Former_Semi_Supervised2.ipynb b/Mask2Former_Semi_Supervised2.ipynb new file mode 100644 index 0000000..e69de29 From 70b4b3ebddc2aeecad53c65e8e12e7a9a339d06d Mon Sep 17 00:00:00 2001 From: Haokun Date: Sun, 11 Jan 2026 16:21:50 +0100 Subject: [PATCH 10/11] explanation --- DPT.ipynb | 3 +- Mask2Former_Semi_Supervised2.ipynb | 780 +++++++++++++++++++++ UPerNet_Swin_ViT_backbone.ipynb | 3 +- attention_u_net.ipynb | 11 +- deeplabv3.ipynb | 11 +- "mask2_former_semi-supervis\303\251.ipynb" | 8 + mask2former.ipynb | 11 +- resnet.ipynb | 93 +-- swin-unet.ipynb | 11 +- upernet.ipynb | 162 ++--- 10 files changed, 948 insertions(+), 145 deletions(-) diff --git a/DPT.ipynb b/DPT.ipynb index 6adc67a..a85d708 100644 --- a/DPT.ipynb +++ b/DPT.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "id": "d12efdc1", "metadata": {}, "outputs": [ @@ -63,7 +63,6 @@ } ], "source": [ - "# -*- coding: utf-8 -*-\n", "\"\"\"\n", "DPT (Dense Prediction Transformer) - Full runnable code (NO cv2)\n", "Using segmentation-models-pytorch (smp) built-in DPT.\n", diff --git a/Mask2Former_Semi_Supervised2.ipynb b/Mask2Former_Semi_Supervised2.ipynb index e69de29..76572bc 100644 --- a/Mask2Former_Semi_Supervised2.ipynb +++ b/Mask2Former_Semi_Supervised2.ipynb @@ -0,0 +1,780 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": 1, + "id": "8fca60f8", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "c:\\Users\\lenovo\\anaconda3\\envs\\deep-torch\\lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n", + " from .autonotebook import tqdm as notebook_tqdm\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "DEVICE: cuda | AMP: True\n", + "Labeled train dir: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\X_train_uDRk9z9\\images\n", + "Unlabeled dir: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\X_unlabeled_mtkxUlo\\images\n", + "Test dir: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\X_test_xNbnvIa\\images\n", + "Pretrained: facebook/mask2former-swin-tiny-ade-semantic\n", + "EPOCHS=5, pseudo_update_every=2, unlab_ratio=0.5\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Using a slow image processor as `use_fast` is unset and a slow processor was saved with this model. `use_fast=True` will be the default behavior in v4.52, even if the model was saved with a slow processor. This will result in minor differences in outputs. You'll still be able to use a slow processor with `use_fast=False`.\n", + "c:\\Users\\lenovo\\anaconda3\\envs\\deep-torch\\lib\\site-packages\\transformers\\image_processing_base.py:417: UserWarning: The following named arguments are not valid for `Mask2FormerImageProcessor.__init__` and were ignored: '_max_size', 'reduce_labels'\n", + " image_processor = cls(**image_processor_dict)\n", + "Some weights of Mask2FormerForUniversalSegmentation were not initialized from the model checkpoint at facebook/mask2former-swin-tiny-ade-semantic and are newly initialized because the shapes did not match:\n", + "- class_predictor.bias: found shape torch.Size([151]) in the checkpoint and torch.Size([4]) in the model instantiated\n", + "- class_predictor.weight: found shape torch.Size([151, 256]) in the checkpoint and torch.Size([4, 256]) in the model instantiated\n", + "- criterion.empty_weight: found shape torch.Size([151]) in the checkpoint and torch.Size([4]) in the model instantiated\n", + "You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference.\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Labeled train: 2790 | Val: 1620 | Unlabeled used: 990\n", + "\n", + "Epoch 01/5 | lambda_u=0.050 | tau=0.95\n", + "[Info] Building pseudo-label cache...\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "C:\\Users\\lenovo\\AppData\\Local\\Temp\\ipykernel_24600\\1117854030.py:375: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n", + " with torch.cuda.amp.autocast(enabled=AMP):\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[Info] Pseudo cache size: 990\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "C:\\Users\\lenovo\\AppData\\Local\\Temp\\ipykernel_24600\\1117854030.py:408: FutureWarning: `torch.cuda.amp.GradScaler(args...)` is deprecated. Please use `torch.amp.GradScaler('cuda', args...)` instead.\n", + " scaler = torch.cuda.amp.GradScaler(enabled=AMP)\n", + "C:\\Users\\lenovo\\AppData\\Local\\Temp\\ipykernel_24600\\1117854030.py:441: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n", + " with torch.cuda.amp.autocast(enabled=AMP):\n", + "C:\\Users\\lenovo\\AppData\\Local\\Temp\\ipykernel_24600\\1117854030.py:483: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n", + " with torch.cuda.amp.autocast(enabled=AMP):\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[Loss] train=20.0172 | val(well6)=14.3855\n", + "[OK] Best model saved: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_mask2former_semi.pth\n", + "\n", + "Epoch 02/5 | lambda_u=0.163 | tau=0.92\n", + "[Loss] train=18.9800 | val(well6)=14.1031\n", + "[OK] Best model saved: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\best_mask2former_semi.pth\n", + "\n", + "Epoch 03/5 | lambda_u=0.275 | tau=0.90\n", + "[Info] Building pseudo-label cache...\n", + "[Info] Pseudo cache size: 990\n", + "[Loss] train=19.6246 | val(well6)=15.1079\n", + "\n", + "Epoch 04/5 | lambda_u=0.388 | tau=0.88\n", + "[Loss] train=20.9259 | val(well6)=14.5028\n", + "\n", + "Epoch 05/5 | lambda_u=0.500 | tau=0.85\n", + "[Info] Building pseudo-label cache...\n", + "[Info] Pseudo cache size: 990\n", + "[Loss] train=21.1578 | val(well6)=14.7531\n", + "\n", + "[Info] Loading best model and generating submission...\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "C:\\Users\\lenovo\\AppData\\Local\\Temp\\ipykernel_24600\\1117854030.py:515: FutureWarning: `torch.cuda.amp.autocast(args...)` is deprecated. Please use `torch.amp.autocast('cuda', args...)` instead.\n", + " with torch.cuda.amp.autocast(enabled=AMP):\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[OK] submission saved to: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\submission.csv\n" + ] + } + ], + "source": [ + "\"\"\"\n", + "Mask2Former + Semi-Supervised Pseudo-Label (Full runnable, 5-epoch quick run)\n", + "\n", + "Data:\n", + "- Labeled train images: X_train_uDRk9z9/images (well1-6)\n", + "- Unlabeled images: X_unlabeled_mtkxUlo/images (well12-14)\n", + "- Test images: X_test_xNbnvIa/images (well7-11)\n", + "- Train labels (CSV): Y_train_T9NrBYo.csv (flattened + -1 padding)\n", + "\n", + "Split (avoid leakage by well):\n", + "- Train: well1-5\n", + "- Val: well6\n", + "- Test: well7-11 (submission only)\n", + "\n", + "Output:\n", + "- submission.csv (one row per patch name, flattened, padded to 160*272 with -1)\n", + "\n", + "Notes:\n", + "- Model input uses 224x224 (Mask2Former pretrained backbone-friendly).\n", + "- Images are single-channel; we repeat channel to 3.\n", + "- Semi-supervised from epoch 1: supervised + lambda_u * unsup_loss\n", + "- Pseudo labels are filtered by pixel confidence threshold tau.\n", + "\n", + "Install (in your env):\n", + " pip install -U transformers accelerate\n", + " # if transformers complains about huggingface-hub:\n", + " # pip install \"huggingface-hub<1.0,>=0.34.0\"\n", + "\"\"\"\n", + "\n", + "import re\n", + "from pathlib import Path\n", + "from typing import Dict, List, Tuple\n", + "\n", + "import numpy as np\n", + "import pandas as pd\n", + "\n", + "import torch\n", + "import torch.nn.functional as F\n", + "from torch.utils.data import Dataset, DataLoader, Subset\n", + "\n", + "from transformers import AutoImageProcessor, Mask2FormerForUniversalSegmentation\n", + "\n", + "\n", + "# =========================\n", + "# 0) Paths\n", + "# =========================\n", + "DATA_ROOT = Path(r\"C:\\Users\\lenovo\\Desktop\\deep_datachallenge\")\n", + "\n", + "TRAIN_IMAGES_DIR = DATA_ROOT / \"X_train_uDRk9z9\" / \"images\"\n", + "UNLAB_IMAGES_DIR = DATA_ROOT / \"X_unlabeled_mtkxUlo\" / \"images\"\n", + "TEST_IMAGES_DIR = DATA_ROOT / \"X_test_xNbnvIa\" / \"images\"\n", + "Y_TRAIN_CSV = DATA_ROOT / \"Y_train_T9NrBYo.csv\"\n", + "\n", + "# Submission target size\n", + "TARGET_H, TARGET_W = 160, 272\n", + "IGNORE_INDEX = -1\n", + "NUM_CLASSES = 3 # classes are 0/1/2\n", + "\n", + "# Model input size\n", + "MODEL_H, MODEL_W = 224, 224\n", + "\n", + "# Pretrained checkpoint (semantic)\n", + "PRETRAINED = \"facebook/mask2former-swin-tiny-ade-semantic\"\n", + "\n", + "# =========================\n", + "# 1) Hyperparameters (5-epoch quick run)\n", + "# =========================\n", + "EPOCHS = 5\n", + "BATCH_SIZE_L = 2 # labeled batch size\n", + "BATCH_SIZE_U = 2 # unlabeled batch size (keep small)\n", + "LR = 5e-5\n", + "WEIGHT_DECAY = 1e-4\n", + "\n", + "# Semi-supervised schedule (start from epoch 1)\n", + "LAMBDA_U_MIN = 0.05\n", + "LAMBDA_U_MAX = 0.50\n", + "RAMP_EPOCHS = 5\n", + "\n", + "# Pseudo label threshold schedule\n", + "PSEUDO_TH_START = 0.95\n", + "PSEUDO_TH_END = 0.85\n", + "\n", + "# Cache pseudo labels and refresh every N epochs\n", + "PSEUDO_UPDATE_EVERY = 2\n", + "\n", + "# Unlabeled sampling ratio (0~1): use a subset of unlabeled per epoch for speed\n", + "UNLAB_SAMPLE_RATIO = 0.50\n", + "\n", + "# Speed\n", + "NUM_WORKERS = 0 # set to 2~4 if your Windows setup is stable\n", + "\n", + "\n", + "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", + "AMP = (DEVICE == \"cuda\")\n", + "\n", + "\n", + "# =========================\n", + "# 2) Utils\n", + "# =========================\n", + "def parse_well_id(name: str) -> int:\n", + " m = re.search(r\"well_(\\d+)_\", name)\n", + " return int(m.group(1)) if m else -1\n", + "\n", + "\n", + "def minmax_normalize(x: np.ndarray) -> np.ndarray:\n", + " x = x.astype(np.float32)\n", + " x = np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)\n", + " x_min = float(x.min())\n", + " x_max = float(x.max())\n", + " if x_max - x_min < 1e-6:\n", + " return np.zeros_like(x, dtype=np.float32)\n", + " return (x - x_min) / (x_max - x_min)\n", + "\n", + "\n", + "def pad_to_160x272(img: np.ndarray, fill_value: float = 0.0) -> np.ndarray:\n", + " h, w = img.shape\n", + " assert h == TARGET_H, f\"Expected height {TARGET_H}, got {h}\"\n", + " if w == TARGET_W:\n", + " return img\n", + " if w < TARGET_W:\n", + " out = np.full((TARGET_H, TARGET_W), fill_value, dtype=img.dtype)\n", + " out[:, :w] = img\n", + " return out\n", + " return img[:, :TARGET_W]\n", + "\n", + "\n", + "def decode_mask_from_csv_row(row_values: np.ndarray) -> np.ndarray:\n", + " valid = row_values[row_values != IGNORE_INDEX]\n", + " assert len(valid) % TARGET_H == 0, f\"Valid mask length {len(valid)} not divisible by {TARGET_H}\"\n", + " w = len(valid) // TARGET_H\n", + " return valid.reshape(TARGET_H, w).astype(np.int64)\n", + "\n", + "\n", + "def pad_mask_to_160x272(mask: np.ndarray) -> np.ndarray:\n", + " h, w = mask.shape\n", + " assert h == TARGET_H\n", + " if w == TARGET_W:\n", + " return mask\n", + " out = np.full((TARGET_H, TARGET_W), IGNORE_INDEX, dtype=np.int64)\n", + " out[:, :w] = mask\n", + " return out\n", + "\n", + "\n", + "def resize_image_torch(img_1hw: torch.Tensor, h: int, w: int) -> torch.Tensor:\n", + " x = img_1hw.unsqueeze(0) # (1,1,H,W)\n", + " x = F.interpolate(x, size=(h, w), mode=\"bilinear\", align_corners=False)\n", + " return x.squeeze(0) # (1,h,w)\n", + "\n", + "\n", + "def resize_mask_torch(mask_hw: torch.Tensor, h: int, w: int) -> torch.Tensor:\n", + " y = mask_hw.unsqueeze(0).unsqueeze(0).float()\n", + " y = F.interpolate(y, size=(h, w), mode=\"nearest\")\n", + " return y.squeeze(0).squeeze(0).long()\n", + "\n", + "\n", + "def semantic_to_mask2former_targets(\n", + " semantic_mask: torch.Tensor,\n", + " num_classes: int,\n", + " ignore_index: int = -1,\n", + ") -> Tuple[torch.Tensor, torch.Tensor]:\n", + " \"\"\"\n", + " Convert semantic (H,W) to set-based:\n", + " - class_labels: (N,)\n", + " - mask_labels: (N,H,W) float 0/1\n", + " Ignore pixels == ignore_index.\n", + " \"\"\"\n", + " valid = semantic_mask != ignore_index\n", + " if valid.sum() == 0:\n", + " class_labels = torch.tensor([0], dtype=torch.long)\n", + " mask_labels = torch.zeros((1, semantic_mask.shape[0], semantic_mask.shape[1]), dtype=torch.float32)\n", + " return class_labels, mask_labels\n", + "\n", + " present = torch.unique(semantic_mask[valid]).tolist()\n", + " present = [int(c) for c in present if 0 <= int(c) < num_classes]\n", + " if len(present) == 0:\n", + " class_labels = torch.tensor([0], dtype=torch.long)\n", + " mask_labels = torch.zeros((1, semantic_mask.shape[0], semantic_mask.shape[1]), dtype=torch.float32)\n", + " return class_labels, mask_labels\n", + "\n", + " masks, classes = [], []\n", + " for c in present:\n", + " m = (semantic_mask == c) & valid\n", + " if m.sum() == 0:\n", + " continue\n", + " masks.append(m.float())\n", + " classes.append(c)\n", + "\n", + " if len(classes) == 0:\n", + " class_labels = torch.tensor([0], dtype=torch.long)\n", + " mask_labels = torch.zeros((1, semantic_mask.shape[0], semantic_mask.shape[1]), dtype=torch.float32)\n", + " return class_labels, mask_labels\n", + "\n", + " class_labels = torch.tensor(classes, dtype=torch.long)\n", + " mask_labels = torch.stack(masks, dim=0).float()\n", + " return class_labels, mask_labels\n", + "\n", + "\n", + "def get_lambda_u(epoch: int) -> float:\n", + " if RAMP_EPOCHS <= 1:\n", + " return LAMBDA_U_MAX\n", + " t = (epoch - 1) / (RAMP_EPOCHS - 1)\n", + " t = max(0.0, min(1.0, t))\n", + " return LAMBDA_U_MIN + (LAMBDA_U_MAX - LAMBDA_U_MIN) * t\n", + "\n", + "\n", + "def get_pseudo_th(epoch: int) -> float:\n", + " if EPOCHS <= 1:\n", + " return PSEUDO_TH_END\n", + " t = (epoch - 1) / (EPOCHS - 1)\n", + " t = max(0.0, min(1.0, t))\n", + " return PSEUDO_TH_START + (PSEUDO_TH_END - PSEUDO_TH_START) * t\n", + "\n", + "\n", + "# =========================\n", + "# 3) Datasets\n", + "# =========================\n", + "class LabeledWellDataset(Dataset):\n", + " def __init__(self, images_dir: Path, y_csv_path: Path):\n", + " self.image_paths = sorted(images_dir.glob(\"*.npy\"))\n", + " self.names = [p.stem for p in self.image_paths]\n", + " self.y_df = pd.read_csv(y_csv_path, index_col=0)\n", + "\n", + " def __len__(self):\n", + " return len(self.image_paths)\n", + "\n", + " def __getitem__(self, idx: int) -> Dict:\n", + " name = self.names[idx]\n", + " img = np.load(self.image_paths[idx]) # (160,160) or (160,272)\n", + " raw_w = int(img.shape[1])\n", + "\n", + " img = minmax_normalize(img)\n", + " img = pad_to_160x272(img, fill_value=0.0)\n", + "\n", + " img_t = torch.from_numpy(img).unsqueeze(0).float() # (1,160,272)\n", + " img_t = resize_image_torch(img_t, MODEL_H, MODEL_W) # (1,224,224)\n", + "\n", + " row = self.y_df.loc[name].values.astype(np.int64)\n", + " mask = decode_mask_from_csv_row(row) # (160,w)\n", + " mask = pad_mask_to_160x272(mask) # (160,272)\n", + " mask_t = torch.from_numpy(mask).long() # (160,272)\n", + " mask_t = resize_mask_torch(mask_t, MODEL_H, MODEL_W) # (224,224)\n", + "\n", + " return {\"name\": name, \"image\": img_t, \"mask\": mask_t, \"raw_w\": raw_w}\n", + "\n", + "\n", + "class UnlabeledWellDataset(Dataset):\n", + " def __init__(self, images_dir: Path):\n", + " self.image_paths = sorted(images_dir.glob(\"*.npy\"))\n", + " self.names = [p.stem for p in self.image_paths]\n", + "\n", + " def __len__(self):\n", + " return len(self.image_paths)\n", + "\n", + " def __getitem__(self, idx: int) -> Dict:\n", + " name = self.names[idx]\n", + " img = np.load(self.image_paths[idx])\n", + " raw_w = int(img.shape[1])\n", + "\n", + " img = minmax_normalize(img)\n", + " img = pad_to_160x272(img, fill_value=0.0)\n", + "\n", + " img_t = torch.from_numpy(img).unsqueeze(0).float() # (1,160,272)\n", + " img_t = resize_image_torch(img_t, MODEL_H, MODEL_W) # (1,224,224)\n", + "\n", + " return {\"name\": name, \"image\": img_t, \"raw_w\": raw_w, \"idx\": idx}\n", + "\n", + "\n", + "# =========================\n", + "# 4) Collate\n", + "# =========================\n", + "def collate_labeled(batch: List[Dict]) -> Dict:\n", + " names = [b[\"name\"] for b in batch]\n", + " raw_ws = torch.tensor([b[\"raw_w\"] for b in batch], dtype=torch.long)\n", + "\n", + " imgs_1 = torch.stack([b[\"image\"] for b in batch], dim=0) # (B,1,224,224)\n", + " pixel_values = imgs_1.repeat(1, 3, 1, 1) # (B,3,224,224)\n", + " pixel_mask = torch.ones((pixel_values.shape[0], MODEL_H, MODEL_W), dtype=torch.long)\n", + "\n", + " class_labels_list, mask_labels_list = [], []\n", + " for b in batch:\n", + " cls, msk = semantic_to_mask2former_targets(b[\"mask\"], NUM_CLASSES, IGNORE_INDEX)\n", + " class_labels_list.append(cls)\n", + " mask_labels_list.append(msk)\n", + "\n", + " return {\n", + " \"names\": names,\n", + " \"raw_ws\": raw_ws,\n", + " \"pixel_values\": pixel_values,\n", + " \"pixel_mask\": pixel_mask,\n", + " \"class_labels\": class_labels_list,\n", + " \"mask_labels\": mask_labels_list,\n", + " }\n", + "\n", + "\n", + "def collate_unlabeled(batch: List[Dict]) -> Dict:\n", + " names = [b[\"name\"] for b in batch]\n", + " raw_ws = torch.tensor([b[\"raw_w\"] for b in batch], dtype=torch.long)\n", + " idxs = torch.tensor([b[\"idx\"] for b in batch], dtype=torch.long)\n", + "\n", + " imgs_1 = torch.stack([b[\"image\"] for b in batch], dim=0)\n", + " pixel_values = imgs_1.repeat(1, 3, 1, 1)\n", + " pixel_mask = torch.ones((pixel_values.shape[0], MODEL_H, MODEL_W), dtype=torch.long)\n", + "\n", + " return {\n", + " \"names\": names,\n", + " \"raw_ws\": raw_ws,\n", + " \"idxs\": idxs,\n", + " \"pixel_values\": pixel_values,\n", + " \"pixel_mask\": pixel_mask,\n", + " }\n", + "\n", + "\n", + "# =========================\n", + "# 5) Build model & processor\n", + "# =========================\n", + "def build_model_and_processor() -> Tuple[Mask2FormerForUniversalSegmentation, AutoImageProcessor]:\n", + " id2label = {0: \"class0\", 1: \"class1\", 2: \"class2\"}\n", + " label2id = {v: k for k, v in id2label.items()}\n", + "\n", + " processor = AutoImageProcessor.from_pretrained(PRETRAINED)\n", + "\n", + " model = Mask2FormerForUniversalSegmentation.from_pretrained(\n", + " PRETRAINED,\n", + " ignore_mismatched_sizes=True,\n", + " id2label=id2label,\n", + " label2id=label2id,\n", + " num_labels=NUM_CLASSES,\n", + " use_safetensors=True, # safer + avoids torch.load restrictions\n", + " )\n", + " return model, processor\n", + "\n", + "\n", + "# =========================\n", + "# 6) Pseudo label from Mask2Former outputs\n", + "# =========================\n", + "@torch.no_grad()\n", + "def outputs_to_semantic_and_conf(outputs) -> Tuple[torch.Tensor, torch.Tensor]:\n", + " \"\"\"\n", + " From Mask2Former outputs -> semantic label map + confidence map (both HxW, H=W=224).\n", + " Using:\n", + " class_probs: softmax over (num_classes + no-object), drop last\n", + " mask_probs: sigmoid\n", + " per-pixel per-class score = sum_q class_probs[q,c] * mask_probs[q, y, x]\n", + " \"\"\"\n", + " class_logits = outputs.class_queries_logits[0] # (Q, C+1)\n", + " mask_logits = outputs.masks_queries_logits[0] # (Q, H, W)\n", + "\n", + " class_probs = F.softmax(class_logits, dim=-1)[..., :NUM_CLASSES] # (Q, C)\n", + " mask_probs = torch.sigmoid(mask_logits) # (Q, H, W)\n", + "\n", + " # score: (C, H, W)\n", + " score = torch.einsum(\"qc,qhw->chw\", class_probs, mask_probs)\n", + "\n", + " conf, pred = torch.max(score, dim=0) # (H,W), (H,W)\n", + " return pred.to(torch.int64), conf.to(torch.float32)\n", + "\n", + "\n", + "@torch.no_grad()\n", + "def build_pseudo_cache(\n", + " model: torch.nn.Module,\n", + " unlab_loader: DataLoader,\n", + " tau: float,\n", + ") -> Dict[int, torch.Tensor]:\n", + " \"\"\"\n", + " Create a cache: {unlabeled_idx -> pseudo_mask(H,W) int64}, with low-confidence pixels set to IGNORE_INDEX.\n", + " \"\"\"\n", + " model.eval()\n", + " cache: Dict[int, torch.Tensor] = {}\n", + "\n", + " for batch in unlab_loader:\n", + " pixel_values = batch[\"pixel_values\"].to(DEVICE)\n", + " pixel_mask = batch[\"pixel_mask\"].to(DEVICE)\n", + " idxs = batch[\"idxs\"].tolist()\n", + "\n", + " with torch.cuda.amp.autocast(enabled=AMP):\n", + " outputs = model(pixel_values=pixel_values, pixel_mask=pixel_mask)\n", + "\n", + " # process each sample in batch\n", + " for b_i, u_idx in enumerate(idxs):\n", + " # take sample b_i outputs by slicing\n", + " out_i = type(outputs)(\n", + " **{k: (v[b_i:b_i+1] if torch.is_tensor(v) else v) for k, v in outputs.items()}\n", + " )\n", + "\n", + " pred, conf = outputs_to_semantic_and_conf(out_i)\n", + " pseudo = pred.clone()\n", + " pseudo[conf < tau] = IGNORE_INDEX\n", + " cache[u_idx] = pseudo.cpu()\n", + "\n", + " return cache\n", + "\n", + "\n", + "# =========================\n", + "# 7) Train / Eval\n", + "# =========================\n", + "def train_one_epoch(\n", + " model: torch.nn.Module,\n", + " labeled_loader: DataLoader,\n", + " unlabeled_loader: DataLoader,\n", + " pseudo_cache: Dict[int, torch.Tensor],\n", + " optimizer: torch.optim.Optimizer,\n", + " lambda_u: float,\n", + ") -> float:\n", + " model.train()\n", + " total_loss = 0.0\n", + " n_seen = 0\n", + "\n", + " scaler = torch.cuda.amp.GradScaler(enabled=AMP)\n", + "\n", + " # iterate by min length (simple)\n", + " it_u = iter(unlabeled_loader)\n", + " for batch_l in labeled_loader:\n", + " try:\n", + " batch_u = next(it_u)\n", + " except StopIteration:\n", + " it_u = iter(unlabeled_loader)\n", + " batch_u = next(it_u)\n", + "\n", + " # ----- labeled -----\n", + " pv_l = batch_l[\"pixel_values\"].to(DEVICE)\n", + " pm_l = batch_l[\"pixel_mask\"].to(DEVICE)\n", + " cls_l = [x.to(DEVICE) for x in batch_l[\"class_labels\"]]\n", + " msk_l = [x.to(DEVICE) for x in batch_l[\"mask_labels\"]]\n", + "\n", + " # ----- unlabeled -----\n", + " pv_u = batch_u[\"pixel_values\"].to(DEVICE)\n", + " pm_u = batch_u[\"pixel_mask\"].to(DEVICE)\n", + " idxs_u = batch_u[\"idxs\"].tolist()\n", + "\n", + " # build pseudo targets for this unlabeled batch from cache\n", + " class_labels_u, mask_labels_u = [], []\n", + " for u_idx in idxs_u:\n", + " pseudo = pseudo_cache[u_idx] # (224,224) on CPU\n", + " pseudo_t = pseudo.to(torch.int64)\n", + " cls, msk = semantic_to_mask2former_targets(pseudo_t, NUM_CLASSES, IGNORE_INDEX)\n", + " class_labels_u.append(cls.to(DEVICE))\n", + " mask_labels_u.append(msk.to(DEVICE))\n", + "\n", + " optimizer.zero_grad(set_to_none=True)\n", + "\n", + " with torch.cuda.amp.autocast(enabled=AMP):\n", + " out_l = model(\n", + " pixel_values=pv_l,\n", + " pixel_mask=pm_l,\n", + " class_labels=cls_l,\n", + " mask_labels=msk_l,\n", + " )\n", + " loss_sup = out_l.loss\n", + "\n", + " out_u = model(\n", + " pixel_values=pv_u,\n", + " pixel_mask=pm_u,\n", + " class_labels=class_labels_u,\n", + " mask_labels=mask_labels_u,\n", + " )\n", + " loss_unsup = out_u.loss\n", + "\n", + " loss = loss_sup + lambda_u * loss_unsup\n", + "\n", + " scaler.scale(loss).backward()\n", + " scaler.step(optimizer)\n", + " scaler.update()\n", + "\n", + " bs = pv_l.size(0)\n", + " total_loss += float(loss.item()) * bs\n", + " n_seen += bs\n", + "\n", + " return total_loss / max(1, n_seen)\n", + "\n", + "\n", + "@torch.no_grad()\n", + "def eval_one_epoch(model: torch.nn.Module, val_loader: DataLoader) -> float:\n", + " model.eval()\n", + " total_loss = 0.0\n", + " n_seen = 0\n", + "\n", + " for batch in val_loader:\n", + " pv = batch[\"pixel_values\"].to(DEVICE)\n", + " pm = batch[\"pixel_mask\"].to(DEVICE)\n", + " cls = [x.to(DEVICE) for x in batch[\"class_labels\"]]\n", + " msk = [x.to(DEVICE) for x in batch[\"mask_labels\"]]\n", + "\n", + " with torch.cuda.amp.autocast(enabled=AMP):\n", + " out = model(pixel_values=pv, pixel_mask=pm, class_labels=cls, mask_labels=msk)\n", + " loss = out.loss\n", + "\n", + " bs = pv.size(0)\n", + " total_loss += float(loss.item()) * bs\n", + " n_seen += bs\n", + "\n", + " return total_loss / max(1, n_seen)\n", + "\n", + "\n", + "# =========================\n", + "# 8) Predict & submission\n", + "# =========================\n", + "@torch.no_grad()\n", + "def predict_and_submit(model: torch.nn.Module, processor: AutoImageProcessor, out_csv_path: Path):\n", + " model.eval()\n", + "\n", + " test_ds = UnlabeledWellDataset(TEST_IMAGES_DIR) # same structure, no labels\n", + " test_loader = DataLoader(\n", + " test_ds, batch_size=1, shuffle=False, num_workers=NUM_WORKERS, collate_fn=collate_unlabeled\n", + " )\n", + "\n", + " preds_dict = {}\n", + "\n", + " for batch in test_loader:\n", + " name = batch[\"names\"][0]\n", + " raw_w = int(batch[\"raw_ws\"][0].item())\n", + "\n", + " pv = batch[\"pixel_values\"].to(DEVICE)\n", + " pm = batch[\"pixel_mask\"].to(DEVICE)\n", + "\n", + " with torch.cuda.amp.autocast(enabled=AMP):\n", + " outputs = model(pixel_values=pv, pixel_mask=pm)\n", + "\n", + " # semantic map at 224x224\n", + " pred_224, _ = outputs_to_semantic_and_conf(outputs) # (224,224)\n", + "\n", + " # upsample back to 160x272\n", + " pred_224_f = pred_224.unsqueeze(0).unsqueeze(0).float()\n", + " pred_160_272 = F.interpolate(pred_224_f, size=(TARGET_H, TARGET_W), mode=\"nearest\").squeeze(0).squeeze(0)\n", + " pred_160_272 = pred_160_272.cpu().numpy().astype(np.int64)\n", + "\n", + " # crop to original width\n", + " pred = pred_160_272[:, :raw_w]\n", + "\n", + " # flatten + pad -1 to 160*272\n", + " if raw_w < TARGET_W:\n", + " padded = np.full((TARGET_H * TARGET_W,), IGNORE_INDEX, dtype=np.int64)\n", + " padded[: TARGET_H * raw_w] = pred.flatten()\n", + " preds_dict[name] = padded\n", + " else:\n", + " preds_dict[name] = pred.flatten()\n", + "\n", + " sub = pd.DataFrame(preds_dict, dtype=\"int64\").T\n", + " sub.to_csv(out_csv_path)\n", + " print(f\"[OK] submission saved to: {out_csv_path}\")\n", + "\n", + "\n", + "# =========================\n", + "# 9) Main\n", + "# =========================\n", + "def main():\n", + " print(f\"DEVICE: {DEVICE} | AMP: {AMP}\")\n", + " print(f\"Labeled train dir: {TRAIN_IMAGES_DIR}\")\n", + " print(f\"Unlabeled dir: {UNLAB_IMAGES_DIR}\")\n", + " print(f\"Test dir: {TEST_IMAGES_DIR}\")\n", + " print(f\"Pretrained: {PRETRAINED}\")\n", + " print(f\"EPOCHS={EPOCHS}, pseudo_update_every={PSEUDO_UPDATE_EVERY}, unlab_ratio={UNLAB_SAMPLE_RATIO}\")\n", + "\n", + " # ----- labeled dataset (well1-6) -----\n", + " labeled_all = LabeledWellDataset(TRAIN_IMAGES_DIR, Y_TRAIN_CSV)\n", + "\n", + " # split: well6 as val, well1-5 as train\n", + " VAL_WELLS = {6}\n", + " tr_idx, va_idx = [], []\n", + " for i, name in enumerate(labeled_all.names):\n", + " w = parse_well_id(name)\n", + " if w in VAL_WELLS:\n", + " va_idx.append(i)\n", + " else:\n", + " tr_idx.append(i)\n", + "\n", + " train_ds = Subset(labeled_all, tr_idx)\n", + " val_ds = Subset(labeled_all, va_idx)\n", + "\n", + " train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE_L, shuffle=True,\n", + " num_workers=NUM_WORKERS, collate_fn=collate_labeled, drop_last=True)\n", + " val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE_L, shuffle=False,\n", + " num_workers=NUM_WORKERS, collate_fn=collate_labeled)\n", + "\n", + " # ----- unlabeled dataset (well12-14) -----\n", + " unlab_all = UnlabeledWellDataset(UNLAB_IMAGES_DIR)\n", + " n_unlab = len(unlab_all)\n", + " n_use = max(1, int(n_unlab * UNLAB_SAMPLE_RATIO))\n", + "\n", + " # sample a fixed subset for this run (fast & reproducible)\n", + " rng = np.random.RandomState(42)\n", + " use_indices = rng.choice(n_unlab, size=n_use, replace=False).tolist()\n", + " unlab_ds = Subset(unlab_all, use_indices)\n", + "\n", + " unlab_loader = DataLoader(unlab_ds, batch_size=BATCH_SIZE_U, shuffle=True,\n", + " num_workers=NUM_WORKERS, collate_fn=collate_unlabeled, drop_last=True)\n", + "\n", + " # ----- model -----\n", + " model, processor = build_model_and_processor()\n", + " model = model.to(DEVICE)\n", + "\n", + " optimizer = torch.optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)\n", + "\n", + " print(f\"Labeled train: {len(train_ds)} | Val: {len(val_ds)} | Unlabeled used: {len(unlab_ds)}\")\n", + "\n", + " best_val = 1e9\n", + " best_path = DATA_ROOT / \"best_mask2former_semi.pth\"\n", + "\n", + " pseudo_cache: Dict[int, torch.Tensor] = {}\n", + "\n", + " for epoch in range(1, EPOCHS + 1):\n", + " lambda_u = get_lambda_u(epoch)\n", + " tau = get_pseudo_th(epoch)\n", + " print(f\"\\nEpoch {epoch:02d}/{EPOCHS} | lambda_u={lambda_u:.3f} | tau={tau:.2f}\")\n", + "\n", + " # refresh pseudo cache\n", + " if (epoch == 1) or ((epoch - 1) % PSEUDO_UPDATE_EVERY == 0) or (len(pseudo_cache) == 0):\n", + " print(\"[Info] Building pseudo-label cache...\")\n", + " # build cache over current unlabeled subset loader (teacher = current model)\n", + " pseudo_cache = build_pseudo_cache(model, unlab_loader, tau=tau)\n", + " print(f\"[Info] Pseudo cache size: {len(pseudo_cache)}\")\n", + "\n", + " tr_loss = train_one_epoch(model, train_loader, unlab_loader, pseudo_cache, optimizer, lambda_u=lambda_u)\n", + " va_loss = eval_one_epoch(model, val_loader)\n", + "\n", + " print(f\"[Loss] train={tr_loss:.4f} | val(well6)={va_loss:.4f}\")\n", + "\n", + " if va_loss < best_val:\n", + " best_val = va_loss\n", + " torch.save(model.state_dict(), best_path)\n", + " print(f\"[OK] Best model saved: {best_path}\")\n", + "\n", + " # ----- inference -----\n", + " print(\"\\n[Info] Loading best model and generating submission...\")\n", + " state = torch.load(best_path, map_location=DEVICE, weights_only=True)\n", + " model.load_state_dict(state)\n", + "\n", + " out_csv = DATA_ROOT / \"submission.csv\"\n", + " predict_and_submit(model, processor, out_csv)\n", + "\n", + "\n", + "if __name__ == \"__main__\":\n", + " main()\n" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "deep-torch", + "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.10.19" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/UPerNet_Swin_ViT_backbone.ipynb b/UPerNet_Swin_ViT_backbone.ipynb index fb66ab5..aa51b45 100644 --- a/UPerNet_Swin_ViT_backbone.ipynb +++ b/UPerNet_Swin_ViT_backbone.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "code", - "execution_count": 4, + "execution_count": null, "id": "1d87bb20", "metadata": {}, "outputs": [ @@ -38,7 +38,6 @@ } ], "source": [ - "# -*- coding: utf-8 -*-\n", "\"\"\"\n", "UPerNet + Swin Transformer backbone\n", "FINAL VERSION (NO cv2, PyTorch interpolate only)\n", diff --git a/attention_u_net.ipynb b/attention_u_net.ipynb index 01beaa9..4857e32 100644 --- a/attention_u_net.ipynb +++ b/attention_u_net.ipynb @@ -1,8 +1,16 @@ { "cells": [ + { + "cell_type": "markdown", + "id": "210c55c1", + "metadata": {}, + "source": [ + "## Attention U-Net " + ] + }, { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "id": "9bd2e078", "metadata": {}, "outputs": [ @@ -44,7 +52,6 @@ } ], "source": [ - "# -*- coding: utf-8 -*-\n", "\"\"\"\n", "Attention U-Net - Full runnable semantic segmentation code (no extra libs)\n", "\n", diff --git a/deeplabv3.ipynb b/deeplabv3.ipynb index 093df65..a9435f5 100644 --- a/deeplabv3.ipynb +++ b/deeplabv3.ipynb @@ -1,8 +1,16 @@ { "cells": [ + { + "cell_type": "markdown", + "id": "48b2e6d5", + "metadata": {}, + "source": [ + "## DeepLabV3" + ] + }, { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "id": "35e48916", "metadata": {}, "outputs": [ @@ -51,7 +59,6 @@ } ], "source": [ - "# -*- coding: utf-8 -*-\n", "\"\"\"\n", "DeepLabV3（torchvision）分割：可直接跑（简单版）\n", "- 训练图像目录：X_train_uDRk9z9/images（well1-6）\n", diff --git "a/mask2_former_semi-supervis\303\251.ipynb" "b/mask2_former_semi-supervis\303\251.ipynb" index 141e7f0..e703985 100644 --- "a/mask2_former_semi-supervis\303\251.ipynb" +++ "b/mask2_former_semi-supervis\303\251.ipynb" @@ -1,5 +1,13 @@ { "cells": [ + { + "cell_type": "markdown", + "id": "bccb277e", + "metadata": {}, + "source": [ + "## Mask2former(semi-supervised)" + ] + }, { "cell_type": "code", "execution_count": null, diff --git a/mask2former.ipynb b/mask2former.ipynb index 98659bd..1cb7dd5 100644 --- a/mask2former.ipynb +++ b/mask2former.ipynb @@ -1,8 +1,16 @@ { "cells": [ + { + "cell_type": "markdown", + "id": "99f77dc9", + "metadata": {}, + "source": [ + "## Mask2former" + ] + }, { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "id": "e5988541", "metadata": {}, "outputs": [ @@ -73,7 +81,6 @@ } ], "source": [ - "# -*- coding: utf-8 -*-\n", "\"\"\"\n", "Mask2Former (HuggingFace Transformers) - Full runnable semantic segmentation code\n", "\n", diff --git a/resnet.ipynb b/resnet.ipynb index 73839af..5ca15d5 100644 --- a/resnet.ipynb +++ b/resnet.ipynb @@ -1,8 +1,16 @@ { "cells": [ + { + "cell_type": "markdown", + "id": "77e145ce", + "metadata": {}, + "source": [ + "## Res Net34-Unet" + ] + }, { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "id": "02e3cbc1", "metadata": {}, "outputs": [ @@ -52,14 +60,13 @@ } ], "source": [ - "# -*- coding: utf-8 -*-\n", "\"\"\"\n", - "ResNet34-UNet（简化版，可直接跑）\n", - "- 训练图像目录：X_train_uDRk9z9/images（well1-6）\n", - "- 测试图像目录：X_test_xNbnvIa/images（well7-11）\n", - "- 训练标签：Y_train_T9NrBYo.csv（flatten + -1 padding）\n", - "- 验证：从训练集中按井划分（例：well6 为 val，其余为 train）\n", - "- 输出：submission.csv（每行一个 patch，flatten，pad 到 160*272 用 -1）\n", + "ResNet34-UNet (simplified, runnable)\n", + "- Training images: X_train_uDRk9z9/images (well1-6)\n", + "- Test images: X_test_xNbnvIa/images (well7-11)\n", + "- Training labels: Y_train_T9NrBYo.csv (flatten + -1 padding)\n", + "- Validation: split by well from training set (e.g. well6 as val, others as train)\n", + "- Output: submission.csv (one row per patch, flattened, padded to 160*272 with -1)\n", "\"\"\"\n", "\n", "import re\n", @@ -77,9 +84,9 @@ "\n", "\n", "# =========================\n", - "# 0. 超参数与路径\n", + "# 0. Hyperparameters & Paths\n", "# =========================\n", - "DATA_ROOT = Path(r\"C:\\Users\\lenovo\\Desktop\\deep_datachallenge\") # 改成你的真实路径\n", + "DATA_ROOT = Path(r\"C:\\Users\\lenovo\\Desktop\\deep_datachallenge\") # change to your actual path\n", "\n", "TRAIN_IMAGES_DIR = DATA_ROOT / \"X_train_uDRk9z9\" / \"images\"\n", "TEST_IMAGES_DIR = DATA_ROOT / \"X_test_xNbnvIa\" / \"images\"\n", @@ -88,28 +95,28 @@ "TARGET_H = 160\n", "TARGET_W = 272\n", "\n", - "NUM_CLASSES = 3 # 你确认 CSV 里只有 0/1/2\n", + "NUM_CLASSES = 3 # you confirmed the CSV has only 0/1/2\n", "IGNORE_INDEX = -1 # CSV padding\n", "\n", "BATCH_SIZE = 8\n", "LR = 1e-3\n", "WEIGHT_DECAY = 1e-4\n", - "EPOCHS = 20 # 训练慢就先设 5，跑通后再加\n", + "EPOCHS = 20 # if training is slow, set to 5 first; increase after it runs\n", "\n", "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", "\n", "\n", "# =========================\n", - "# 1. 工具函数\n", + "# 1. Utility Functions\n", "# =========================\n", "def parse_well_id(name: str) -> int:\n", - " \"\"\"从 well_1_section_0_patch_0 提取 well id=1\"\"\"\n", + " \"\"\"Extract well id=1 from name like well_1_section_0_patch_0\"\"\"\n", " m = re.search(r\"well_(\\d+)_\", name)\n", " return int(m.group(1)) if m else -1\n", "\n", "\n", "def minmax_normalize(x: np.ndarray) -> np.ndarray:\n", - " \"\"\"min-max 归一化；NaN/inf 置 0\"\"\"\n", + " \"\"\"Min-max normalization; set NaN/inf to 0.\"\"\"\n", " x = x.astype(np.float32)\n", " x = np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)\n", " x_min = float(x.min())\n", @@ -120,9 +127,9 @@ "\n", "\n", "def pad_to_160x272(img: np.ndarray, fill_value: float = 0.0) -> np.ndarray:\n", - " \"\"\"把 (160,160) 或 (160,272) pad 到 (160,272)\"\"\"\n", + " \"\"\"Pad (160,160) or (160,272) to (160,272).\"\"\"\n", " h, w = img.shape\n", - " assert h == TARGET_H, f\"期望高度 {TARGET_H}，但拿到 {h}\"\n", + " assert h == TARGET_H, f\"Expected height {TARGET_H}, but got {h}\"\n", " if w == TARGET_W:\n", " return img\n", " if w < TARGET_W:\n", @@ -134,18 +141,18 @@ "\n", "def decode_mask_from_csv_row(row_values: np.ndarray) -> np.ndarray:\n", " \"\"\"\n", - " 从 CSV 一行恢复 mask：\n", + " Restore mask from one CSV row:\n", " - row_values: flatten + -1 padding\n", - " - 去掉 -1 后 reshape 成 (160, w)\n", + " - remove -1 then reshape to (160, w)\n", " \"\"\"\n", " valid = row_values[row_values != IGNORE_INDEX]\n", - " assert len(valid) % TARGET_H == 0, f\"mask 有效长度 {len(valid)} 不能被 160 整除\"\n", + " assert len(valid) % TARGET_H == 0, f\"Valid mask length {len(valid)} is not divisible by 160\"\n", " w = len(valid) // TARGET_H\n", " return valid.reshape(TARGET_H, w).astype(np.int64)\n", "\n", "\n", "def pad_mask_to_160x272(mask: np.ndarray) -> np.ndarray:\n", - " \"\"\"把 (160,w) pad 到 (160,272)，pad 用 -1（ignore）\"\"\"\n", + " \"\"\"Pad (160,w) to (160,272), pad with -1 (ignore).\"\"\"\n", " h, w = mask.shape\n", " assert h == TARGET_H\n", " if w == TARGET_W:\n", @@ -156,12 +163,12 @@ "\n", "\n", "# =========================\n", - "# 2. Dataset（训练/测试共用）\n", + "# 2. Dataset (shared for train/test)\n", "# =========================\n", "class WellSegDataset(Dataset):\n", " def __init__(self, images_dir: Path, y_csv_path: Path = None):\n", " \"\"\"\n", - " y_csv_path=None 表示无标签（测试）\n", + " y_csv_path=None indicates unlabeled data (test set).\n", " \"\"\"\n", " self.images_dir = images_dir\n", " self.has_label = y_csv_path is not None\n", @@ -170,7 +177,7 @@ " self.names = [p.stem for p in self.image_paths]\n", "\n", " if self.has_label:\n", - " # CSV index 通常就是 patch 名（不含 .npy）\n", + " # CSV index is usually patch name (without .npy)\n", " self.y_df = pd.read_csv(y_csv_path, index_col=0)\n", " else:\n", " self.y_df = None\n", @@ -183,7 +190,7 @@ " img_path = self.image_paths[idx]\n", "\n", " img = np.load(img_path) # (160,160) or (160,272)\n", - " raw_w = img.shape[1] # 记录原始宽度（推理时裁回去）\n", + " raw_w = img.shape[1] # record original width (crop back during inference)\n", " img = minmax_normalize(img)\n", " img = pad_to_160x272(img, fill_value=0.0)\n", " img_t = torch.from_numpy(img).unsqueeze(0).float() # (1,160,272)\n", @@ -200,7 +207,7 @@ "\n", "\n", "# =========================\n", - "# 3. ResNet34-UNet（简化实现）\n", + "# 3. ResNet34-UNet (simplified implementation)\n", "# =========================\n", "class ConvRelu(nn.Module):\n", " def __init__(self, in_ch, out_ch):\n", @@ -234,7 +241,7 @@ " super().__init__()\n", " backbone = resnet34(weights=ResNet34_Weights.IMAGENET1K_V1)\n", "\n", - " # 输入单通道：第一层卷积改成 1 通道（用原权重均值初始化）\n", + " # Single-channel input: change first conv to 1 channel (initialize by mean of original weights)\n", " old_conv1 = backbone.conv1\n", " new_conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)\n", " with torch.no_grad():\n", @@ -275,7 +282,7 @@ "\n", "\n", "# =========================\n", - "# 4. 训练与验证（最简单）\n", + "# 4. Training & Validation (minimal)\n", "# =========================\n", "def train_one_epoch(model, loader, optimizer):\n", " model.train()\n", @@ -312,15 +319,15 @@ "\n", "\n", "# =========================\n", - "# 5. 推理并生成提交 CSV\n", + "# 5. Inference & Submission CSV\n", "# =========================\n", "@torch.no_grad()\n", "def predict_and_make_submission(model, test_images_dir: Path, out_csv_path: Path):\n", " \"\"\"\n", - " 对 test_images_dir 全部 npy 预测并生成提交 CSV\n", - " - 每行：一个 patch\n", - " - 长度：160*272\n", - " - 如果原始宽度 < 272，剩余用 -1 padding\n", + " Predict all npy files in test_images_dir and generate submission CSV.\n", + " - one row per patch\n", + " - length: 160*272\n", + " - if original width < 272, pad the rest with -1\n", " \"\"\"\n", " model.eval()\n", "\n", @@ -337,7 +344,7 @@ " logits = model(x)\n", " pred_full = torch.argmax(logits, dim=1).squeeze(0).cpu().numpy().astype(np.int64) # (160,272)\n", "\n", - " pred = pred_full[:, :raw_w] # 裁回原始宽度\n", + " pred = pred_full[:, :raw_w] # crop back to original width\n", " if raw_w < TARGET_W:\n", " padded = np.full((TARGET_H * TARGET_W,), IGNORE_INDEX, dtype=np.int64)\n", " padded[: TARGET_H * raw_w] = pred.flatten()\n", @@ -347,17 +354,17 @@ "\n", " sub = pd.DataFrame(preds_dict, dtype=\"int64\").T\n", " sub.to_csv(out_csv_path)\n", - " print(f\"[OK] submission 已保存: {out_csv_path}\")\n", + " print(f\"[OK] Submission saved: {out_csv_path}\")\n", "\n", "\n", "# =========================\n", - "# 6. 主函数：训练(井1-5) + 验证(井6) + 预测test(井7-11目录)\n", + "# 6. Main: train (well1-5) + validate (well6) + predict test (well7-11)\n", "# =========================\n", "def main():\n", - " # (A) 构建训练集（well1-6）\n", + " # (A) Build the training dataset (well1-6)\n", " train_ds_all = WellSegDataset(TRAIN_IMAGES_DIR, Y_TRAIN_CSV)\n", "\n", - " # (B) 按井划分 train/val：well6 做验证\n", + " # (B) Split train/val by well: use well6 as validation\n", " VAL_WELLS = {6}\n", " train_indices, val_indices = [], []\n", " for i, name in enumerate(train_ds_all.names):\n", @@ -373,13 +380,13 @@ " train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)\n", " val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)\n", "\n", - " print(f\"训练样本数: {len(train_ds)} | 验证样本数: {len(val_ds)} | val_wells={VAL_WELLS}\")\n", + " print(f\"Train samples: {len(train_ds)} | Val samples: {len(val_ds)} | val_wells={VAL_WELLS}\")\n", "\n", - " # (C) 模型与优化器\n", + " # (C) Model & optimizer\n", " model = ResNet34UNet(num_classes=NUM_CLASSES).to(DEVICE)\n", " optimizer = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)\n", "\n", - " # (D) 训练\n", + " # (D) Train\n", " best_val = 1e9\n", " best_path = DATA_ROOT / \"best_resnet34_unet.pth\"\n", "\n", @@ -392,9 +399,9 @@ " if va_loss < best_val:\n", " best_val = va_loss\n", " torch.save(model.state_dict(), best_path)\n", - " print(f\" -> 保存最优模型: {best_path}\")\n", + " print(f\" -> Best model saved: {best_path}\")\n", "\n", - " # (E) 生成提交（测试目录 well7-11）\n", + " # (E) Generate submission (test directory well7-11)\n", " out_csv = DATA_ROOT / \"submission.csv\"\n", " model.load_state_dict(torch.load(best_path, map_location=DEVICE))\n", " predict_and_make_submission(model, TEST_IMAGES_DIR, out_csv)\n", diff --git a/swin-unet.ipynb b/swin-unet.ipynb index 3074856..a339616 100644 --- a/swin-unet.ipynb +++ b/swin-unet.ipynb @@ -1,8 +1,16 @@ { "cells": [ + { + "cell_type": "markdown", + "id": "55c7edca", + "metadata": {}, + "source": [ + "## Swin-Unet" + ] + }, { "cell_type": "code", - "execution_count": 2, + "execution_count": null, "id": "e433fe32", "metadata": {}, "outputs": [ @@ -44,7 +52,6 @@ } ], "source": [ - "# -*- coding: utf-8 -*-\n", "\"\"\"\n", "Swin-UNet (Swin Transformer encoder + UNet decoder) - Full runnable code (224x224 input for Swin)\n", "\n", diff --git a/upernet.ipynb b/upernet.ipynb index e47a758..19566ca 100644 --- a/upernet.ipynb +++ b/upernet.ipynb @@ -2,7 +2,7 @@ "cells": [ { "cell_type": "code", - "execution_count": 2, + "execution_count": null, "id": "6aaeb914", "metadata": {}, "outputs": [ @@ -53,18 +53,18 @@ } ], "source": [ - "# -*- coding: utf-8 -*-\n", "\"\"\"\n", - "UPerNet（ResNet50 backbone，torchvision预训练）完整可跑版本（简单实现）\n", - "- 训练图像目录：X_train_uDRk9z9/images（well1-6）\n", - "- 测试图像目录：X_test_xNbnvIa/images（well7-11）\n", - "- 训练标签：Y_train_T9NrBYo.csv（flatten + -1 padding）\n", - "- 验证：按井划分（例：well6 为 val，其余为 train）\n", - "- 输出：submission.csv（每行一个 patch，flatten，pad 到 160*272 用 -1）\n", - "\n", - "你只需要确认/修改：\n", - "1) DATA_ROOT 路径\n", - "2) EPOCHS / BATCH_SIZE（训练慢就先设小）\n", + "UPerNet (ResNet50 backbone, torchvision pretrained) - Full runnable version (simple implementation)\n", + "\n", + "- Training images: X_train_uDRk9z9/images (well1–6)\n", + "- Test images: X_test_xNbnvIa/images (well7–11)\n", + "- Training labels: Y_train_T9NrBYo.csv (flattened with -1 padding)\n", + "- Validation split: by well (e.g. well6 as validation, others as training)\n", + "- Output: submission.csv (one row per patch, flattened, padded to 160*272 with -1)\n", + "\n", + "You only need to check / modify:\n", + "1) DATA_ROOT path\n", + "2) EPOCHS / BATCH_SIZE (set smaller if training is slow)\n", "\"\"\"\n", "\n", "import re\n", @@ -82,9 +82,9 @@ "\n", "\n", "# =========================\n", - "# 0. 超参数与路径\n", + "# 0. Hyperparameters & Paths\n", "# =========================\n", - "DATA_ROOT = Path(r\"C:\\Users\\lenovo\\Desktop\\deep_datachallenge\") # 改成你的真实路径\n", + "DATA_ROOT = Path(r\"C:\\Users\\lenovo\\Desktop\\deep_datachallenge\") # change to your actual path\n", "\n", "TRAIN_IMAGES_DIR = DATA_ROOT / \"X_train_uDRk9z9\" / \"images\"\n", "TEST_IMAGES_DIR = DATA_ROOT / \"X_test_xNbnvIa\" / \"images\"\n", @@ -93,28 +93,28 @@ "TARGET_H = 160\n", "TARGET_W = 272\n", "\n", - "NUM_CLASSES = 3 # CSV 里只有 0/1/2\n", - "IGNORE_INDEX = -1 # CSV padding\n", + "NUM_CLASSES = 3 # only classes 0/1/2 in CSV\n", + "IGNORE_INDEX = -1 # padding value in CSV\n", "\n", - "BATCH_SIZE = 4 # UPerNet 比较吃显存，4060(8GB)建议先 2~4\n", - "LR = 1e-4 # UPerNet + 预训练 backbone 通常 lr 小一点更稳\n", + "BATCH_SIZE = 4 # UPerNet is memory-heavy; for RTX 4060 (8GB) start with 2–4\n", + "LR = 1e-4 # smaller LR is usually more stable with pretrained backbones\n", "WEIGHT_DECAY = 1e-4\n", - "EPOCHS = 20 # 训练慢就先设 20，跑通后再加\n", + "EPOCHS = 20 # reduce for quick debugging\n", "\n", "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", "\n", "\n", "# =========================\n", - "# 1. 工具函数\n", + "# 1. Utility Functions\n", "# =========================\n", "def parse_well_id(name: str) -> int:\n", - " \"\"\"从 well_1_section_0_patch_0 提取 well id=1\"\"\"\n", + " \"\"\"Extract well id from name like: well_1_section_0_patch_0 -> 1\"\"\"\n", " m = re.search(r\"well_(\\d+)_\", name)\n", " return int(m.group(1)) if m else -1\n", "\n", "\n", "def minmax_normalize(x: np.ndarray) -> np.ndarray:\n", - " \"\"\"min-max 归一化；NaN/inf 置 0\"\"\"\n", + " \"\"\"Min-max normalization; replace NaN/Inf with 0.\"\"\"\n", " x = x.astype(np.float32)\n", " x = np.nan_to_num(x, nan=0.0, posinf=0.0, neginf=0.0)\n", " x_min = float(x.min())\n", @@ -125,9 +125,9 @@ "\n", "\n", "def pad_to_160x272(img: np.ndarray, fill_value: float = 0.0) -> np.ndarray:\n", - " \"\"\"把 (160,160) 或 (160,272) pad 到 (160,272)\"\"\"\n", + " \"\"\"Pad (160,160) or (160,272) image to (160,272).\"\"\"\n", " h, w = img.shape\n", - " assert h == TARGET_H, f\"期望高度 {TARGET_H}，但拿到 {h}\"\n", + " assert h == TARGET_H, f\"Expected height {TARGET_H}, got {h}\"\n", " if w == TARGET_W:\n", " return img\n", " if w < TARGET_W:\n", @@ -139,18 +139,18 @@ "\n", "def decode_mask_from_csv_row(row_values: np.ndarray) -> np.ndarray:\n", " \"\"\"\n", - " 从 CSV 一行恢复 mask：\n", - " - row_values: flatten + -1 padding\n", - " - 去掉 -1 后 reshape 成 (160, w)\n", + " Restore mask from one CSV row:\n", + " - row_values: flattened mask with -1 padding\n", + " - remove -1 and reshape to (160, w)\n", " \"\"\"\n", " valid = row_values[row_values != IGNORE_INDEX]\n", - " assert len(valid) % TARGET_H == 0, f\"mask 有效长度 {len(valid)} 不能被 160 整除\"\n", + " assert len(valid) % TARGET_H == 0, f\"Valid mask length {len(valid)} not divisible by 160\"\n", " w = len(valid) // TARGET_H\n", " return valid.reshape(TARGET_H, w).astype(np.int64)\n", "\n", "\n", "def pad_mask_to_160x272(mask: np.ndarray) -> np.ndarray:\n", - " \"\"\"把 (160,w) pad 到 (160,272)，pad 用 -1（ignore）\"\"\"\n", + " \"\"\"Pad (160, w) mask to (160,272) using -1 as ignore index.\"\"\"\n", " h, w = mask.shape\n", " assert h == TARGET_H\n", " if w == TARGET_W:\n", @@ -161,12 +161,12 @@ "\n", "\n", "# =========================\n", - "# 2. Dataset（训练/测试共用）\n", + "# 2. Dataset (shared for train/test)\n", "# =========================\n", "class WellSegDataset(Dataset):\n", " def __init__(self, images_dir: Path, y_csv_path: Path = None):\n", " \"\"\"\n", - " y_csv_path=None 表示无标签（测试）\n", + " y_csv_path=None indicates unlabeled data (test set).\n", " \"\"\"\n", " self.images_dir = images_dir\n", " self.has_label = y_csv_path is not None\n", @@ -187,7 +187,7 @@ " img_path = self.image_paths[idx]\n", "\n", " img = np.load(img_path) # (160,160) or (160,272)\n", - " raw_w = img.shape[1] # 推理时裁回原始宽度\n", + " raw_w = img.shape[1] # used to crop back during inference\n", " img = minmax_normalize(img)\n", " img = pad_to_160x272(img, fill_value=0.0)\n", " img_t = torch.from_numpy(img).unsqueeze(0).float() # (1,160,272)\n", @@ -204,7 +204,7 @@ "\n", "\n", "# =========================\n", - "# 3. UPerNet Head（简化实现）\n", + "# 3. UPerNet Head (simplified)\n", "# =========================\n", "class ConvBNReLU(nn.Module):\n", " def __init__(self, in_ch, out_ch, k=3, p=1):\n", @@ -221,7 +221,7 @@ "\n", "class PSPModule(nn.Module):\n", " \"\"\"\n", - " PSP: 对最高层特征做多尺度池化，再concat回去\n", + " PSP module: multi-scale pooling on the highest-level feature.\n", " \"\"\"\n", " def __init__(self, in_ch, out_ch=256, pool_sizes=(1, 2, 3, 6)):\n", " super().__init__()\n", @@ -243,79 +243,68 @@ " y = F.interpolate(y, size=(h, w), mode=\"bilinear\", align_corners=False)\n", " priors.append(y)\n", " x = torch.cat(priors, dim=1)\n", - " return self.bottleneck(x) # (B, out_ch, h, w)\n", + " return self.bottleneck(x)\n", "\n", "\n", "class UPerHead(nn.Module):\n", " \"\"\"\n", " UPerNet Head = PSP + FPN\n", - " 输入特征：c2(1/4), c3(1/8), c4(1/16), c5(1/32)\n", + " Input features: c2(1/4), c3(1/8), c4(1/16), c5(1/32)\n", " \"\"\"\n", " def __init__(self, in_channels=(256, 512, 1024, 2048), fpn_dim=256, num_classes=3):\n", " super().__init__()\n", "\n", " c2, c3, c4, c5 = in_channels\n", "\n", - " # PSP 在最高层 c5 上\n", " self.psp = PSPModule(c5, out_ch=fpn_dim)\n", "\n", - " # lateral 1x1：把各层通道统一到 fpn_dim\n", " self.lateral_c2 = nn.Conv2d(c2, fpn_dim, kernel_size=1, bias=False)\n", " self.lateral_c3 = nn.Conv2d(c3, fpn_dim, kernel_size=1, bias=False)\n", " self.lateral_c4 = nn.Conv2d(c4, fpn_dim, kernel_size=1, bias=False)\n", "\n", - " # FPN 3x3：融合后再平滑\n", - " self.fpn_c2 = ConvBNReLU(fpn_dim, fpn_dim, k=3, p=1)\n", - " self.fpn_c3 = ConvBNReLU(fpn_dim, fpn_dim, k=3, p=1)\n", - " self.fpn_c4 = ConvBNReLU(fpn_dim, fpn_dim, k=3, p=1)\n", - " self.fpn_c5 = ConvBNReLU(fpn_dim, fpn_dim, k=3, p=1)\n", + " self.fpn_c2 = ConvBNReLU(fpn_dim, fpn_dim)\n", + " self.fpn_c3 = ConvBNReLU(fpn_dim, fpn_dim)\n", + " self.fpn_c4 = ConvBNReLU(fpn_dim, fpn_dim)\n", + " self.fpn_c5 = ConvBNReLU(fpn_dim, fpn_dim)\n", "\n", - " # 融合输出：concat(4层) -> 256\n", - " self.fuse = ConvBNReLU(fpn_dim * 4, fpn_dim, k=3, p=1)\n", + " self.fuse = ConvBNReLU(fpn_dim * 4, fpn_dim)\n", " self.cls = nn.Conv2d(fpn_dim, num_classes, kernel_size=1)\n", "\n", " def forward(self, c2, c3, c4, c5):\n", - " # PSP 输出 p5（仍在 1/32 分辨率）\n", - " p5 = self.psp(c5) # (B,256, H/32, W/32)\n", + " p5 = self.psp(c5)\n", "\n", - " # lateral\n", - " p4 = self.lateral_c4(c4) # (B,256, H/16, W/16)\n", - " p3 = self.lateral_c3(c3) # (B,256, H/8, W/8 )\n", - " p2 = self.lateral_c2(c2) # (B,256, H/4, W/4 )\n", + " p4 = self.lateral_c4(c4)\n", + " p3 = self.lateral_c3(c3)\n", + " p2 = self.lateral_c2(c2)\n", "\n", - " # top-down 融合\n", " p4 = p4 + F.interpolate(p5, size=p4.shape[-2:], mode=\"bilinear\", align_corners=False)\n", " p3 = p3 + F.interpolate(p4, size=p3.shape[-2:], mode=\"bilinear\", align_corners=False)\n", " p2 = p2 + F.interpolate(p3, size=p2.shape[-2:], mode=\"bilinear\", align_corners=False)\n", "\n", - " # 平滑\n", " p5 = self.fpn_c5(p5)\n", " p4 = self.fpn_c4(p4)\n", " p3 = self.fpn_c3(p3)\n", " p2 = self.fpn_c2(p2)\n", "\n", - " # 全部上采样到 p2 尺寸，再 concat\n", " p5_u = F.interpolate(p5, size=p2.shape[-2:], mode=\"bilinear\", align_corners=False)\n", " p4_u = F.interpolate(p4, size=p2.shape[-2:], mode=\"bilinear\", align_corners=False)\n", " p3_u = F.interpolate(p3, size=p2.shape[-2:], mode=\"bilinear\", align_corners=False)\n", "\n", " x = torch.cat([p2, p3_u, p4_u, p5_u], dim=1)\n", " x = self.fuse(x)\n", - " logits = self.cls(x) # (B, num_classes, H/4, W/4)\n", + " logits = self.cls(x)\n", " return logits\n", "\n", "\n", "# =========================\n", - "# 4. UPerNet（ResNet50 backbone）\n", + "# 4. UPerNet (ResNet50 backbone)\n", "# =========================\n", "class UPerNet(nn.Module):\n", " def __init__(self, num_classes: int):\n", " super().__init__()\n", "\n", - " # ResNet50 预训练\n", " backbone = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)\n", "\n", - " # 输入单通道：conv1 从 3 通道改为 1 通道（用均值初始化）\n", " old_conv1 = backbone.conv1\n", " new_conv1 = nn.Conv2d(\n", " in_channels=1,\n", @@ -330,32 +319,31 @@ " backbone.conv1 = new_conv1\n", "\n", " self.backbone = backbone\n", - "\n", - " # UPerHead：resnet50 的输出通道 (layer1..4) = 256,512,1024,2048\n", - " self.head = UPerHead(in_channels=(256, 512, 1024, 2048), fpn_dim=256, num_classes=num_classes)\n", + " self.head = UPerHead(\n", + " in_channels=(256, 512, 1024, 2048),\n", + " fpn_dim=256,\n", + " num_classes=num_classes,\n", + " )\n", "\n", " def forward(self, x):\n", - " # ResNet 前几层\n", " x = self.backbone.conv1(x)\n", " x = self.backbone.bn1(x)\n", " x = self.backbone.relu(x)\n", " x = self.backbone.maxpool(x)\n", "\n", - " # 取各层特征\n", - " c2 = self.backbone.layer1(x) # 1/4\n", - " c3 = self.backbone.layer2(c2) # 1/8\n", - " c4 = self.backbone.layer3(c3) # 1/16\n", - " c5 = self.backbone.layer4(c4) # 1/32\n", + " c2 = self.backbone.layer1(x)\n", + " c3 = self.backbone.layer2(c2)\n", + " c4 = self.backbone.layer3(c3)\n", + " c5 = self.backbone.layer4(c4)\n", "\n", " logits_1_4 = self.head(c2, c3, c4, c5)\n", - "\n", - " # 上采样到原图尺寸 (160,272)\n", - " logits = F.interpolate(logits_1_4, size=(TARGET_H, TARGET_W), mode=\"bilinear\", align_corners=False)\n", + " logits = F.interpolate(logits_1_4, size=(TARGET_H, TARGET_W),\n", + " mode=\"bilinear\", align_corners=False)\n", " return logits\n", "\n", "\n", "# =========================\n", - "# 5. 训练与验证（最简单）\n", + "# 5. Training & Validation\n", "# =========================\n", "def train_one_epoch(model, loader, optimizer):\n", " model.train()\n", @@ -392,15 +380,15 @@ "\n", "\n", "# =========================\n", - "# 6. 推理并生成提交 CSV\n", + "# 6. Inference & Submission\n", "# =========================\n", "@torch.no_grad()\n", "def predict_and_make_submission(model, test_images_dir: Path, out_csv_path: Path):\n", " \"\"\"\n", - " 对 test_images_dir 全部 npy 预测并生成提交 CSV\n", - " - 每行：一个 patch\n", - " - 长度：160*272\n", - " - 如果原始宽度 < 272，剩余用 -1 padding\n", + " Run inference on all npy files and generate submission CSV:\n", + " - one row per patch\n", + " - length = 160*272\n", + " - pad with -1 if original width < 272\n", " \"\"\"\n", " model.eval()\n", "\n", @@ -415,9 +403,9 @@ " x = batch[\"image\"].to(DEVICE)\n", "\n", " logits = model(x)\n", - " pred_full = torch.argmax(logits, dim=1).squeeze(0).cpu().numpy().astype(np.int64) # (160,272)\n", + " pred_full = torch.argmax(logits, dim=1).squeeze(0).cpu().numpy().astype(np.int64)\n", "\n", - " pred = pred_full[:, :raw_w] # 裁回原始宽度\n", + " pred = pred_full[:, :raw_w]\n", " if raw_w < TARGET_W:\n", " padded = np.full((TARGET_H * TARGET_W,), IGNORE_INDEX, dtype=np.int64)\n", " padded[: TARGET_H * raw_w] = pred.flatten()\n", @@ -427,17 +415,15 @@ "\n", " sub = pd.DataFrame(preds_dict, dtype=\"int64\").T\n", " sub.to_csv(out_csv_path)\n", - " print(f\"[OK] submission 已保存: {out_csv_path}\")\n", + " print(f\"[OK] Submission saved to: {out_csv_path}\")\n", "\n", "\n", "# =========================\n", - "# 7. 主函数：训练(井1-5) + 验证(井6) + 预测test(井7-11目录)\n", + "# 7. Main\n", "# =========================\n", "def main():\n", - " # (A) 构建训练集（well1-6）\n", " train_ds_all = WellSegDataset(TRAIN_IMAGES_DIR, Y_TRAIN_CSV)\n", "\n", - " # (B) 按井划分 train/val：well6 做验证\n", " VAL_WELLS = {6}\n", " train_indices, val_indices = [], []\n", " for i, name in enumerate(train_ds_all.names):\n", @@ -447,35 +433,31 @@ " else:\n", " train_indices.append(i)\n", "\n", - " train_ds = Subset(train_ds_all, train_indices) # well1-5\n", - " val_ds = Subset(train_ds_all, val_indices) # well6\n", + " train_ds = Subset(train_ds_all, train_indices)\n", + " val_ds = Subset(train_ds_all, val_indices)\n", "\n", " train_loader = DataLoader(train_ds, batch_size=BATCH_SIZE, shuffle=True, num_workers=0)\n", " val_loader = DataLoader(val_ds, batch_size=BATCH_SIZE, shuffle=False, num_workers=0)\n", "\n", - " print(f\"训练样本数: {len(train_ds)} | 验证样本数: {len(val_ds)} | val_wells={VAL_WELLS}\")\n", + " print(f\"Training samples: {len(train_ds)} | Validation samples: {len(val_ds)} | val_wells={VAL_WELLS}\")\n", " print(f\"DEVICE: {DEVICE}\")\n", "\n", - " # (C) 模型与优化器\n", " model = UPerNet(num_classes=NUM_CLASSES).to(DEVICE)\n", " optimizer = torch.optim.Adam(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)\n", "\n", - " # (D) 训练\n", " best_val = 1e9\n", " best_path = DATA_ROOT / \"best_upernet.pth\"\n", "\n", " for epoch in range(1, EPOCHS + 1):\n", " tr_loss = train_one_epoch(model, train_loader, optimizer)\n", " va_loss = valid_one_epoch(model, val_loader)\n", - "\n", " print(f\"Epoch {epoch:02d}/{EPOCHS} | train_loss={tr_loss:.4f} | val_loss={va_loss:.4f}\")\n", "\n", " if va_loss < best_val:\n", " best_val = va_loss\n", " torch.save(model.state_dict(), best_path)\n", - " print(f\" -> 保存最优模型: {best_path}\")\n", + " print(f\" -> Best model saved: {best_path}\")\n", "\n", - " # (E) 生成提交（测试目录 well7-11）\n", " out_csv = DATA_ROOT / \"submission.csv\"\n", " state_dict = torch.load(best_path, map_location=DEVICE, weights_only=True)\n", " model.load_state_dict(state_dict)\n", From 3cbf916d67293ad9bb81261f90f79337d750db39 Mon Sep 17 00:00:00 2001 From: Haokun Date: Sun, 11 Jan 2026 20:21:28 +0100 Subject: [PATCH 11/11] =?UTF-8?q?Update=20mask2=5Fformer=5Fsemi-supervis?= =?UTF-8?q?=C3=A9.ipynb?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- "mask2_former_semi-supervis\303\251.ipynb" | 28 +++++++++++----------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git "a/mask2_former_semi-supervis\303\251.ipynb" "b/mask2_former_semi-supervis\303\251.ipynb" index e703985..d0e8421 100644 --- "a/mask2_former_semi-supervis\303\251.ipynb" +++ "b/mask2_former_semi-supervis\303\251.ipynb" @@ -61,7 +61,6 @@ } ], "source": [ - "\n", "\"\"\"\n", "Mask2Former + Semi-Supervised Pseudo Label (Full runnable)0.6663556049927326\n", "\n", @@ -128,8 +127,8 @@ "EPOCHS = 10\n", "\n", "# semi-supervised hyperparams\n", - "PSEUDO_TH = 0.85 # 伪标签置信度阈值(越高越保守)\n", - "LAMBDA_U = 0.5 # 无标签loss权重(0.2~1.0可调)\n", + "PSEUDO_TH = 0.85 # Pseudo-label confidence threshold (higher = more conservative)\n", + "LAMBDA_U = 0.5 # Unlabeled loss weight (tunable in ~0.2~1.0)\n", "\n", "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", "\n", @@ -241,10 +240,10 @@ "\n", "\n", "# =========================\n", - "# 2) 简单增强（无标签用）\n", + "# 2) Simple augmentations (for unlabeled data)\n", "# =========================\n", "def aug_weak(x: torch.Tensor) -> torch.Tensor:\n", - " \"\"\"弱增强：随机左右翻转 + 轻噪声\"\"\"\n", + " \"\"\"Weak aug: random horizontal flip + light noise\"\"\"\n", " # x: (1,224,224)\n", " if torch.rand(1).item() < 0.5:\n", " x = torch.flip(x, dims=[2])\n", @@ -253,7 +252,7 @@ "\n", "\n", "def aug_strong(x: torch.Tensor) -> torch.Tensor:\n", - " \"\"\"强增强：随机翻转 + 更强噪声 + 亮度对比度扰动\"\"\"\n", + " \"\"\"Strong aug: random flip + stronger noise + brightness/contrast jitter\"\"\"\n", " if torch.rand(1).item() < 0.5:\n", " x = torch.flip(x, dims=[2])\n", " # brightness/contrast\n", @@ -315,7 +314,7 @@ " img_t = torch.from_numpy(img).unsqueeze(0).float() # (1,160,272)\n", " img_t = resize_image_torch(img_t, MODEL_H, MODEL_W) # (1,224,224)\n", "\n", - " # 返回 base 图像（增强在 collate 做）\n", + " # Return the base image (augmentations are applied in collate)\n", " return {\"name\": name, \"image\": img_t, \"raw_w\": raw_w}\n", "\n", "\n", @@ -393,10 +392,10 @@ "@torch.no_grad()\n", "def pseudo_from_outputs(outputs, num_classes: int) -> Tuple[torch.Tensor, torch.Tensor]:\n", " \"\"\"\n", - " 从 Mask2Former 输出构造像素级类别概率：\n", + " Build pixel-level class scores from Mask2Former outputs:\n", " class_probs (softmax) * mask_probs (sigmoid) -> per-pixel scores\n", "\n", - " outputs.class_queries_logits: (B, Q, C+1) (最后一个通常是 no-object)\n", + " outputs.class_queries_logits: (B, Q, C+1) (last one is usually no-object)\n", " outputs.masks_queries_logits: (B, Q, H, W)\n", "\n", " return:\n", @@ -457,11 +456,11 @@ " out_u_teacher = model(pixel_values=pv_w, pixel_mask=pm_u)\n", " pseudo, conf = pseudo_from_outputs(out_u_teacher, NUM_CLASSES) # (B,224,224)\n", "\n", - " # 置信度过滤：低于阈值的像素设为 IGNORE\n", + " # Confidence filtering: pixels below threshold are set to IGNORE\n", " pseudo = pseudo.clone()\n", " pseudo[conf < PSEUDO_TH] = IGNORE_INDEX\n", "\n", - " # 将 pseudo semantic mask -> mask2former targets(list)\n", + " # Convert pseudo semantic mask -> mask2former targets(list)\n", " class_labels_u, mask_labels_u = [], []\n", " for i in range(pseudo.shape[0]):\n", " cls_i, msk_i = semantic_to_mask2former_targets(pseudo[i], NUM_CLASSES, IGNORE_INDEX)\n", @@ -523,7 +522,7 @@ "def predict_and_make_submission(model, processor, test_images_dir: Path, out_csv_path: Path):\n", " model.eval()\n", "\n", - " # 这里复用 UnlabeledDataset 结构（只有图像，无标签）\n", + " # Reuse UnlabeledDataset structure (image-only, no labels)\n", " test_ds = UnlabeledDataset(test_images_dir)\n", "\n", " def collate_test(batch: List[Dict]) -> Dict:\n", @@ -547,7 +546,7 @@ "\n", " outputs = model(pixel_values=pixel_values, pixel_mask=pixel_mask)\n", "\n", - " # 使用 processor 的语义后处理（稳定）\n", + " # Use processor's semantic post-processing (stable)\n", " seg_list = processor.post_process_semantic_segmentation(outputs, target_sizes=[(MODEL_H, MODEL_W)])\n", " seg_224 = seg_list[0].to(torch.int64) # (224,224)\n", "\n", @@ -645,7 +644,8 @@ "\n", "\n", "if __name__ == \"__main__\":\n", - " main()\n" + " main()\n", + "```\n" ] } ],