{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":8089,"databundleVersionId":44321,"sourceType":"competition"}],"dockerImageVersionId":31154,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# 0. Unzip Data","metadata":{}},{"cell_type":"code","source":"# Unzip data\nimport zipfile\nwith zipfile.ZipFile('/kaggle/input/data-science-bowl-2018/'+ 'stage1_train.zip', 'r') as zip_ref:\n    zip_ref.extractall('./train')\n    \nwith zipfile.ZipFile('/kaggle/input/data-science-bowl-2018/' + 'stage1_test.zip', 'r') as zip_ref:\n    zip_ref.extractall('./test')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T16:46:33.618222Z","iopub.execute_input":"2025-11-03T16:46:33.618498Z","iopub.status.idle":"2025-11-03T16:46:38.538904Z","shell.execute_reply.started":"2025-11-03T16:46:33.618471Z","shell.execute_reply":"2025-11-03T16:46:38.537794Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 1. Imports và Cấu hình","metadata":{}},{"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader, random_split\nimport torchvision.transforms.functional as TF\n\nimport os\nimport numpy as np\nimport random\nfrom tqdm import tqdm\nfrom skimage.io import imread\nfrom skimage.transform import resize\nimport matplotlib.pyplot as plt\n\n# --- ⚙️ CONFIGURATION ---\nTRAIN_DIR = \"/kaggle/working/train\"\nTEST_DIR = \"/kaggle/working/test\"\n\nIMG_WIDTH = 128\nIMG_HEIGHT = 128\nIMG_CHANNELS = 3\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nBATCH_SIZE = 16\nLEARNING_RATE = 1e-4\nEPOCHS = 100\nSEED = 42\n\nVALIDATION_SPLIT = 0.1 # 10%\nTEST_SPLIT = 0.2       # 20%\nEARLY_STOP_PATIENCE = 15\n\n# --- THAY ĐỔI: Đường dẫn lưu model riêng biệt ---\nUNET_SAVE_PATH = \"/kaggle/working/best_unet_model.pth\"\nUNET_PLUS_PLUS_SAVE_PATH = \"/kaggle/working/best_unet_pp_model.pth\"\n# -----------------------------------------------\n\n# Cài đặt seed\nrandom.seed(SEED)\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\nif torch.cuda.is_available():\n    torch.cuda.manual_seed(SEED)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n\nprint(f\"Using device: {DEVICE}\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-11-03T16:46:38.540908Z","iopub.execute_input":"2025-11-03T16:46:38.541192Z","iopub.status.idle":"2025-11-03T16:46:46.169911Z","shell.execute_reply.started":"2025-11-03T16:46:38.541162Z","shell.execute_reply":"2025-11-03T16:46:46.169144Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 2. EDA, Visualize và Pre-load dữ liệu","metadata":{}},{"cell_type":"code","source":"print(\"--- 2. Starting EDA & Data Pre-loading ---\")\n\ntrain_ids = next(os.walk(TRAIN_DIR))[1]\n\n# Lists to store all data\nX_data = [] # Original images (resized)\nY_data = [] # Combined masks (resized)\n\n# Stats for EDA\noriginal_sizes = []\nmask_counts = []\n\nprint('Loading, resizing, and collecting stats...')\nfor id_ in tqdm(train_ids):\n    path = os.path.join(TRAIN_DIR, id_)\n    \n    # --- Load Image & Stats ---\n    img_path = os.path.join(path, 'images', id_ + '.png')\n    img_original = imread(img_path)\n    original_sizes.append(img_original.shape[:2]) # (H, W)\n    \n    # Resize image\n    img = resize(img_original[:,:,:IMG_CHANNELS], (IMG_HEIGHT, IMG_WIDTH), mode='constant', preserve_range=True)\n    X_data.append(img.astype(np.uint8))\n\n    # --- Load Masks & Stats ---\n    mask_path = os.path.join(path, 'masks')\n    mask_files = next(os.walk(mask_path))[2]\n    mask_counts.append(len(mask_files))\n    \n    # Combine masks\n    mask_combined = np.zeros((IMG_HEIGHT, IMG_WIDTH, 1), dtype=np.bool_)\n    for mask_file in mask_files:\n        mask_ = imread(os.path.join(mask_path, mask_file))\n        mask_ = np.expand_dims(resize(mask_, (IMG_HEIGHT, IMG_WIDTH), mode='constant', preserve_range=True), axis=-1)\n        mask_combined = np.maximum(mask_combined, mask_)\n    Y_data.append(mask_combined)\n\nprint(f\"Data pre-loading complete. Loaded {len(X_data)} images.\")\n\n# --- Visualize Stats ---\nprint(\"\\n--- EDA Visualizations ---\")\n\n# 1. Histogram of mask counts (Số lượng hạt nhân trên mỗi ảnh)\nplt.figure(figsize=(14, 6))\nplt.subplot(1, 2, 1)\nplt.hist(mask_counts, bins=30, color='blue', alpha=0.7)\nplt.title('Histogram of Nuclei Count per Image')\nplt.xlabel('Number of Nuclei (Masks)')\nplt.ylabel('Frequency')\nplt.grid(True, linestyle='--', alpha=0.5)\n\n# 2. Scatter plot of original image sizes (Kích thước ảnh gốc)\nplt.subplot(1, 2, 2)\nsizes_np = np.array(original_sizes)\nplt.scatter(sizes_np[:, 1], sizes_np[:, 0], alpha=0.3) # x=Width, y=Height\nplt.title('Original Image Sizes (Width vs Height)')\nplt.xlabel('Width (pixels)')\nplt.ylabel('Height (pixels)')\nplt.grid(True, linestyle='--', alpha=0.5)\nplt.tight_layout()\nplt.show()\n\n# 3. Visualize a random sample (Xem ảnh và mask mẫu)\nprint(\"\\n--- Sample Image and Combined Mask ---\")\nix = random.randint(0, len(X_data) - 1)\nplt.figure(figsize=(10, 5))\nplt.subplot(1, 2, 1)\nplt.imshow(X_data[ix])\nplt.title(f'Sample Image')\nplt.axis('off')\n\nplt.subplot(1, 2, 2)\nplt.imshow(np.squeeze(Y_data[ix]), cmap='gray')\nplt.title(f'Combined Mask ({mask_counts[ix]} nuclei)')\nplt.axis('off')\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T16:46:46.170739Z","iopub.execute_input":"2025-11-03T16:46:46.171036Z","iopub.status.idle":"2025-11-03T16:51:18.923417Z","shell.execute_reply.started":"2025-11-03T16:46:46.171019Z","shell.execute_reply":"2025-11-03T16:51:18.922404Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 3. Visualize 5 random samples\nprint(\"\\n--- 5 Sample Images and Combined Masks ---\")\n\n# Lấy 5 index ngẫu nhiên (không trùng lặp) từ bộ dữ liệu\nindices = random.sample(range(len(X_data)), 5)\n\nfor ix in indices:\n    # Lấy dữ liệu\n    image = X_data[ix]\n    mask = Y_data[ix]\n    mask_count = mask_counts[ix]\n    image_id = train_ids[ix]\n\n    # Vẽ\n    plt.figure(figsize=(10, 5))\n    \n    # Ảnh gốc\n    plt.subplot(1, 2, 1)\n    plt.imshow(image)\n    plt.title(f'Sample Image ({ix})')\n    plt.axis('off')\n\n    # Mask tổng hợp\n    plt.subplot(1, 2, 2)\n    plt.imshow(np.squeeze(mask), cmap='gray')\n    plt.title(f'Combined Mask ({mask_count} nuclei)')\n    plt.axis('off')\n    \n    plt.show() # Hiển thị từng cặp ảnh","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T16:51:18.924424Z","iopub.execute_input":"2025-11-03T16:51:18.924753Z","iopub.status.idle":"2025-11-03T16:51:20.008359Z","shell.execute_reply.started":"2025-11-03T16:51:18.924722Z","shell.execute_reply":"2025-11-03T16:51:20.007537Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 3. Định nghĩa Model U-Net","metadata":{}},{"cell_type":"code","source":"class ConvBlock(nn.Module):\n    \"\"\"(Conv2d -> BN -> ReLU) x 2\"\"\"\n    # LƯU Ý: Đây là block giống hệt như trong UNet++\n    # Nếu bạn đã định nghĩa nó ở Cell 3 (UNet++) rồi thì không cần định nghĩa lại\n    # Nhưng để cell này độc lập, tôi sẽ thêm nó vào đây.\n    def __init__(self, in_channels, out_channels, mid_channels=None):\n        super().__init__()\n        if not mid_channels:\n            mid_channels = out_channels\n        self.double_conv = nn.Sequential(\n            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),\n            nn.BatchNorm2d(mid_channels),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True)\n        )\n\n    def forward(self, x):\n        return self.double_conv(x)\n\nclass UNet(nn.Module):\n    \"\"\"\n    Implementation of U-Net (Viết theo style của UNet++)\n    Sử dụng filter count: [32, 64, 128, 256, 512]\n    \"\"\"\n    def __init__(self, in_channels=3, out_channels=1, features=[32, 64, 128, 256, 512]):\n        super(UNet, self).__init__()\n        nb_filter = features\n\n        self.pool = nn.MaxPool2d(2, 2)\n        # Sử dụng ConvTranspose2d cho nhất quán với U-Net gốc\n        # (UNet++ dùng Upsample, nhưng ConvTranspose2d cũng rất phổ biến)\n        # Bạn có thể đổi self.up4_t thành nn.Upsample(...) nếu muốn\n        \n        # --- Encoder ---\n        self.conv0_0 = ConvBlock(in_channels, nb_filter[0])\n        self.conv1_0 = ConvBlock(nb_filter[0], nb_filter[1])\n        self.conv2_0 = ConvBlock(nb_filter[1], nb_filter[2])\n        self.conv3_0 = ConvBlock(nb_filter[2], nb_filter[3])\n        \n        # --- Bottleneck ---\n        self.conv4_0 = ConvBlock(nb_filter[3], nb_filter[4])\n\n        # --- Decoder ---\n        # Up-sampling + ConvBlock\n        self.up4_t = nn.ConvTranspose2d(nb_filter[4], nb_filter[3], kernel_size=2, stride=2)\n        self.conv3_1 = ConvBlock(nb_filter[4], nb_filter[3]) # (skip + up) -> conv\n        \n        self.up3_t = nn.ConvTranspose2d(nb_filter[3], nb_filter[2], kernel_size=2, stride=2)\n        self.conv2_1 = ConvBlock(nb_filter[3], nb_filter[2])\n        \n        self.up2_t = nn.ConvTranspose2d(nb_filter[2], nb_filter[1], kernel_size=2, stride=2)\n        self.conv1_1 = ConvBlock(nb_filter[2], nb_filter[1])\n        \n        self.up1_t = nn.ConvTranspose2d(nb_filter[1], nb_filter[0], kernel_size=2, stride=2)\n        self.conv0_1 = ConvBlock(nb_filter[1], nb_filter[0])\n\n        # --- Final Output ---\n        self.final = nn.Conv2d(nb_filter[0], out_channels, kernel_size=1)\n\n    def forward(self, x):\n        # --- Encoder path ---\n        x0_0 = self.conv0_0(x)       # Skip 1 (32 filters)\n        x1_0 = self.conv1_0(self.pool(x0_0)) # Skip 2 (64 filters)\n        x2_0 = self.conv2_0(self.pool(x1_0)) # Skip 3 (128 filters)\n        x3_0 = self.conv3_0(self.pool(x2_0)) # Skip 4 (256 filters)\n\n        # --- Bottleneck ---\n        x4_0 = self.conv4_0(self.pool(x3_0)) # (512 filters)\n\n        # --- Decoder path ---\n        # Tầng 4\n        x3_1_up = self.up4_t(x4_0)\n        # Cắt (crop) nếu kích thước không khớp (ít xảy ra với padding='same')\n        if x3_1_up.shape != x3_0.shape:\n             x3_1_up = TF.resize(x3_1_up, size=x3_0.shape[2:])\n        x3_1_cat = torch.cat([x3_0, x3_1_up], 1) # Concat (256 + 256)\n        x3_1 = self.conv3_1(x3_1_cat)\n        \n        # Tầng 3\n        x2_1_up = self.up3_t(x3_1)\n        if x2_1_up.shape != x2_0.shape:\n             x2_1_up = TF.resize(x2_1_up, size=x2_0.shape[2:])\n        x2_1_cat = torch.cat([x2_0, x2_1_up], 1) # Concat (128 + 128)\n        x2_1 = self.conv2_1(x2_1_cat)\n\n        # Tầng 2\n        x1_1_up = self.up2_t(x2_1)\n        if x1_1_up.shape != x1_0.shape:\n             x1_1_up = TF.resize(x1_1_up, size=x1_0.shape[2:])\n        x1_1_cat = torch.cat([x1_0, x1_1_up], 1) # Concat (64 + 64)\n        x1_1 = self.conv1_1(x1_1_cat)\n        \n        # Tầng 1\n        x0_1_up = self.up1_t(x1_1)\n        if x0_1_up.shape != x0_0.shape:\n             x0_1_up = TF.resize(x0_1_up, size=x0_0.shape[2:])\n        x0_1_cat = torch.cat([x0_0, x0_1_up], 1) # Concat (32 + 32)\n        x0_1 = self.conv0_1(x0_1_cat)\n        \n        # --- Output ---\n        output = self.final(x0_1)\n        return torch.sigmoid(output) # Sigmoid vì dùng BCELoss","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T16:51:20.009467Z","iopub.execute_input":"2025-11-03T16:51:20.009725Z","iopub.status.idle":"2025-11-03T16:51:20.022962Z","shell.execute_reply.started":"2025-11-03T16:51:20.009706Z","shell.execute_reply":"2025-11-03T16:51:20.022400Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 4. Định nghĩa Model UNet++","metadata":{}},{"cell_type":"code","source":"class ConvBlock(nn.Module):\n    \"\"\"(Conv2d -> BN -> ReLU) x 2\"\"\"\n    def __init__(self, in_channels, out_channels, mid_channels=None):\n        super().__init__()\n        if not mid_channels:\n            mid_channels = out_channels\n        self.double_conv = nn.Sequential(\n            nn.Conv2d(in_channels, mid_channels, kernel_size=3, padding=1, bias=False),\n            nn.BatchNorm2d(mid_channels),\n            nn.ReLU(inplace=True),\n            nn.Conv2d(mid_channels, out_channels, kernel_size=3, padding=1, bias=False),\n            nn.BatchNorm2d(out_channels),\n            nn.ReLU(inplace=True)\n        )\n\n    def forward(self, x):\n        return self.double_conv(x)\n\nclass UNetPlusPlus(nn.Module):\n    \"\"\"\n    Implementation of UNet++ (Nested U-Net)\n    Sử dụng filter count từ Table 2: [32, 64, 128, 256, 512]\n    \"\"\"\n    def __init__(self, in_channels=3, out_channels=1, \n                 filters=[32, 64, 128, 256, 512], deep_supervision=False):\n        super(UNetPlusPlus, self).__init__()\n        \n        self.deep_supervision = deep_supervision\n        nb_filter = filters\n\n        self.pool = nn.MaxPool2d(2, 2)\n        self.up = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True)\n\n        # Encoder path (X^i,0)\n        self.conv0_0 = ConvBlock(in_channels, nb_filter[0])\n        self.conv1_0 = ConvBlock(nb_filter[0], nb_filter[1])\n        self.conv2_0 = ConvBlock(nb_filter[1], nb_filter[2])\n        self.conv3_0 = ConvBlock(nb_filter[2], nb_filter[3])\n        self.conv4_0 = ConvBlock(nb_filter[3], nb_filter[4])\n\n        # Nested skip paths (X^i,j)\n        # j=1\n        self.conv0_1 = ConvBlock(nb_filter[0] + nb_filter[1], nb_filter[0])\n        self.conv1_1 = ConvBlock(nb_filter[1] + nb_filter[2], nb_filter[1])\n        self.conv2_1 = ConvBlock(nb_filter[2] + nb_filter[3], nb_filter[2])\n        self.conv3_1 = ConvBlock(nb_filter[3] + nb_filter[4], nb_filter[3])\n        \n        # j=2\n        self.conv0_2 = ConvBlock(nb_filter[0]*2 + nb_filter[1], nb_filter[0])\n        self.conv1_2 = ConvBlock(nb_filter[1]*2 + nb_filter[2], nb_filter[1])\n        self.conv2_2 = ConvBlock(nb_filter[2]*2 + nb_filter[3], nb_filter[2])\n        \n        # j=3\n        self.conv0_3 = ConvBlock(nb_filter[0]*3 + nb_filter[1], nb_filter[0])\n        self.conv1_3 = ConvBlock(nb_filter[1]*3 + nb_filter[2], nb_filter[1])\n\n        # j=4\n        self.conv0_4 = ConvBlock(nb_filter[0]*4 + nb_filter[1], nb_filter[0])\n\n        # Final output convs (cho deep supervision)\n        # Áp dụng sigmoid ở cuối cùng\n        self.final1 = nn.Conv2d(nb_filter[0], out_channels, kernel_size=1)\n        self.final2 = nn.Conv2d(nb_filter[0], out_channels, kernel_size=1)\n        self.final3 = nn.Conv2d(nb_filter[0], out_channels, kernel_size=1)\n        self.final4 = nn.Conv2d(nb_filter[0], out_channels, kernel_size=1)\n\n    def forward(self, x):\n        # Encoder (X^i,0)\n        x0_0 = self.conv0_0(x)\n        x1_0 = self.conv1_0(self.pool(x0_0))\n        x2_0 = self.conv2_0(self.pool(x1_0))\n        x3_0 = self.conv3_0(self.pool(x2_0))\n        x4_0 = self.conv4_0(self.pool(x3_0))\n\n        # Nested path j=1\n        x0_1 = self.conv0_1(torch.cat([x0_0, self.up(x1_0)], 1))\n        x1_1 = self.conv1_1(torch.cat([x1_0, self.up(x2_0)], 1))\n        x2_1 = self.conv2_1(torch.cat([x2_0, self.up(x3_0)], 1))\n        x3_1 = self.conv3_1(torch.cat([x3_0, self.up(x4_0)], 1))\n\n        # Nested path j=2\n        x0_2 = self.conv0_2(torch.cat([x0_0, x0_1, self.up(x1_1)], 1))\n        x1_2 = self.conv1_2(torch.cat([x1_0, x1_1, self.up(x2_1)], 1))\n        x2_2 = self.conv2_2(torch.cat([x2_0, x2_1, self.up(x3_1)], 1))\n\n        # Nested path j=3\n        x0_3 = self.conv0_3(torch.cat([x0_0, x0_1, x0_2, self.up(x1_2)], 1))\n        x1_3 = self.conv1_3(torch.cat([x1_0, x1_1, x1_2, self.up(x2_2)], 1))\n        \n        # Nested path j=4\n        x0_4 = self.conv0_4(torch.cat([x0_0, x0_1, x0_2, x0_3, self.up(x1_3)], 1))\n\n        # Áp dụng sigmoid cho các output\n        output1 = torch.sigmoid(self.final1(x0_1))\n        output2 = torch.sigmoid(self.final2(x0_2))\n        output3 = torch.sigmoid(self.final3(x0_3))\n        output4 = torch.sigmoid(self.final4(x0_4))\n\n        if self.deep_supervision:\n            return [output1, output2, output3, output4]\n        else:\n            return output4","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T16:51:20.023803Z","iopub.execute_input":"2025-11-03T16:51:20.024054Z","iopub.status.idle":"2025-11-03T16:51:20.044331Z","shell.execute_reply.started":"2025-11-03T16:51:20.024038Z","shell.execute_reply":"2025-11-03T16:51:20.043718Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 5. Dataset, Loss và Metrics","metadata":{}},{"cell_type":"code","source":"class DSB2018Dataset(Dataset):\n    def __init__(self, X_data_list, Y_data_list, transform=None):\n        self.transform = transform\n        # Nhận dữ liệu đã được pre-load\n        self.X_train_list = X_data_list\n        self.Y_train_list = Y_data_list\n\n    def __len__(self):\n        return len(self.X_train_list)\n\n    def __getitem__(self, index):\n        image = self.X_train_list[index]\n        mask = self.Y_train_list[index]\n        \n        # Chuẩn hóa ảnh (về [0, 1])\n        image = image.astype(np.float32) / 255.0\n        # Chuyển mask (về [0, 1])\n        mask = (mask > 0).astype(np.float32)\n        \n        # Chuyển sang tensor (C, H, W)\n        image = torch.from_numpy(image).permute(2, 0, 1)\n        mask = torch.from_numpy(mask).permute(2, 0, 1)\n\n        if self.transform:\n            # Bạn có thể thêm augmentation ở đây\n            pass\n            \n        return image, mask\n\n# --- LOSS FUNCTIONS --- (Không thay đổi)\nclass DiceLoss(nn.Module):\n    # ... (giữ nguyên code)\n    def __init__(self, smooth=1e-6):\n        super(DiceLoss, self).__init__()\n        self.smooth = smooth\n    def forward(self, inputs, targets):\n        inputs = inputs.view(-1)\n        targets = targets.view(-1)\n        intersection = (inputs * targets).sum()\n        dice = (2. * intersection + self.smooth) / (inputs.sum() + targets.sum() + self.smooth)\n        return 1 - dice\n\nclass BCEDiceLoss(nn.Module):\n    # ... (giữ nguyên code)\n    def __init__(self, bce_weight=0.5, dice_weight=0.5):\n        super(BCEDiceLoss, self).__init__()\n        self.bce_weight = bce_weight\n        self.dice_weight = dice_weight\n        self.bce_loss = nn.BCELoss()\n        self.dice_loss = DiceLoss()\n    def forward(self, inputs, targets):\n        bce = self.bce_loss(inputs, targets)\n        dice = self.dice_loss(inputs, targets)\n        return self.bce_weight * bce + self.dice_weight * dice\n\nclass DeepSupervisionLoss(nn.Module):\n    # ... (giữ nguyên code)\n    def __init__(self, bce_weight=0.5, dice_weight=0.5, weights=[0.25, 0.25, 0.25, 0.25]):\n        super(DeepSupervisionLoss, self).__init__()\n        self.base_loss = BCEDiceLoss(bce_weight, dice_weight)\n        self.weights = weights\n    def forward(self, inputs: list, target):\n        total_loss = 0\n        for i, pred in enumerate(inputs):\n            if pred.shape != target.shape:\n                 target_resized = F.interpolate(target, size=pred.shape[2:], mode='bilinear', align_corners=False)\n            else:\n                 target_resized = target\n            total_loss += self.weights[i] * self.base_loss(pred, target_resized)\n        return total_loss\n\n# --- METRIC --- (Không thay đổi)\ndef iou_metric(pred, target, threshold=0.5, smooth=1e-6):\n    # ... (giữ nguyên code)\n    pred = (pred > threshold).float()\n    target = (target > 0.5).float()\n    intersection = (pred * target).sum(dim=(1, 2, 3))\n    union = pred.sum(dim=(1, 2, 3)) + target.sum(dim=(1, 2, 3)) - intersection\n    iou = (intersection + smooth) / (union + smooth)\n    return iou.mean()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T16:51:20.046607Z","iopub.execute_input":"2025-11-03T16:51:20.046886Z","iopub.status.idle":"2025-11-03T16:51:20.063308Z","shell.execute_reply.started":"2025-11-03T16:51:20.046871Z","shell.execute_reply":"2025-11-03T16:51:20.062702Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 6. Hàm Training và Validation","metadata":{}},{"cell_type":"code","source":"def train_fn(loader, model, optimizer, loss_fn, device, deep_supervision):\n    model.train()\n    loop = tqdm(loader, desc=\"Training\")\n    \n    total_loss = 0\n    total_iou = 0\n    \n    for batch_idx, (data, targets) in enumerate(loop):\n        data = data.to(device=device)\n        targets = targets.to(device=device)\n\n        predictions = model(data)\n        \n        if deep_supervision:\n            loss = loss_fn(predictions, targets)\n            iou_preds = predictions[-1] \n        else:\n            loss = loss_fn(predictions, targets)\n            iou_preds = predictions\n\n        optimizer.zero_grad()\n        loss.backward()\n        optimizer.step()\n\n        iou = iou_metric(iou_preds, targets)\n        total_loss += loss.item()\n        total_iou += iou.item()\n        loop.set_postfix(loss=loss.item(), iou=iou.item())\n        \n    avg_loss = total_loss / len(loader)\n    avg_iou = total_iou / len(loader)\n    print(f\"Epoch Train Loss: {avg_loss:.4f}, Epoch Train IoU: {avg_iou:.4f}\")\n    # --- THÊM DÒNG NÀY ---\n    return {\"loss\": avg_loss, \"iou\": avg_iou}\n\ndef val_fn(loader, model, loss_fn, device, deep_supervision):\n    model.eval()\n    loop = tqdm(loader, desc=\"Validation\")\n    val_loss = 0\n    val_iou = 0\n    \n    with torch.no_grad():\n        for batch_idx, (data, targets) in enumerate(loop):\n            data = data.to(device=device)\n            targets = targets.to(device=device)\n            predictions = model(data)\n            \n            if deep_supervision:\n                loss = loss_fn(predictions, targets)\n                iou_preds = predictions[-1]\n            else:\n                loss = loss_fn(predictions, targets)\n                iou_preds = predictions\n\n            val_loss += loss.item()\n            val_iou += iou_metric(iou_preds, targets).item()\n\n    avg_loss = val_loss / len(loader)\n    avg_iou = val_iou / len(loader)\n    print(f\"==> Avg Val Loss: {avg_loss:.4f}, Avg Val IoU: {avg_iou:.4f}\")\n    # --- THÊM DÒNG NÀY ---\n    return {\"loss\": avg_loss, \"iou\": avg_iou}\n\n# [HÀM MỚI] Dùng để đánh giá trên tập test cuối cùng\ndef evaluate_fn(loader, model, loss_fn, device, deep_supervision):\n    model.eval()\n    loop = tqdm(loader, desc=\"Evaluating Test Set\")\n    eval_loss = 0\n    eval_iou = 0\n    \n    with torch.no_grad():\n        for data, targets in loop:\n            data, targets = data.to(device), targets.to(device)\n            predictions = model(data)\n            \n            if deep_supervision:\n                loss = loss_fn(predictions, targets)\n                iou_preds = predictions[-1]\n            else:\n                loss = loss_fn(predictions, targets)\n                iou_preds = predictions\n                \n            eval_loss += loss.item()\n            eval_iou += iou_metric(iou_preds, targets).item()\n            \n    avg_loss = eval_loss / len(loader)\n    avg_iou = eval_iou / len(loader)\n    return avg_loss, avg_iou","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T16:51:20.063914Z","iopub.execute_input":"2025-11-03T16:51:20.064203Z","iopub.status.idle":"2025-11-03T16:51:20.080888Z","shell.execute_reply.started":"2025-11-03T16:51:20.064185Z","shell.execute_reply":"2025-11-03T16:51:20.080326Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 7. Khởi tạo và Huấn luyện","metadata":{}},{"cell_type":"code","source":"print(\"--- 3. Initializing Dataset & Splitting ---\")\n\n# (Code tải và chia data... giữ nguyên)\n# ...\nfull_dataset = DSB2018Dataset(X_data_list=X_data, Y_data_list=Y_data)\ntotal_size = len(full_dataset)\nval_size = int(total_size * VALIDATION_SPLIT)\ntest_size = int(total_size * TEST_SPLIT)\ntrain_size = total_size - val_size - test_size\nprint(f\"Splitting data: {train_size} Train, {val_size} Validation, {test_size} Test\")\ntrain_dataset, val_dataset, test_dataset = random_split(\n    full_dataset, \n    [train_size, val_size, test_size],\n    generator=torch.Generator().manual_seed(SEED)\n)\ntrain_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\nval_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)\ntest_loader = DataLoader(test_dataset, batch_size=BATCH_SIZE, shuffle=False)\nprint(\"DataLoaders created.\")\n\n# --- Nơi lưu kết quả test cuối cùng ---\nfinal_test_results = {}\n# --- MỚI: Nơi lưu lịch sử training ---\ntraining_histories = {\n    \"unet\": {\"train_loss\": [], \"train_iou\": [], \"val_loss\": [], \"val_iou\": []},\n    \"unet_pp\": {\"train_loss\": [], \"train_iou\": [], \"val_loss\": [], \"val_iou\": []}\n}\n\n# ======================================================================\n# --- BƯỚC 4A: HUẤN LUYỆN VÀ ĐÁNH GIÁ UNET ---\n# ======================================================================\nprint(f\"\\n\" + \"=\"*50)\nprint(f\"--- 4. STARTING UNET TRAINING ---\")\nprint(\"=\"*50)\n\nmodel_unet = UNet(in_channels=IMG_CHANNELS, out_channels=1).to(DEVICE)\nloss_fn_unet = BCEDiceLoss()\noptimizer_unet = optim.Adam(model_unet.parameters(), lr=LEARNING_RATE)\nbest_val_iou = -1.0\nearly_stop_counter = 0\n\nfor epoch in range(EPOCHS):\n    print(f\"\\nEpoch {epoch+1}/{EPOCHS}\")\n    \n    train_metrics = train_fn(train_loader, model_unet, optimizer_unet, loss_fn_unet, DEVICE, deep_supervision=False)\n    val_metrics = val_fn(val_loader, model_unet, loss_fn_unet, DEVICE, deep_supervision=False)\n    \n    # --- MỚI: Lưu lịch sử ---\n    training_histories[\"unet\"][\"train_loss\"].append(train_metrics[\"loss\"])\n    training_histories[\"unet\"][\"train_iou\"].append(train_metrics[\"iou\"])\n    training_histories[\"unet\"][\"val_loss\"].append(val_metrics[\"loss\"])\n    training_histories[\"unet\"][\"val_iou\"].append(val_metrics[\"iou\"])\n    # -----------------------\n\n    current_val_iou = val_metrics[\"iou\"]\n    if current_val_iou > best_val_iou:\n        best_val_iou = current_val_iou\n        print(f\"🎉 New best UNet model! Val IoU: {best_val_iou:.4f}. Saving model to {UNET_SAVE_PATH}\")\n        torch.save(model_unet.state_dict(), UNET_SAVE_PATH)\n        early_stop_counter = 0\n    else:\n        early_stop_counter += 1\n        print(f\"Val IoU did not improve. Counter: {early_stop_counter}/{EARLY_STOP_PATIENCE}\")\n\n    if early_stop_counter >= EARLY_STOP_PATIENCE:\n        print(f\"--- UNet Early stopping triggered after {epoch+1} epochs ---\")\n        break\n\nprint(\"\\n--- UNet Training Finished ---\")\nprint(f\"\\n--- Evaluating UNet on Test Set ---\")\nmodel_unet.load_state_dict(torch.load(UNET_SAVE_PATH))\ntest_loss, test_iou = evaluate_fn(test_loader, model_unet, loss_fn_unet, DEVICE, deep_supervision=False)\nfinal_test_results[\"UNet\"] = {\"loss\": test_loss, \"iou\": test_iou}\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T16:51:20.081612Z","iopub.execute_input":"2025-11-03T16:51:20.081856Z","iopub.status.idle":"2025-11-03T16:53:05.282399Z","shell.execute_reply.started":"2025-11-03T16:51:20.081841Z","shell.execute_reply":"2025-11-03T16:53:05.281728Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======================================================================\n# --- BƯỚC 4B: HUẤN LUYỆN VÀ ĐÁNH GIÁ UNET++ ---\n# ======================================================================\nprint(f\"\\n\" + \"=\"*50)\nprint(f\"--- 5. STARTING UNET++ TRAINING ---\")\nprint(\"=\"*50)\n\nUSE_DEEP_SUPERVISION = True\nmodel_unet_pp = UNetPlusPlus(in_channels=IMG_CHANNELS, out_channels=1, \n                             deep_supervision=USE_DEEP_SUPERVISION).to(DEVICE)\nif USE_DEEP_SUPERVISION:\n    print(\"Using Deep Supervision Loss for UNet++.\")\n    loss_fn_unet_pp = DeepSupervisionLoss()\nelse:\n    print(\"Using standard BCEDice Loss for UNet++.\")\n    loss_fn_unet_pp = BCEDiceLoss()\noptimizer_unet_pp = optim.Adam(model_unet_pp.parameters(), lr=LEARNING_RATE)\nbest_val_iou = -1.0\nearly_stop_counter = 0\n\nfor epoch in range(EPOCHS):\n    print(f\"\\nEpoch {epoch+1}/{EPOCHS}\")\n    \n    train_metrics = train_fn(train_loader, model_unet_pp, optimizer_unet_pp, loss_fn_unet_pp, DEVICE, USE_DEEP_SUPERVISION)\n    val_metrics = val_fn(val_loader, model_unet_pp, loss_fn_unet_pp, DEVICE, USE_DEEP_SUPERVISION)\n    \n    # --- MỚI: Lưu lịch sử ---\n    training_histories[\"unet_pp\"][\"train_loss\"].append(train_metrics[\"loss\"])\n    training_histories[\"unet_pp\"][\"train_iou\"].append(train_metrics[\"iou\"])\n    training_histories[\"unet_pp\"][\"val_loss\"].append(val_metrics[\"loss\"])\n    training_histories[\"unet_pp\"][\"val_iou\"].append(val_metrics[\"iou\"])\n    # -----------------------\n\n    current_val_iou = val_metrics[\"iou\"]\n    if current_val_iou > best_val_iou:\n        best_val_iou = current_val_iou\n        print(f\"🎉 New best UNet++ model! Val IoU: {best_val_iou:.4f}. Saving model to {UNET_PLUS_PLUS_SAVE_PATH}\")\n        torch.save(model_unet_pp.state_dict(), UNET_PLUS_PLUS_SAVE_PATH)\n        early_stop_counter = 0\n    else:\n        early_stop_counter += 1\n        print(f\"Val IoU did not improve. Counter: {early_stop_counter}/{EARLY_STOP_PATIENCE}\")\n\n    if early_stop_counter >= EARLY_STOP_PATIENCE:\n        print(f\"--- UNet++ Early stopping triggered after {epoch+1} epochs ---\")\n        break\n\nprint(\"\\n--- UNet++ Training Finished ---\")\nprint(f\"\\n--- Evaluating UNet++ on Test Set ---\")\nmodel_unet_pp.load_state_dict(torch.load(UNET_PLUS_PLUS_SAVE_PATH))\ntest_loss, test_iou = evaluate_fn(test_loader, model_unet_pp, loss_fn_unet_pp, DEVICE, USE_DEEP_SUPERVISION)\nfinal_test_results[\"UNet++\"] = {\"loss\": test_loss, \"iou\": test_iou}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T16:53:05.283169Z","iopub.execute_input":"2025-11-03T16:53:05.283437Z","iopub.status.idle":"2025-11-03T16:58:50.916096Z","shell.execute_reply.started":"2025-11-03T16:53:05.283420Z","shell.execute_reply":"2025-11-03T16:58:50.915424Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ======================================================================\n# --- BƯỚC 6: TỔNG KẾT ---\n# ======================================================================\nprint(f\"\\n\" + \"=\"*50)\nprint(f\"== FINAL TEST SET RESULTS ==\")\nprint(f\"UNet:     Test Loss: {final_test_results['UNet']['loss']:.4f}, Test IoU: {final_test_results['UNet']['iou']:.4f}\")\nprint(f\"UNet++:   Test Loss: {final_test_results['UNet++']['loss']:.4f}, Test IoU: {final_test_results['UNet++']['iou']:.4f}\")\nprint(f\"================================\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T16:58:50.916794Z","iopub.execute_input":"2025-11-03T16:58:50.917003Z","iopub.status.idle":"2025-11-03T16:58:50.922167Z","shell.execute_reply.started":"2025-11-03T16:58:50.916970Z","shell.execute_reply":"2025-11-03T16:58:50.921346Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 8. Visualize Lịch sử Training","metadata":{}},{"cell_type":"code","source":"print(\"--- 8. Visualizing Training History ---\")\n\n# Lấy dữ liệu từ dictionary\nunet_hist = training_histories[\"unet\"]\nunet_pp_hist = training_histories[\"unet_pp\"]\n\n# --- 1. Biểu đồ Loss ---\nplt.figure(figsize=(14, 6))\nplt.subplot(1, 2, 1)\n# UNet\nplt.plot(unet_hist[\"train_loss\"], label='UNet - Train Loss', color='blue', linestyle='--')\nplt.plot(unet_hist[\"val_loss\"], label='UNet - Val Loss', color='blue')\n# UNet++\nplt.plot(unet_pp_hist[\"train_loss\"], label='UNet++ - Train Loss', color='green', linestyle='--')\nplt.plot(unet_pp_hist[\"val_loss\"], label='UNet++ - Val Loss', color='green')\n\nplt.title('Model Loss History')\nplt.xlabel('Epochs')\nplt.ylabel('Loss (BCE + Dice)')\nplt.legend()\nplt.grid(True, linestyle='--', alpha=0.5)\n\n# --- 2. Biểu đồ IoU ---\nplt.subplot(1, 2, 2)\n# UNet\nplt.plot(unet_hist[\"train_iou\"], label='UNet - Train IoU', color='blue', linestyle='--')\nplt.plot(unet_hist[\"val_iou\"], label='UNet - Val IoU', color='blue')\n# UNet++\nplt.plot(unet_pp_hist[\"train_iou\"], label='UNet++ - Train IoU', color='green', linestyle='--')\nplt.plot(unet_pp_hist[\"val_iou\"], label='UNet++ - Val IoU', color='green')\n\nplt.title('Model IoU History')\nplt.xlabel('Epochs')\nplt.ylabel('IoU (Metric)')\nplt.legend()\nplt.grid(True, linestyle='--', alpha=0.5)\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T16:58:50.922961Z","iopub.execute_input":"2025-11-03T16:58:50.923217Z","iopub.status.idle":"2025-11-03T16:58:51.344140Z","shell.execute_reply.started":"2025-11-03T16:58:50.923195Z","shell.execute_reply":"2025-11-03T16:58:51.343409Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 9. Trực quan hóa kết quả","metadata":{}},{"cell_type":"code","source":"print(\"--- 7. Visualizing Results on TEST Set (Comparing UNet vs UNet++) ---\")\n\n# --- 1. Tải cả hai model tốt nhất ---\nprint(\"Loading best saved models...\")\n# Tải UNet\nmodel_unet = UNet(in_channels=IMG_CHANNELS, out_channels=1).to(DEVICE)\nmodel_unet.load_state_dict(torch.load(UNET_SAVE_PATH))\nmodel_unet.eval()\n\n# Tải UNet++\n# (Giả sử model UNet++ tốt nhất đã được lưu với deep_supervision=True)\nDEEP_SUPERVISION_FOR_LOADING = True \nmodel_unet_pp = UNetPlusPlus(in_channels=IMG_CHANNELS, out_channels=1, \n                             deep_supervision=DEEP_SUPERVISION_FOR_LOADING).to(DEVICE)\nmodel_unet_pp.load_state_dict(torch.load(UNET_PLUS_PLUS_SAVE_PATH))\nmodel_unet_pp.eval()\n\nprint(\"Models loaded. Visualizing samples...\")\n\n# --- 2. Lấy mẫu và dự đoán ---\nnum_samples = 5\nindices = random.sample(range(len(test_dataset)), num_samples)\n\nwith torch.no_grad():\n    for i, idx in enumerate(indices):\n        # Lấy ảnh và mask từ test_dataset\n        img, mask = test_dataset[idx]\n        \n        # Thêm chiều batch (1, C, H, W) và đưa lên GPU\n        img_tensor = img.unsqueeze(0).to(DEVICE)\n        \n        # --- Dự đoán ---\n        # UNet\n        pred_unet = model_unet(img_tensor)\n        pred_unet_np = (pred_unet.squeeze().cpu().numpy() > 0.5).astype(float)\n        \n        # UNet++\n        pred_unet_pp = model_unet_pp(img_tensor)\n        if DEEP_SUPERVISION_FOR_LOADING:\n            pred_unet_pp = pred_unet_pp[-1] # Lấy output cuối cùng\n        pred_unet_pp_np = (pred_unet_pp.squeeze().cpu().numpy() > 0.5).astype(float)\n        \n        # --- Chuyển đổi để vẽ ---\n        img_np = img.permute(1, 2, 0).cpu().numpy()\n        mask_np = mask.squeeze().cpu().numpy()\n        \n        # --- Vẽ 4 ảnh ---\n        plt.figure(figsize=(20, 5))\n        \n        plt.subplot(1, 4, 1)\n        plt.title(\"Image (Test)\")\n        plt.imshow(img_np)\n        plt.axis('off')\n        \n        plt.subplot(1, 4, 2)\n        plt.title(\"Ground Truth Mask\")\n        plt.imshow(mask_np, cmap='gray')\n        plt.axis('off')\n\n        plt.subplot(1, 4, 3)\n        plt.title(\"UNet Prediction\")\n        plt.imshow(pred_unet_np, cmap='gray')\n        plt.axis('off')\n        \n        plt.subplot(1, 4, 4)\n        plt.title(\"UNet++ Prediction\")\n        plt.imshow(pred_unet_pp_np, cmap='gray')\n        plt.axis('off')\n        \n        plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-03T16:58:51.344950Z","iopub.execute_input":"2025-11-03T16:58:51.345176Z","iopub.status.idle":"2025-11-03T16:58:53.174415Z","shell.execute_reply.started":"2025-11-03T16:58:51.345160Z","shell.execute_reply":"2025-11-03T16:58:53.173716Z"}},"outputs":[],"execution_count":null}]}