diff --git a/DPT.ipynb b/DPT.ipynb new file mode 100644 index 0000000..a85d708 --- /dev/null +++ b/DPT.ipynb @@ -0,0 +1,401 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "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": [ + "\"\"\"\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/Mask2Former_Semi_Supervised2.ipynb b/Mask2Former_Semi_Supervised2.ipynb new file mode 100644 index 0000000..76572bc --- /dev/null +++ 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/README.md b/README.md index 14a4873..31645e3 100644 --- a/README.md +++ b/README.md @@ -1 +1,150 @@ -# datachallenge_deep \ No newline at end of file +# datachallenge_deep + +transformer最强 + + +下面按“类型”把**已经用过**的模型和**还可以用**的模型列出来（只列结构名，不展开细节）。 + +--- + +## 1) U-Net 系列（Encoder–Decoder / Skip Connection） + +### 你已用 + +* **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）** +* **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 系列语义分割 + +### 你已用 + +* **Mask2Former + Semi-Supervised Pseudo Label (Full runnable)0.6663556049927326 +* **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** + +### 仍可用 + + +--- + +## 4) Pyramid / FPN 多尺度融合（CNN 语义分割头） + +### 你已用 + +* **UPerNet（PSP + FPN）0.5387714111687977** + +### 仍可用 + +* **PSPNet** +* **FPN-Seg** +* **HRNet + OCR（高分辨率）**3 +* **PAN（Path Aggregation Network）** + +--- + +## 5) “边界/形状友好”的结构（适合细长界面） + +### 你已用 + +* （暂无） + +### 仍可用 + +* **Boundary Head（分割 + 边界分支）** +* **Distance Transform 回归 + 分割** +* **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..aa51b45 --- /dev/null +++ b/UPerNet_Swin_ViT_backbone.ipynb @@ -0,0 +1,329 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "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": [ + "\"\"\"\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..4857e32 --- /dev/null +++ b/attention_u_net.ipynb @@ -0,0 +1,508 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "210c55c1", + "metadata": {}, + "source": [ + "## Attention U-Net " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "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": [ + "\"\"\"\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/deeplabv3.ipynb b/deeplabv3.ipynb new file mode 100644 index 0000000..a9435f5 --- /dev/null +++ b/deeplabv3.ipynb @@ -0,0 +1,398 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "48b2e6d5", + "metadata": {}, + "source": [ + "## DeepLabV3" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "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": [ + "\"\"\"\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/mask2_former_semi-supervis\303\251.ipynb" "b/mask2_former_semi-supervis\303\251.ipynb" new file mode 100644 index 0000000..d0e8421 --- /dev/null +++ "b/mask2_former_semi-supervis\303\251.ipynb" @@ -0,0 +1,673 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "bccb277e", + "metadata": {}, + "source": [ + "## Mask2former(semi-supervised)" + ] + }, + { + "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", + "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 # 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", + "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) Simple augmentations (for unlabeled data)\n", + "# =========================\n", + "def aug_weak(x: torch.Tensor) -> torch.Tensor:\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", + " 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", + " \"\"\"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", + " 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", + " # Return the base image (augmentations are applied in 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", + " 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) (last one is usually 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", + " # Confidence filtering: pixels below threshold are set to IGNORE\n", + " pseudo = pseudo.clone()\n", + " pseudo[conf < PSEUDO_TH] = IGNORE_INDEX\n", + "\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", + " 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", + " # Reuse UnlabeledDataset structure (image-only, no labels)\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", + " # 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", + " 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", + "```\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..1cb7dd5 --- /dev/null +++ b/mask2former.ipynb @@ -0,0 +1,599 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "99f77dc9", + "metadata": {}, + "source": [ + "## Mask2former" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "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": [ + "\"\"\"\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/resnet.ipynb b/resnet.ipynb new file mode 100644 index 0000000..5ca15d5 --- /dev/null +++ b/resnet.ipynb @@ -0,0 +1,436 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "77e145ce", + "metadata": {}, + "source": [ + "## Res Net34-Unet" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "02e3cbc1", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "训练样本数: 2790 | 验证样本数: 1620 | val_wells={6}\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/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 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_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" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "[OK] submission 已保存: C:\\Users\\lenovo\\Desktop\\deep_datachallenge\\submission.csv\n" + ] + } + ], + "source": [ + "\"\"\"\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", + "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. Hyperparameters & Paths\n", + "# =========================\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", + "Y_TRAIN_CSV = DATA_ROOT / \"Y_train_T9NrBYo.csv\"\n", + "\n", + "TARGET_H = 160\n", + "TARGET_W = 272\n", + "\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 # 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. Utility Functions\n", + "# =========================\n", + "def parse_well_id(name: str) -> int:\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 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", + " 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}, but 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", + " Restore mask from one CSV row:\n", + " - row_values: flatten + -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)} 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", + " \"\"\"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", + " 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 (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 indicates unlabeled data (test set).\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 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", + "\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] # 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", + "\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 (simplified implementation)\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", + " # 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", + " 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. Training & Validation (minimal)\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. 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", + " 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", + " 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 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", + " 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: {out_csv_path}\")\n", + "\n", + "\n", + "# =========================\n", + "# 6. Main: train (well1-5) + validate (well6) + predict test (well7-11)\n", + "# =========================\n", + "def main():\n", + " # (A) Build the training dataset (well1-6)\n", + " train_ds_all = WellSegDataset(TRAIN_IMAGES_DIR, Y_TRAIN_CSV)\n", + "\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", + " 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 = ResNet34UNet(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_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 model saved: {best_path}\")\n", + "\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", + "\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..a339616 --- /dev/null +++ b/swin-unet.ipynb @@ -0,0 +1,441 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "55c7edca", + "metadata": {}, + "source": [ + "## Swin-Unet" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "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": [ + "\"\"\"\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 +} diff --git a/upernet.ipynb b/upernet.ipynb new file mode 100644 index 0000000..19566ca --- /dev/null +++ b/upernet.ipynb @@ -0,0 +1,493 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "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": [ + "\"\"\"\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", + "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. Hyperparameters & Paths\n", + "# =========================\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", + "Y_TRAIN_CSV = DATA_ROOT / \"Y_train_T9NrBYo.csv\"\n", + "\n", + "TARGET_H = 160\n", + "TARGET_W = 272\n", + "\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 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 # reduce for quick debugging\n", + "\n", + "DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n", + "\n", + "\n", + "# =========================\n", + "# 1. Utility Functions\n", + "# =========================\n", + "def parse_well_id(name: str) -> int:\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 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", + " 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) image 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", + " 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\"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) 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", + " 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 (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 indicates unlabeled data (test set).\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 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", + "\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 (simplified)\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 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", + " 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)\n", + "\n", + "\n", + "class UPerHead(nn.Module):\n", + " \"\"\"\n", + " UPerNet Head = PSP + FPN\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", + " self.psp = PSPModule(c5, out_ch=fpn_dim)\n", + "\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", + " 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", + " 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", + " p5 = self.psp(c5)\n", + "\n", + " p4 = self.lateral_c4(c4)\n", + " p3 = self.lateral_c3(c3)\n", + " p2 = self.lateral_c2(c2)\n", + "\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", + " 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", + " 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)\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", + " backbone = resnet50(weights=ResNet50_Weights.IMAGENET1K_V2)\n", + "\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", + " 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", + " 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", + " 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", + " 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. Training & Validation\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. Inference & Submission\n", + "# =========================\n", + "@torch.no_grad()\n", + "def predict_and_make_submission(model, test_images_dir: Path, out_csv_path: Path):\n", + " \"\"\"\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", + " 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)\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 saved to: {out_csv_path}\")\n", + "\n", + "\n", + "# =========================\n", + "# 7. Main\n", + "# =========================\n", + "def main():\n", + " train_ds_all = WellSegDataset(TRAIN_IMAGES_DIR, Y_TRAIN_CSV)\n", + "\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)\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\"Training samples: {len(train_ds)} | Validation samples: {len(val_ds)} | val_wells={VAL_WELLS}\")\n", + " print(f\"DEVICE: {DEVICE}\")\n", + "\n", + " model = UPerNet(num_classes=NUM_CLASSES).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_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", + " 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", + " 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 +}