{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":71549,"databundleVersionId":8561470,"sourceType":"competition"}],"dockerImageVersionId":30776,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport sys\nimport glob\nimport re\nimport copy\nimport numpy as np\nimport pandas as pd\nimport logging\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nimport torch.optim.lr_scheduler as lr_scheduler\nfrom torch import Tensor\nfrom pathlib import Path\nfrom torch.utils.data import DataLoader\nfrom torch import Tensor\nfrom torch.nn.modules.loss import _Loss as BaseNNLossModule\nfrom typing import Callable, List, Tuple, Dict, Union, Optional, Literal\nfrom tqdm import tqdm\nfrom math import floor, ceil\nimport pydicom\nfrom pydicom.pixel_data_handlers.util import apply_voi_lut\nfrom torch.utils.data import Dataset, random_split\nfrom torchvision import transforms\nfrom torchvision.models.video import (\n    r3d_18, R3D_18_Weights,\n    s3d, S3D_Weights,\n    mc3_18, MC3_18_Weights\n)","metadata":{"execution":{"iopub.status.busy":"2025-01-15T21:13:26.427828Z","iopub.execute_input":"2025-01-15T21:13:26.428062Z","iopub.status.idle":"2025-01-15T21:14:49.040366Z","shell.execute_reply.started":"2025-01-15T21:13:26.428036Z","shell.execute_reply":"2025-01-15T21:14:49.039366Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ============================ Configuration Paths =========================== #\nPROJECT_PATH = Path(\"/kaggle/working\")\nDATA_PATH = Path(\"/kaggle/input/rsna-2024-lumbar-spine-degenerative-classification\")\nTRAIN_IMAGES_PATH = DATA_PATH / \"train_images\"\nTEST_IMAGES_PATH = DATA_PATH / \"test_images\"\nTRAIN_CSV_PATH = DATA_PATH / \"train.csv\"\nTEST_CSV_PATH = DATA_PATH / \"test.csv\"\nMODELS_PATH = PROJECT_PATH / \"saved_models\"\nSUBMISSION_PATH = PROJECT_PATH\n# ======================== End Of Configuration Paths ======================== #\n\n\n# =============================== Encoding Maps ============================== #\nSERIES_S2I = {\"Sagittal T1\": 0, \"Sagittal T2/STIR\": 0, \"Axial T2\": 1}\n# =========================== End Of Encoding Maps =========================== #\n\n# Set the root directory of the project.\nif str(PROJECT_PATH) not in sys.path:\n    sys.path.append(str(PROJECT_PATH))\n","metadata":{"execution":{"iopub.status.busy":"2025-01-15T21:14:49.043045Z","iopub.execute_input":"2025-01-15T21:14:49.043851Z","iopub.status.idle":"2025-01-15T21:14:49.048883Z","shell.execute_reply.started":"2025-01-15T21:14:49.043810Z","shell.execute_reply":"2025-01-15T21:14:49.047974Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def add_channel_dim(x: torch.Tensor) -> torch.Tensor:\n    return x.unsqueeze(1)\n\n\nclass BaseLumbarSpineDataset(Dataset):\n\n    conditions_i2s: List[str] = [\n        \"spinal_canal_stenosis_l1_l2\", \"spinal_canal_stenosis_l2_l3\", \"spinal_canal_stenosis_l3_l4\",\n        \"spinal_canal_stenosis_l4_l5\", \"spinal_canal_stenosis_l5_s1\", \"left_neural_foraminal_narrowing_l1_l2\",\n        \"left_neural_foraminal_narrowing_l2_l3\", \"left_neural_foraminal_narrowing_l3_l4\",\n        \"left_neural_foraminal_narrowing_l4_l5\", \"left_neural_foraminal_narrowing_l5_s1\",\n        \"right_neural_foraminal_narrowing_l1_l2\", \"right_neural_foraminal_narrowing_l2_l3\",\n        \"right_neural_foraminal_narrowing_l3_l4\", \"right_neural_foraminal_narrowing_l4_l5\",\n        \"right_neural_foraminal_narrowing_l5_s1\", \"left_subarticular_stenosis_l1_l2\",\n        \"left_subarticular_stenosis_l2_l3\", \"left_subarticular_stenosis_l3_l4\",\n        \"left_subarticular_stenosis_l4_l5\", \"left_subarticular_stenosis_l5_s1\",\n        \"right_subarticular_stenosis_l1_l2\", \"right_subarticular_stenosis_l2_l3\",\n        \"right_subarticular_stenosis_l3_l4\", \"right_subarticular_stenosis_l4_l5\",\n        \"right_subarticular_stenosis_l5_s1\"\n    ]  # Acts as enumeration for the conditions.\n    severity_i2s: List[str] = [\"Normal/Mild\", \"Moderate\", \"Severe\"]  # Acts as enumeration for the severity.\n\n    conditions_s2i: Dict[str, int] = {condition: i for i, condition in enumerate(conditions_i2s)}\n    severity_s2i: Dict[Union[str, float], int] = {severity: i for i, severity in enumerate(severity_i2s)}\n    nan_class = severity_s2i[\"Normal/Mild\"]  # Will be 0.\n\n    default_preprocess = transforms.Compose([\n        # Accepts a numpy array of dimensions (D, H, W).\n        torch.from_numpy,\n        transforms.ConvertImageDtype(torch.float32),\n        # Add a channel dimension, (D, 1, H, W).\n        add_channel_dim,\n    ])\n\n    def __init__(\n        self,\n        train: bool,\n        *,\n        preprocess: Optional[transforms.Compose] = None,\n        augs: Optional[transforms.Compose] = None\n    ) -> None:\n        super().__init__()\n        self.train = train\n        self.preprocess: transforms.Compose = preprocess or self.default_preprocess\n        self.augs = augs\n        self._apply_augs = augs is not None\n        self.series_description = pd.read_csv(\n            DATA_PATH / f\"{'train' if train else 'test'}_series_descriptions.csv\"\n        )\n        self.images_path: Path = TRAIN_IMAGES_PATH if train else TEST_IMAGES_PATH\n        if train:\n            self.train_csv = pd.read_csv(TRAIN_CSV_PATH, index_col=\"study_id\")\n\n    def split(self, val_size: float = 0.2) -> tuple[\"BaseLumbarSpineDataset\", \"BaseLumbarSpineDataset\"]:\n        train, val = random_split(self, [1 - val_size, val_size])  # type: ignore\n        val._apply_augs = False\n        return train, val","metadata":{"execution":{"iopub.status.busy":"2025-01-15T21:14:49.050063Z","iopub.execute_input":"2025-01-15T21:14:49.050343Z","iopub.status.idle":"2025-01-15T21:14:49.064398Z","shell.execute_reply.started":"2025-01-15T21:14:49.050316Z","shell.execute_reply":"2025-01-15T21:14:49.063773Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for handler in logging.root.handlers[:]:\n    logging.root.removeHandler(handler)\nlogging.basicConfig(stream=sys.stdout, level=logging.DEBUG, format=\"%(asctime)s - %(name)s.%(levelname)s: %(message)s\")\nlogger = logging.getLogger(\"kaggle.submission\")\n\n# ============================== BaseModel Class ============================= #\nclass BaseModel(nn.Module):\n    \"\"\"Base model class.\"\"\"\n    def __init__(self, name: str = \"\") -> None:\n        \"\"\"Constructor.\"\"\"\n        if not isinstance(name, str):\n            raise TypeError(\"Model name must be a string.\")\n        elif name == \"\":\n            raise ValueError(\"Model name must be provided.\")\n        elif \"=\" in name:\n            raise ValueError(\"Model name cannot contain the '=' character.\")\n        elif name == \"do_not_save\":\n            pass\n        else:\n            self.make_model_dir(name, num_tries=5)\n        super().__init__()\n        self.best_weights: Optional[dict[str, Tensor]] = None\n        self.global_epoch: int = 0\n        \n        self.train_costs: list[float] = []\n        self.val_costs: list[float] = []\n        self.train_accs: list[float] = []\n        self.val_accs: list[float] = []\n\n        self.criterion: BaseNNLossModule = nn.NLLLoss()\n        self.name: str = name\n\n    def forward(self, *args, **kwargs) -> Tensor:\n        \"\"\"Forward pass.\"\"\"\n        raise NotImplementedError\n\n    def fit(self, train_loader: DataLoader, val_loader: Optional[DataLoader] = None, \n            num_epochs: int = 30, lr: float = 0.001, momentum: float = 0.9, wd: float = 0.,\n            try_cuda: bool = True, verbose: bool = True, print_stride: int = 1) -> None:\n        \"\"\"\n        Base function for training a model.\n\n        Args:\n            train_loader (DataLoader) - The dataloader to fit the model to.\n            val_loader (DataLoader) - The dataloader to validate the model on.\n            num_epochs (int) - Number of epochs.\n            lr (float) - Learning rate.\n            momentum (float) - Momentum for SGD.\n            wd (float) - Weight decay.\n            try_cuda (bool) - Try to use CUDA.\n            verbose (bool) - Verbose flag.\n            print_stride (int) - Print stride (in epochs).\n        \"\"\"\n        use_cuda = try_cuda and torch.cuda.is_available()\n        if use_cuda:\n            self.cuda()\n            print(\"Using CUDA for training.\")\n        else:\n            self.cpu()\n            print(\"Using CPU for training.\")\n\n        # Create the optimizer.\n        optimizer = optim.Adam(self.parameters(), lr=lr, weight_decay=wd)\n\n        start_epoch = self.global_epoch\n        for epoch in range(num_epochs):\n            running_loss: float = 0.\n            total_corrects: int = 0\n            self.global_epoch += 1\n            for mb, (x, y) in enumerate(train_loader):\n                if use_cuda:\n                    x: Tensor = x.cuda(); y: Tensor = y.cuda()\n\n                y_hat, loss = self._train_step(x, y, optimizer)\n\n                running_loss, total_corrects = self._calc_running_metrics(\n                    x, y_hat, y, loss, running_loss, total_corrects)\n\n                # TODO: Use tqdm instead of print.\n                if verbose:\n                    print(f\"\\r[epoch: {self.global_epoch:02d}/{start_epoch + num_epochs:02d} \"\n                          f\"mb: {mb + 1:03d}/{len(train_loader):03d}]  \" \n                          f\"[Train loss: {loss:.6f}]\\033[J\", end=\"\")\n\n            train_epoch_loss = running_loss / len(train_loader.dataset)  # type: ignore\n            train_total_acc = total_corrects / len(train_loader.dataset)  # type: ignore\n            self.train_costs.append(train_epoch_loss)\n            self.train_accs.append(train_total_acc)\n            \n            if val_loader is not None:\n                val_epoch_loss, val_total_acc = self.calc_metrics(val_loader, use_cuda)\n            else:\n                val_epoch_loss, val_total_acc = torch.nan, torch.nan\n\n            self.val_costs.append(val_epoch_loss)\n            self.val_accs.append(val_total_acc)\n\n            self.save_best_weights()\n\n            if verbose and (epoch % print_stride == 0 or epoch == num_epochs - 1):\n                print(f\"\\r[epoch: {self.global_epoch:02d}/{start_epoch + num_epochs:02d}] \"\n                      f\"[Train loss: {train_epoch_loss:.6f} \"\n                      f\" Train acc: {100 * train_total_acc:.6f}%]  \"\n                      f\"[Val loss: {val_epoch_loss:.6f} \"\n                      f\" Val acc: {100 * val_total_acc:.6f}%]\")\n        self.cpu()\n\n    def _train_step(self, x: Tensor, y: Tensor, optimizer: optim.Optimizer) -> Tuple[Tensor, float]:\n        \"\"\"\n        Performs a single training step.\n\n        Args:\n            x (Tensor) - Input tensor.\n            y (Tensor) - Target tensor.\n            optimizer (Optimizer) - Optimizer.\n        \n        Returns:\n            Tuple[Tensor, float] - The model's output and the loss of the model.\n        \"\"\"\n        # zero the parameter gradients.\n        optimizer.zero_grad()\n\n        # forward + backward + optimize.\n        y_hat: Tensor = self(x)\n\n        loss: Tensor = self.criterion(y_hat.view(-1, y_hat.size(-1)), y.view(-1))\n        loss.backward()\n        optimizer.step()\n\n        lloss = loss.item()\n\n        return y_hat, lloss\n\n    def calc_metrics(self, data_loader: DataLoader, use_cuda: bool) -> Tuple[float, float]:\n        \"\"\"\n        Calculates and returns the loss and the accuracy of the model on a given dataset.\n\n        Args:\n            data_loader (DataLoader) - Data loader.\n            use_cuda (bool) - Use CUDA flag.\n\n        Returns:\n            Tuple[float, float] - The loss and score of the model on the given dataset.            \n        \"\"\"\n        running_loss: float = 0.\n        total_corrects: int = 0\n        self.eval()\n        with torch.no_grad():\n            for x, y in tqdm(data_loader, desc=\"Validation\", total=len(data_loader)):\n                if use_cuda:\n                    x: Tensor = x.cuda(); y: Tensor = y.cuda()\n\n                y_hat: Tensor = self(x)\n                loss: float = self.criterion(y_hat.view(-1, y_hat.size(-1)), y.view(-1)).item()\n\n                running_loss, total_corrects = self._calc_running_metrics(\n                    x, y_hat, y, loss, running_loss, total_corrects)\n\n        self.train()\n        total_loss = running_loss / len(data_loader.dataset)  # type: ignore\n        total_acc = total_corrects / len(data_loader.dataset)  # type: ignore\n        return total_loss, total_acc\n\n    def _calc_running_metrics(self, x: Tensor, y_hat: Tensor, y: Tensor, loss: float,\n                              running_loss: float, total_corrects: int) -> Tuple[float, int]:\n        corrects = int((y_hat.argmax(-1) == y).sum().item())\n\n        running_loss += loss * x.size(0)\n        total_corrects += corrects\n        return running_loss, total_corrects\n\n    def save_best_weights(self) -> None:\n        \"\"\"Saves the best weights of the model.\"\"\"\n        if self.best_weights is None or self.val_costs[-1] < min(self.val_costs[:-1]):\n            self.best_weights = copy.deepcopy(self.state_dict())\n            self.best_val_cost = self.val_costs[-1]\n            self.best_epoch = len(self.val_costs)\n            self._cleanup_old_models()\n            logger.info(f\"Saving model weights at epoch: {self.global_epoch}\")\n            torch.save(self.best_weights, self.model_dir / f\"{self.name}_e={self.global_epoch}.pt\")\n\n    def load_best_weights(self):\n        \"\"\"Loads the best weights of the model.\"\"\"\n        assert self.best_weights is not None, \"No weights to load.\"\n        logger.info(f\"Loading best model weights at epoch: {self.best_epoch} and val cost: {self.best_val_cost}\")\n        self.load_state_dict(self.best_weights)\n\n    def _cleanup_old_models(self):\n        \"\"\"Removes old model files, keeping only the last 5.\"\"\"\n        model_files = glob.glob(str(self.model_dir / f\"{self.name}_e=*.pt\"))\n        if len(model_files) > 5:\n            # Extract integer values from filenames\n            def extract_epoch(filename: str) -> Union[int, float]:\n                match = re.search(r\"_e=(\\d+)\\.pt\", filename)\n                return int(match.group(1)) if match else float('inf')\n\n            # Sort files by extracted integer values\n            model_files.sort(key=extract_epoch)\n\n            # Delete old files, keeping only the last 5\n            for old_model in model_files[:-5]:\n                logger.warning(f\"Too many model checkpoints, Removing old model: {old_model}\")\n                os.remove(old_model)\n\n    def get_outputs(self, data_loader: DataLoader, try_cuda: bool = True) -> Tuple[Tensor, Tensor]:\n        \"\"\"\n        Calculates and returns the outputs of the model on a given dataset.\n\n        Args:\n            data_loader (DataLoader) - Data loader.\n            try_cuda (bool) - Try to use CUDA flag.\n\n        Returns:\n            Tuple[Tensor, Tensor] - The outputs (class, probability) of the model on the given dataset.\n        \"\"\"\n        use_cuda = try_cuda and torch.cuda.is_available()\n        if use_cuda:\n            self.cuda()\n        self.eval()\n        outputs: list[Tensor] = []\n        with torch.no_grad():\n            for x, _ in data_loader:\n                if use_cuda:\n                    x: Tensor = x.cuda()\n                y_hat: Tensor = self(x)\n                outputs.append(y_hat)\n        self.train()\n        logits, preds = torch.cat(outputs).max()\n        probs = torch.softmax(logits, dim=1)\n\n        return preds, probs\n    \n    @classmethod\n    def load_model(cls, name: str, epoch: Union[int, str]) -> \"BaseModel\":\n        \"\"\"\n        Loads a model from a file.\n\n        Args:\n            name (str) - The model's name.\n            epoch (Union[int, str]) - The epoch to load.\n\n        Returns:\n            BaseModel - The loaded model.\n        \"\"\"\n        model = cls(name)\n        model.load_state_dict(torch.load(MODELS_PATH / f\"{name}_ckpts\" / f\"{name}_e={epoch}.pt\"))\n        return model\n\n    def make_model_dir(self, name: str, num_tries: int = 3) -> None:\n        \"\"\"Creates the model directory.\"\"\"\n        model_dir = MODELS_PATH / f\"{name}_ckpts\"\n        if not model_dir.exists():\n            for try_num in range(num_tries + 1):\n                if try_num == num_tries:\n                    logger.error(f\"Failed to create model directory: {model_dir}\")\n                    raise Exception(f\"Failed to create model directory: {model_dir}\")\n                try:\n                    model_dir.mkdir()\n                    logger.info(f\"Created model directory: {model_dir}\")\n                    break\n                except Exception as e:\n                    logger.error(f\"[Try: {try_num+1}/{num_tries+1}] - Failed to create model directory: {model_dir}.\"\n                                 f\" {e}\")\n        else:\n            logger.warning(f\"Model directory already exists: {model_dir}.\"\n                           f\"\\nMake sure you are not overwriting an existing model.\")\n\n        self.model_dir: Path = model_dir\n# ========================== End Of BaseModel Class ========================== #","metadata":{"execution":{"iopub.status.busy":"2025-01-15T21:14:49.065718Z","iopub.execute_input":"2025-01-15T21:14:49.065962Z","iopub.status.idle":"2025-01-15T21:14:49.094547Z","shell.execute_reply.started":"2025-01-15T21:14:49.065939Z","shell.execute_reply":"2025-01-15T21:14:49.093885Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def normalize(x: Tensor, eps: float = 1e-8) -> Tensor:\n    x /= x.max()\n    return (x - x.mean()) / (x.std() + eps)\n\nclass LumbarSpineStenosisResNet(BaseModel):\n    \"\"\"Lumbar Spine Stenosis ResNet model.\"\"\"\n    num_levels: int = 5\n    num_conditions: int = 5\n    num_severities: int = 3\n    num_total_classes: int = num_levels * num_conditions * num_severities\n    single_channel_mean: List[float] = [0]  # Turn off pre-trained normalization.\n    single_channel_std: List[float] = [1]  # Turn off pre-trained normalization.\n    normalizer: Callable[..., Tensor] = staticmethod(normalize)\n\n    def __init__(\n        self,\n        architecture: Literal[\"R3D_18\", \"S3D\", \"MC3_18\"] = \"MC3_18\",\n        pretrained: bool = False,\n        progress: bool = True,\n        **kwargs\n    ) -> None:\n        super().__init__(name=kwargs.get(\"name\", \"\"))\n        assert isinstance(architecture, str), \"architecture must be a string.\"\n        _hidden_size = kwargs.get(\"hidden_size\", 1024)\n        _dropout_val = kwargs.get(\"dropout\", kwargs.get(\"p\", 0.5))\n        _max_grad_norm = kwargs.get(\"max_grad_norm\", 1.0)\n        self.architecture = architecture\n        if architecture == \"R3D_18\":\n            # Load the pre-trained 3D ResNet18 model.\n            # See: https://pytorch.org/vision/main/models/generated/torchvision.models.video.r3d_18.html\n            # Under: #torchvision.saved_models.video.R3D_18_Weights\n            self.pre_trained_transforms = R3D_18_Weights.KINETICS400_V1.transforms(\n                mean=self.single_channel_mean, std=self.single_channel_std\n            )\n            _weights = R3D_18_Weights.KINETICS400_V1 if pretrained else None\n            self.model = r3d_18(weights=_weights, progress=progress)\n\n            # Modify the first convolutional layer to accept 1 input channel instead of 3.\n            # We do that by taking the average of the 3 input channels (ensemble).\n            first_conv_layer: nn.Conv3d = self.model.stem._modules[\"0\"]  # type: ignore\n            first_conv_layer.weight = nn.Parameter(first_conv_layer.weight.mean(dim=1, keepdim=True))\n            first_conv_layer.__dict__[\"in_channels\"] = 1\n\n            # Modify the fully connected layer.\n            self.model.fc = nn.Sequential(\n                nn.Linear(self.model.fc.in_features, _hidden_size),\n                nn.ReLU(),\n                nn.Dropout(p=_dropout_val),\n                nn.Linear(_hidden_size, self.num_total_classes)\n            )\n\n        elif architecture == \"S3D\":\n            raise NotImplementedError(\"S3D architecture is not supported due to mismatched input size.\")\n            # Load the pre-trained S3D model.\n            # See: https://pytorch.org/vision/main/models/generated/torchvision.models.video.r3d_18.html\n            # Under: #torchvision.saved_models.video.S3D_Weights\n            self.pre_trained_transforms = S3D_Weights.KINETICS400_V1.transforms(\n                mean=self.single_channel_mean, std=self.single_channel_std\n            )\n            _weights = S3D_Weights.KINETICS400_V1 if pretrained else None\n            self.model = s3d(weights=_weights, progress=progress)\n\n            # Modify the first convolutional layer to accept 1 input channel instead of 3.\n            # We do that by taking the average of the 3 input channels (ensemble).\n            first_conv_layer: nn.Conv3d = self.model.features[0][0][0]\n            first_conv_layer.weight = nn.Parameter(first_conv_layer.weight.mean(dim=1, keepdim=True))\n            first_conv_layer.__dict__[\"in_channels\"] = 1\n\n            # Modify the fully connected layer.\n            self.model.classifier = nn.Sequential(\n                nn.Dropout(p=_dropout_val),\n                nn.Conv3d(1024, self.num_total_classes, kernel_size=1, stride=1, bias=True)\n            )\n        elif architecture == \"MC3_18\":\n            # Load the pre-trained MC3_18 model.\n            # See: https://pytorch.org/vision/main/models/generated/torchvision.models.video.r3d_18.html\n            # Under: #torchvision.saved_models.video.MC3_18_Weights\n            self.pre_trained_transforms = MC3_18_Weights.KINETICS400_V1.transforms(\n                mean=self.single_channel_mean, std=self.single_channel_std\n            )\n\n            _weights = MC3_18_Weights.KINETICS400_V1 if pretrained else None\n            self.model = mc3_18(weights=_weights, progress=progress)\n\n            # Modify the first convolutional layer to accept 1 input channel instead of 3.\n            # We do that by taking the average of the 3 input channels (ensemble).\n            first_conv_layer: nn.Conv3d = self.model.stem._modules[\"0\"]  # type: ignore\n            first_conv_layer.weight = nn.Parameter(first_conv_layer.weight.mean(dim=1, keepdim=True))\n            first_conv_layer.__dict__[\"in_channels\"] = 1\n\n            # Modify the fully connected layer.\n            self.model.fc = nn.Sequential(\n                nn.Linear(self.model.fc.in_features, _hidden_size),\n                nn.ReLU(),\n                nn.Dropout(p=_dropout_val),\n                nn.Linear(_hidden_size, self.num_total_classes)\n            )\n\n        else:\n            raise NotImplementedError(f\"architecture '{architecture}' is not supported.\")\n\n        self.log_softmax = nn.LogSoftmax(dim=2)\n\n        torch.nn.utils.clip_grad_norm_(self.parameters(), max_norm=_max_grad_norm)\n        self._add_forward_transforms()\n    \n    def forward(self, x: Tensor) -> Tensor:\n        x = self.model(x)\n        x = x.view(-1, self.num_levels * self.num_conditions, self.num_severities)\n        x = self.log_softmax(x)  # Apply softmax over the severity classes\n        return x\n    \n    def _add_forward_transforms(self) -> None:\n        if hasattr(self, \"pre_trained_transforms\"):\n            def decorator(forward: Callable[[Tensor], Tensor]) -> Callable:\n                def wrapper(x: Tensor) -> Tensor:\n                    x = self.pre_trained_transforms(x)\n                    x = self.normalizer(x)\n                    return forward(x)\n                return wrapper\n            self.forward = decorator(self.forward)\n\n    def _calc_running_metrics(self, x: Tensor, y_hat: Tensor, y: Tensor, loss: float,\n                              running_loss: float, total_corrects: float) -> Tuple[float, float]:\n        \"\"\"Override the base class method to calculate the running metrics.\"\"\"\n        corrects = int((y_hat.argmax(-1) == y).sum().item())\n\n        running_loss += loss\n        total_corrects += corrects / self.num_levels / self.num_conditions\n        return running_loss, total_corrects","metadata":{"execution":{"iopub.status.busy":"2025-01-15T21:14:49.095813Z","iopub.execute_input":"2025-01-15T21:14:49.096143Z","iopub.status.idle":"2025-01-15T21:14:49.116048Z","shell.execute_reply.started":"2025-01-15T21:14:49.096108Z","shell.execute_reply":"2025-01-15T21:14:49.115301Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class SingleModelLumbarSpineDataset(BaseLumbarSpineDataset):\n\n    def __len__(self) -> int:\n        return self.series_description.shape[0]\n\n    def __getitem__(self, idx: int) -> tuple[torch.Tensor, Optional[torch.Tensor]]:\n        \"\"\"returns x (D, 1, H, W), y (75)\"\"\"\n        desc_row = self.series_description.iloc[idx]\n        series_path = self.images_path / str(desc_row[\"study_id\"]) / str(desc_row[\"series_id\"])\n        dicom_files = load_dicom_series(series_path)\n\n        x: torch.Tensor = self.preprocess(dicom_files)  # type: ignore - torch.from_numpy\n        if self.train:\n            if self._apply_augs and self.augs is not None:\n                x = self.augs(x)\n            condition_row = self.train_csv.loc[desc_row[\"study_id\"]]\n            y = torch.tensor(\n                [self.severity_s2i.get(condition_row[col], self.nan_class) for col in self.conditions_i2s],\n                dtype=torch.int64\n            )\n        else:\n            y = None\n\n        return x, y\n","metadata":{"execution":{"iopub.status.busy":"2025-01-15T21:14:49.116957Z","iopub.execute_input":"2025-01-15T21:14:49.117280Z","iopub.status.idle":"2025-01-15T21:14:49.131736Z","shell.execute_reply.started":"2025-01-15T21:14:49.117244Z","shell.execute_reply":"2025-01-15T21:14:49.130870Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class MultiModelLumbarSpineDataset(BaseLumbarSpineDataset):\n    def __init__(\n        self,\n        train: bool, \n        *,\n        preprocess: Optional[transforms.Compose] = None,\n        augs: Optional[transforms.Compose] = None\n    ) -> None:\n        super().__init__(train=train, preprocess=preprocess, augs=augs)\n        self.series_description.set_index(\"study_id\", inplace=True)\n        self.study_ids = self.series_description.index.unique()\n\n    def __len__(self) -> int:\n        return self.study_ids.shape[0]\n\n    def __getitem__(self, idx: int) -> Dict[str, Union[List[torch.Tensor], torch.Tensor]]:\n        \"\"\"returns {\"data\": [(D, 1, H, W)], \"target\": (25), \"series_types\": [(0-2)]}\"\"\"\n        study_id = self.study_ids[idx]\n        study_df = self.series_description.loc[study_id]  # Expect at least 2 MRI series per study.\n        series_paths = study_df[\"series_id\"].apply(\n            lambda series_id: self.images_path / str(study_id) / str(series_id)\n        )\n        data_dict = {\n            \"data\": series_paths.apply(load_dicom_series).apply(self.preprocess).tolist(),\n            \"series_types\": torch.tensor(study_df[\"series_description\"].map(SERIES_S2I).values,\n                                         dtype=torch.int64)\n        }\n        if self.train:\n            if self._apply_augs and self.augs is not None:\n                data_dict[\"data\"] = list(map(self.augs, data_dict[\"data\"]))\n            condition_row = self.train_csv.loc[study_id]\n            data_dict[\"target\"] = torch.tensor(\n                [self.severity_s2i.get(condition_row[col], self.nan_class) for col in self.conditions_i2s],\n                dtype=torch.int64\n            )\n        else:\n            data_dict[\"row_id\"] = study_id\n\n        return data_dict\n\n\ndef load_dicom_series(directory: Union[str, Path]) -> np.ndarray:\n    \"\"\"\n    Load a DICOM series from a directory into a 3D numpy array.\n\n    Args:\n        directory (Union[str, Path]): The directory containing the DICOM series.\n\n    Returns:\n        np.ndarray: The 3D numpy array containing the DICOM series.\n    \"\"\"\n    dicom_files = []\n    for root, _, files in os.walk(directory):\n        # TODO: Raise an exception if no files are found.\n        for file in files:\n            if file.endswith(\".dcm\"):\n                dicom_files.append(pydicom.dcmread(os.path.join(root, file)))\n    dicom_files.sort(key=lambda x: int(x.InstanceNumber))\n    \n    # Stack all slices into a 3D numpy array.\n    image_3d = [apply_voi_lut(dcmf.pixel_array, dcmf) for dcmf in dicom_files]\n    \n    # In rare cases, sometimes the pixel arrays are not of the same shape.\n    if (shapes := np.unique([np.array(img.shape) for img in image_3d], axis=0)).shape[0] > 1:\n        # Calculate the square padding for each image.\n        max_h, max_w = np.max(shapes, axis=0)\n        padding_dict = {\n            (h, w): (\n                (floor((max_h - h) / 2), ceil((max_h - h) / 2)),\n                (floor((max_w - w) / 2), ceil((max_w - w) / 2))\n            )\n            for (h, w) in shapes\n        }\n        stacked_images = np.stack([\n            np.pad(img, padding_dict[tuple(img.shape)], mode=\"constant\", constant_values=0)\n            for img in image_3d\n        ])\n    else:\n        stacked_images = np.stack(image_3d)\n    \n    return stacked_images\n","metadata":{"execution":{"iopub.status.busy":"2025-01-15T21:14:49.133641Z","iopub.execute_input":"2025-01-15T21:14:49.133917Z","iopub.status.idle":"2025-01-15T21:14:49.148812Z","shell.execute_reply.started":"2025-01-15T21:14:49.133894Z","shell.execute_reply":"2025-01-15T21:14:49.148133Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class SingleModelSpineCNN(LumbarSpineStenosisResNet):\n    \"\"\"Single-Model Spine CNN model.\"\"\"\n    def __init__(self, *args, **kwargs) -> None:\n        super().__init__(*args, name=\"do_not_save\", **kwargs)\n        _dropout_val = kwargs.get(\"dropout\", kwargs.get(\"p\", 0.5))\n        _out_features_size = kwargs.get(\"out_features_size\", 512)\n        self.out_features_size = _out_features_size\n\n        # Reinitialize the fully connected layer.\n        if self.architecture == \"R3D_18\":\n            self.model.fc = nn.Linear(self.model.fc[0].in_features, _out_features_size)\n        elif self.architecture == \"S3D\":\n            self.model.classifier = nn.Sequential(\n                nn.Dropout(p=_dropout_val),\n                nn.Conv3d(1024, _out_features_size, kernel_size=1, stride=1, bias=True)\n            )\n        elif self.architecture == \"MC3_18\":\n            self.model.fc = nn.Linear(self.model.fc[0].in_features, _out_features_size)\n        else:\n            raise RuntimeError(\"Architecture must be either 'R3D_18' or 'S3D' or 'MC3_18', this should be unreachable.\")\n\n    def forward(self, x: Tensor) -> Tensor:\n        x = self.model(x).view(-1)  # Batch size must be 1.\n        return x","metadata":{"execution":{"iopub.status.busy":"2025-01-15T21:14:49.149680Z","iopub.execute_input":"2025-01-15T21:14:49.149977Z","iopub.status.idle":"2025-01-15T21:14:49.162549Z","shell.execute_reply.started":"2025-01-15T21:14:49.149952Z","shell.execute_reply":"2025-01-15T21:14:49.161940Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class MultiModelSpineCNN(BaseModel):\n    \"\"\"Multi-Model Spine CNN model.\"\"\"\n    def __init__(self, *model_dicts, last_fc_dim: int = 1024, dropout: float = 0.5, **kwargs) -> None:\n        super().__init__(**kwargs)\n        assert all(isinstance(model_dict, dict) for model_dict in model_dicts), \\\n            \"All model arguments must be dictionaries for model initialization.\"\n        self.models: List[SingleModelSpineCNN] = []\n        for i, model_dict in enumerate(model_dicts):\n            setattr(self, f\"sub_model_{i}\", SingleModelSpineCNN(**model_dict))\n            self.models.append(getattr(self, f\"sub_model_{i}\"))\n        self.fc = nn.Sequential(\n            nn.Linear(sum([model_dict[\"out_features_size\"] for model_dict in model_dicts]), last_fc_dim),\n            nn.ReLU(),\n            nn.Dropout(p=dropout),\n            nn.Linear(last_fc_dim, SingleModelSpineCNN.num_total_classes)\n        )\n        self.log_softmax = nn.LogSoftmax(dim=1)\n\n        self.optimizer: Optional[optim.Optimizer] = None\n        self.scheduler: Optional[lr_scheduler.LRScheduler] = None\n        self._device: Literal[\"cpu\", \"cuda\"] = \"cpu\"\n\n    def forward(self, data_dict: Dict[str, Union[List[Tensor], Tensor]]) -> Tensor:\n        assert \"data\" in data_dict and \"series_types\" in data_dict, \\\n            \"data_dict must contain 'data' and 'series_types' keys.\"\n        assert len(data_dict[\"data\"]) > 0 and len(data_dict[\"series_types\"]) > 0, \\\n            \"data and series_types must not be empty.\"\n        assert len(data_dict[\"data\"]) == len(data_dict[\"series_types\"]), \\\n            \"data and series_types must have the same length.\"\n\n        # Assume series_type is encoded.\n        out_features = [[] for _ in range(len(self.models))]\n        for i, (data, series_type) in enumerate(zip(data_dict[\"data\"], data_dict[\"series_types\"])):\n            out_features[series_type].append(self.models[series_type](data))\n\n        # Handle case when there are no series of a certain type.\n        for i in range(len(self.models)):\n            if out_features[i] == []:\n                out_features[i] = [torch.zeros(self.models[i].out_features_size, device=self._device)]\n\n        # Handle multiple same series types.\n        out_features = [torch.stack(out).mean(dim=0) for out in out_features]\n        total_features = torch.cat(out_features, dim=0)\n\n        y_hat = self.fc(total_features)\n        y_hat = y_hat.view(SingleModelSpineCNN.num_levels * SingleModelSpineCNN.num_conditions,\n                           SingleModelSpineCNN.num_severities)\n        y_hat = self.log_softmax(y_hat)\n        return y_hat.unsqueeze(0)  # Add a batch dimension.\n\n    def fit(self, train_loader: DataLoader, val_loader: Optional[DataLoader] = None,\n            num_epochs: int = 30, lr: float = 0.001, momentum: float = 0.9, wd: float = 0.,\n            try_cuda: bool = True, verbose: bool = True, print_stride: int = 1) -> None:\n        \"\"\"Override function for training a multi-model spine 3D-CNN.\"\"\"\n        use_cuda = try_cuda and torch.cuda.is_available()\n        if use_cuda:\n            self.cuda()\n            self._device = \"cuda\"\n            logger.info(\"Using CUDA for training.\")\n        else:\n            self.cpu()\n            self._device = \"cpu\"\n            logger.info(\"Using CPU for training.\")\n\n        # Create the optimizer.\n        self.optimizer = optim.AdamW(self.parameters(), lr=lr, weight_decay=wd)\n        self.scheduler = lr_scheduler.CosineAnnealingLR(self.optimizer, T_max=num_epochs)\n\n        start_epoch = self.global_epoch\n        for epoch in range(num_epochs):\n            running_loss: float = 0.\n            total_corrects: int = 0\n            self.global_epoch += 1\n            for mb, data_dict in enumerate(train_loader):\n                data_dict: Dict[str, Union[List[Tensor], Tensor]]\n                # Assume batch size is 1, remove the batch dimension.\n                #data_dict[\"data\"] = [x[0] for x in data_dict[\"data\"]]\n                data_dict[\"target\"] = data_dict[\"target\"][0]\n                data_dict[\"series_types\"] = data_dict[\"series_types\"][0]\n                if use_cuda:\n                    data_dict[\"data\"]: List[Tensor] = [x.cuda() for x in data_dict[\"data\"]]\n                    data_dict[\"target\"]: Tensor = data_dict[\"target\"].cuda()\n\n                y_hat, loss = self.train_step(data_dict)\n\n                running_loss, total_corrects = self.calc_running_metrics(\n                    data_dict, y_hat, loss, running_loss, total_corrects)\n\n                # TODO: Use tqdm instead of print.\n                if verbose:\n                    print(f\"\\r[epoch: {self.global_epoch:02d}/{start_epoch + num_epochs:02d} \"\n                          f\"mb: {mb + 1:03d}/{len(train_loader):03d}  lr: {self.scheduler.get_last_lr()[0]:.6f}]  \"\n                          f\"[Train loss: {loss:.6f}]\\033[J\", end=\"\")\n\n            train_epoch_loss = running_loss / len(train_loader.dataset)  # type: ignore\n            train_total_acc = total_corrects / len(train_loader.dataset)  # type: ignore\n            self.train_costs.append(train_epoch_loss)\n            self.train_accs.append(train_total_acc)\n\n            if val_loader is not None:\n                val_epoch_loss, val_total_acc = self.calc_metrics(val_loader, use_cuda)\n            else:\n                val_epoch_loss, val_total_acc = torch.nan, torch.nan\n\n            self.val_costs.append(val_epoch_loss)\n            self.val_accs.append(val_total_acc)\n\n            self.save_best_weights()\n\n            if verbose and (epoch % print_stride == 0 or epoch == num_epochs - 1):\n                logger.info(f\"\\r[epoch: {self.global_epoch:02d}/{start_epoch + num_epochs:02d}\"\n                            f\"  lr: {self.scheduler.get_last_lr()[0]:.6f}] \"\n                            f\"[Train loss: {train_epoch_loss:.6f} \"\n                            f\" Train acc: {100 * train_total_acc:.6f}%]  \"\n                            f\"[Val loss: {val_epoch_loss:.6f} \"\n                            f\" Val acc: {100 * val_total_acc:.6f}%]\")\n            self.scheduler.step()\n\n        self._device = \"cpu\"\n        self.cpu()\n\n    def train_step(\n        self, data_dict: Dict[str, Union[List[Tensor], Tensor]]\n    ) -> Tuple[Tensor, float]:\n        self.optimizer.zero_grad()\n        y_hat = self(data_dict)\n        loss = F.nll_loss(y_hat.view(-1, y_hat.size(-1)), data_dict[\"target\"], weight=torch.tensor([1/7, 2/7, 4/7], device=self._device))\n        loss.backward()\n        self.optimizer.step()\n        return y_hat, loss.item()\n\n    def calc_metrics(self, data_loader: DataLoader, use_cuda: bool) -> Tuple[float, float]:\n        \"\"\"Override function for calculating the validation metrics.\"\"\"\n        running_loss: float = 0.\n        total_corrects: int = 0\n        self.eval()\n        with torch.no_grad():\n            for data_dict in tqdm(data_loader, desc=\"Validation\", total=len(data_loader)):\n                data_dict: Dict[str, Union[List[Tensor], Tensor]]\n                # Assume batch size is 1, remove the batch dimension.\n                # data_dict[\"data\"] = [x[0] for x in data_dict[\"data\"]]\n                data_dict[\"target\"] = data_dict[\"target\"][0]\n                data_dict[\"series_types\"] = data_dict[\"series_types\"][0]\n                if use_cuda:\n                    data_dict[\"data\"]: List[Tensor] = [x.cuda() for x in data_dict[\"data\"]]\n                    data_dict[\"target\"]: Tensor = data_dict[\"target\"].cuda()\n\n                y_hat: Tensor = self(data_dict)\n                loss: float = F.nll_loss(\n                    y_hat.view(-1, y_hat.size(-1)), data_dict[\"target\"]\n                ).item()\n\n                running_loss, total_corrects = self.calc_running_metrics(\n                    data_dict, y_hat, loss, running_loss, total_corrects)\n\n        self.train()\n        total_loss = running_loss / len(data_loader.dataset)  # type: ignore\n        total_acc = total_corrects / len(data_loader.dataset)  # type: ignore\n        return total_loss, total_acc\n\n    def calc_running_metrics(\n        self, data_dict: Dict[str, Union[List[Tensor], Tensor]], y_hat: Tensor, loss: float,\n        running_loss: float, total_corrects: int\n    ) -> Tuple[float, int]:\n        \"\"\"Override function for calculating the running metrics.\"\"\"\n        corrects = int((y_hat.argmax(-1) == data_dict[\"target\"]).sum().item())\n        running_loss += loss\n        total_corrects += corrects / LumbarSpineStenosisResNet.num_levels / LumbarSpineStenosisResNet.num_conditions\n        return running_loss, total_corrects\n","metadata":{"execution":{"iopub.status.busy":"2025-01-15T21:14:49.163900Z","iopub.execute_input":"2025-01-15T21:14:49.164435Z","iopub.status.idle":"2025-01-15T21:14:49.188498Z","shell.execute_reply.started":"2025-01-15T21:14:49.164398Z","shell.execute_reply":"2025-01-15T21:14:49.187903Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def make_submission_with_loaded_multi_model(model: MultiModelSpineCNN, save_suffix: str) -> None:\n    output_df = pd.DataFrame(columns=[\"row_id\", \"normal_mild\", \"moderate\", \"severe\"])\n    test_dataset = MultiModelLumbarSpineDataset(False)\n    dir_path = SUBMISSION_PATH\n    \n    model.load_best_weights()\n    model.eval()\n    with torch.no_grad():\n        for idx in range(len(test_dataset)):\n            test_dict = test_dataset[idx]\n            test_dict[\"data\"] = [torch.unsqueeze(data, 0) for data in test_dict[\"data\"]]\n            y_hat = model(test_dict).view(25, 3)  # Only one sample in the test set.\n            y_hat = F.softmax(y_hat, dim=1)\n            curr_df = pd.DataFrame(\n                {\n                    \"row_id\": [f\"{test_dict['row_id']}_{condition}\" for condition in test_dataset.conditions_i2s],\n                    \"normal_mild\": y_hat[:, 0].numpy().astype(np.float32),\n                    \"moderate\": y_hat[:, 1].numpy().astype(np.float32),\n                    \"severe\": y_hat[:, 2].numpy().astype(np.float32)\n                }\n            )\n            output_df = pd.concat([output_df, curr_df], axis=0, ignore_index=True)\n    output_df.sort_values(by=\"row_id\", inplace=True)\n    output_df.to_csv(dir_path / f\"{save_suffix}.csv\", index=False)","metadata":{"execution":{"iopub.status.busy":"2025-01-15T21:14:49.189464Z","iopub.execute_input":"2025-01-15T21:14:49.189786Z","iopub.status.idle":"2025-01-15T21:14:49.204219Z","shell.execute_reply.started":"2025-01-15T21:14:49.189760Z","shell.execute_reply":"2025-01-15T21:14:49.203427Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def add_gaussian_noise(image: torch.Tensor) -> torch.Tensor:\n    var = image.max() * 0.05 + 1e-7  # TODO: Maybe add variance in each depth slice.\n    noise = torch.randn(image.size(), dtype=torch.float32) * torch.sqrt(var)\n    noisy_tensor = image + noise\n    return torch.clamp(noisy_tensor, image.min(), image.max()) \n\ndataset = MultiModelLumbarSpineDataset(train=True)#, augs=add_gaussian_noise)\ntrain_dataset, val_dataset = dataset.split(val_size=0.15)\ntrain_dataloader = DataLoader(train_dataset, batch_size=1, shuffle=True, num_workers=4)\nval_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=4)\n\nsags_args = dict(architecture=\"R3D_18\", pretrained=False, progress=True, out_features_size=1024)\n#sag_t2_args = dict(architecture=\"MC3_18\", pretrained=True, progress=True, out_features_size=512)\naxial_t2_args = dict(architecture=\"R3D_18\", pretrained=False, progress=True, out_features_size=750)\n\nif not os.path.exists(\"/kaggle/working/saved_models\"):\n    os.mkdir(\"/kaggle/working/saved_models\")\n\nmodel = MultiModelSpineCNN(sags_args, axial_t2_args, last_fc_dim=1024 + 750, dropout=0.5,\n                            name=\"multi_model_submission\")\nmodel.fit(train_loader=train_dataloader, val_loader=val_dataloader, num_epochs=18,\n            lr=0.0005, momentum=0.9, wd=0.0005, try_cuda=True, verbose=True, print_stride=1)","metadata":{"execution":{"iopub.status.busy":"2025-01-15T21:14:49.205091Z","iopub.execute_input":"2025-01-15T21:14:49.205345Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"make_submission_with_loaded_multi_model(model, \"submission\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"pd.read_csv(\"/kaggle/working/submission.csv\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}