{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"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"},"jupytext":{"cell_metadata_filter":"-all","main_language":"python","notebook_metadata_filter":"-all"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":99552,"databundleVersionId":13851420,"sourceType":"competition"},{"sourceId":13444317,"sourceType":"datasetVersion","datasetId":8533725},{"sourceId":13444908,"sourceType":"datasetVersion","datasetId":8534135}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# 7th place solution - 3D nnU-Net + blob regression (again)\n\nWriteup: https://www.kaggle.com/competitions/rsna-intracranial-aneurysm-detection/writeups/7th-place-solution\n\nCode: https://github.com/MIC-DKFZ/kaggle-rsna-intracranial-aneurysm-detection-2025-solution#","metadata":{}},{"cell_type":"markdown","source":"## Preparation","metadata":{}},{"cell_type":"code","source":"import os\nos.listdir('/kaggle/input')","metadata":{"_uuid":"74b7a1c0-43f0-4446-af0b-dd6126fca0e6","_cell_guid":"fcb52739-43aa-4b11-8b71-70e860038032","trusted":true,"collapsed":false,"execution":{"iopub.status.busy":"2025-10-20T15:10:46.02648Z","iopub.execute_input":"2025-10-20T15:10:46.026646Z","iopub.status.idle":"2025-10-20T15:10:46.034688Z","shell.execute_reply.started":"2025-10-20T15:10:46.026632Z","shell.execute_reply":"2025-10-20T15:10:46.033869Z"},"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Installing packages","metadata":{}},{"cell_type":"code","source":"!pip install --no-index --find-links /kaggle/input/rsna-2025-7th-place-packages/ /kaggle/input/rsna-2025-7th-place-packages/*.whl","metadata":{"_uuid":"cc86f678-bc25-451f-8f1e-9eff5a461219","_cell_guid":"77cd227c-5d69-4eb7-9892-e9401bb4fe36","trusted":true,"scrolled":true,"execution":{"iopub.status.busy":"2025-10-20T15:10:46.036138Z","iopub.execute_input":"2025-10-20T15:10:46.036397Z","iopub.status.idle":"2025-10-20T15:15:30.390965Z","shell.execute_reply.started":"2025-10-20T15:10:46.03638Z","shell.execute_reply":"2025-10-20T15:15:30.389975Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport shutil\nfrom collections import defaultdict\n\nimport pandas as pd\nimport polars as pl\nimport pydicom\n\nimport kaggle_evaluation.rsna_inference_server\nLABEL_COLS = ['Left Infraclinoid Internal Carotid Artery', 'Right Infraclinoid Internal Carotid Artery', 'Left Supraclinoid Internal Carotid Artery', 'Right Supraclinoid Internal Carotid Artery', 'Left Middle Cerebral Artery', 'Right Middle Cerebral Artery', 'Anterior Communicating Artery', 'Left Anterior Cerebral Artery', 'Right Anterior Cerebral Artery', 'Left Posterior Communicating Artery', 'Right Posterior Communicating Artery', 'Basilar Tip', 'Other Posterior Circulation', 'Aneurysm Present']","metadata":{"_uuid":"a2a71bde-a95c-4e06-9c9c-42c9d500f02c","_cell_guid":"a2685d5a-ab54-4699-a733-cc407bac1725","trusted":true,"collapsed":false,"execution":{"iopub.status.busy":"2025-10-20T15:17:51.951938Z","iopub.execute_input":"2025-10-20T15:17:51.952259Z","iopub.status.idle":"2025-10-20T15:17:51.957354Z","shell.execute_reply.started":"2025-10-20T15:17:51.952232Z","shell.execute_reply":"2025-10-20T15:17:51.956739Z"},"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Overwriting nnUNetPredictor to split across both GPUs to speed up inference","metadata":{}},{"cell_type":"code","source":"%%writefile /kaggle/working/nnUNetPredictor.py\n\nimport os\nfrom queue import Queue\nfrom threading import Thread\nfrom typing import Tuple, Union, List\nfrom queue import Queue, Empty\nfrom threading import Thread\nimport itertools\n\nimport numpy as np\nimport torch\nfrom acvl_utils.cropping_and_padding.padding import pad_nd_image\nfrom batchgenerators.utilities.file_and_folder_operations import load_json, join, isfile, subdirs\nfrom torch._dynamo import OptimizedModule\nfrom tqdm import tqdm\nimport copy\nimport traceback\nfrom scipy.ndimage import gaussian_filter\n\nimport nnunetv2\nfrom nnunetv2.configuration import default_num_processes\nfrom nnunetv2.inference.nnLandmark.sliding_window_prediction import compute_steps_for_sliding_window\n\nfrom nnunetv2.utilities.find_class_by_name import recursive_find_python_class\nfrom nnunetv2.utilities.helpers import empty_cache, dummy_context\nfrom nnunetv2.utilities.label_handling.label_handling import determine_num_input_channels\nfrom nnunetv2.utilities.plans_handling.plans_handler import PlansManager\n\ndef compute_gaussian(tile_size: Union[Tuple[int, ...], List[int]], sigma_scale: float = 1. / 8,\n                    value_scaling_factor: float = 1, dtype=torch.float16, device=torch.device('cuda', 0)) \\\n        -> torch.Tensor:\n    tmp = np.zeros(tile_size)\n    center_coords = [i // 2 for i in tile_size]\n    sigmas = [i * sigma_scale for i in tile_size]\n    tmp[tuple(center_coords)] = 1\n    gaussian_importance_map = gaussian_filter(tmp, sigmas, 0, mode='constant', cval=0)\n\n    gaussian_importance_map = torch.from_numpy(gaussian_importance_map)\n\n    gaussian_importance_map /= (torch.max(gaussian_importance_map) / value_scaling_factor)\n    gaussian_importance_map = gaussian_importance_map.to(device=device, dtype=dtype)\n    # gaussian_importance_map cannot be 0, otherwise we may end up with nans!\n    mask = gaussian_importance_map == 0\n    gaussian_importance_map[mask] = torch.min(gaussian_importance_map[~mask])\n    return gaussian_importance_map\n\nclass nnUNetPredictor(object):\n    def __init__(self,\n                 tile_step_size: float = 0.5,\n                 use_gaussian: bool = True,\n                 use_mirroring: bool = True,\n                 device: torch.device = torch.device('cuda'),\n                 verbose: bool = False,\n                 verbose_preprocessing: bool = False,\n                 allow_tqdm: bool = True,\n                 ):\n        self.verbose = verbose\n        self.verbose_preprocessing = verbose_preprocessing\n        self.allow_tqdm = allow_tqdm\n\n        self.plans_manager, self.configuration_manager, self.list_of_parameters, self.network, self.dataset_json, \\\n        self.trainer_name, self.allowed_mirroring_axes, self.label_manager = None, None, None, None, None, None, None, None\n\n        self.tile_step_size = tile_step_size\n        self.use_gaussian = use_gaussian\n        self.use_mirroring = use_mirroring\n\n        self._per_device_nets = {}   # {int(device.index): Module on that CUDA device}\n        self._net_uncompiled = None  # filled in initialize_from_trained_model_folder\n\n        \n        assert torch.cuda.is_available(), \"CUDA required. CPU execution is disabled by design.\"\n        torch.backends.cudnn.benchmark = True\n        self.device = device  # primary device (cuda:0 by default)\n\n    def _get_or_make_device_net(self, dev: torch.device):\n        \"\"\"\n        Get a persistent eval (optionally compiled) clone of the base net on `dev`.\n        Reused across images. Thread-safe for inference.\n        \"\"\"\n        didx = dev.index if dev.index is not None else 0\n        net = self._per_device_nets.get(didx, None)\n        if net is not None:\n            return net\n\n        # always clone from the CPU base to avoid extra VRAM copies on cuda:0\n        base = getattr(self, \"_net_uncompiled\", self.network)\n        assert all(not p.is_cuda for p in base.parameters()), \"base net must be on CPU for multi-GPU\"\n        net = copy.deepcopy(base).to(dev)\n        net.eval()\n\n        # compile the per-device net if requested (and not already optimized)\n        if ('nnUNet_compile' in os.environ) and (os.environ['nnUNet_compile'].lower() in ('true','1','t')) \\\n           and not isinstance(net, OptimizedModule):\n            net = torch.compile(net)\n\n        self._per_device_nets[didx] = net\n        return net\n\n    def _clear_device_nets(self):\n        \"\"\"Drop all cached per-device nets (used on hard errors to free VRAM).\"\"\"\n        for didx, net in list(self._per_device_nets.items()):\n            try:\n                dev = torch.device(f'cuda:{didx}')\n                del net\n                torch.cuda.set_device(dev)\n                torch.cuda.empty_cache()\n                torch.cuda.ipc_collect()\n            except Exception:\n                pass\n        self._per_device_nets.clear()\n    \n    def initialize_from_trained_model_folder(self, model_training_output_dir: str,\n                                             use_folds: Union[Tuple[Union[int, str]], None],\n                                             checkpoint_name: str = 'checkpoint_final.pth'):\n        \"\"\"\n        This is used when making predictions with a trained model\n        \"\"\"\n        if use_folds is None:\n            use_folds = nnUNetPredictor.auto_detect_available_folds(model_training_output_dir, checkpoint_name)\n\n        dataset_json = load_json(join(model_training_output_dir, 'dataset.json'))\n        plans = load_json(join(model_training_output_dir, 'plans.json'))\n        plans_manager = PlansManager(plans)\n\n        if isinstance(use_folds, str):\n            use_folds = [use_folds]\n\n        parameters = []\n        for i, f in enumerate(use_folds):\n            f = int(f) if f != 'all' else f\n            checkpoint = torch.load(join(model_training_output_dir, f'fold_{f}', checkpoint_name),\n                                    map_location=torch.device('cpu'), weights_only=False)\n            if i == 0:\n                trainer_name = checkpoint['trainer_name']\n                configuration_name = checkpoint['init_args']['configuration']\n                inference_allowed_mirroring_axes = checkpoint['inference_allowed_mirroring_axes'] if \\\n                    'inference_allowed_mirroring_axes' in checkpoint.keys() else None\n\n            parameters.append(checkpoint['network_weights'])\n\n        configuration_manager = plans_manager.get_configuration(configuration_name)\n\n        # restore network\n        num_input_channels = determine_num_input_channels(plans_manager, configuration_manager, dataset_json)\n        trainer_class = recursive_find_python_class(join(nnunetv2.__path__[0], \"training\", \"nnUNetTrainer\"),\n                                                    trainer_name, 'nnunetv2.training.nnUNetTrainer')\n        if trainer_class is None:\n            raise RuntimeError(f'Unable to locate trainer class {trainer_name} in nnunetv2.training.nnUNetTrainer. '\n                               f'Please place it there (in any .py file)!')\n        network = trainer_class.build_network_architecture(\n            configuration_manager.network_arch_class_name,\n            configuration_manager.network_arch_init_kwargs,\n            configuration_manager.network_arch_init_kwargs_req_import,\n            num_input_channels,\n            plans_manager.get_label_manager(dataset_json).num_segmentation_heads,\n            enable_deep_supervision=False\n        )\n\n        self.plans_manager = plans_manager\n        self.configuration_manager = configuration_manager\n        self.list_of_parameters = parameters\n\n        # initialize network with first set of parameters, also see https://github.com/MIC-DKFZ/nnUNet/issues/2520\n        network.load_state_dict(parameters[0])\n\n        self.network = network\n        self._net_uncompiled = network  # keep a CPU master for cloning\n\n        self.dataset_json = dataset_json\n        self.trainer_name = trainer_name\n        self.allowed_mirroring_axes = inference_allowed_mirroring_axes\n        self.label_manager = plans_manager.get_label_manager(dataset_json)\n        if ('nnUNet_compile' in os.environ.keys()) and (os.environ['nnUNet_compile'].lower() in ('true', '1', 't')) \\\n                and not isinstance(self.network, OptimizedModule):\n            print('Using torch.compile')\n            self.network = torch.compile(self.network)\n\n    @staticmethod\n    def auto_detect_available_folds(model_training_output_dir, checkpoint_name):\n        print('use_folds is None, attempting to auto detect available folds')\n        fold_folders = subdirs(model_training_output_dir, prefix='fold_', join=False)\n        fold_folders = [i for i in fold_folders if i != 'fold_all']\n        fold_folders = [i for i in fold_folders if isfile(join(model_training_output_dir, i, checkpoint_name))]\n        use_folds = [int(i.split('_')[-1]) for i in fold_folders]\n        print(f'found the following folds: {use_folds}')\n        return use_folds\n\n\n    @torch.inference_mode()\n    def predict_logits_from_preprocessed_data(self, data: torch.Tensor, out_device = torch.device('cpu')) -> torch.Tensor:\n        \"\"\"\n        IMPORTANT! IF YOU ARE RUNNING THE CASCADE, THE SEGMENTATION FROM THE PREVIOUS STAGE MUST ALREADY BE STACKED ON\n        TOP OF THE IMAGE AS ONE-HOT REPRESENTATION! SEE PreprocessAdapter ON HOW THIS SHOULD BE DONE!\n\n        RETURNED LOGITS HAVE THE SHAPE OF THE INPUT. THEY MUST BE CONVERTED BACK TO THE ORIGINAL IMAGE SIZE.\n        SEE convert_predicted_logits_to_segmentation_with_correct_shape\n        \"\"\"\n        n_threads = torch.get_num_threads()\n        torch.set_num_threads(default_num_processes if default_num_processes < n_threads else n_threads)\n        prediction = None\n\n        for params in self.list_of_parameters:\n\n            # messing with state dict names...\n            if not isinstance(self.network, OptimizedModule):\n                self.network.load_state_dict(params)\n            else:\n                self.network._orig_mod.load_state_dict(params)\n\n            if prediction is None:\n                prediction = self.predict_sliding_window_return_logits(data).to(out_device)\n            else:\n                prediction += self.predict_sliding_window_return_logits(data).to(out_device)\n\n        if len(self.list_of_parameters) > 1:\n            prediction /= len(self.list_of_parameters)\n\n        if self.verbose: print('Prediction done')\n        torch.set_num_threads(n_threads)\n        return prediction\n\n    def _internal_get_sliding_window_slicers(self, image_size: Tuple[int, ...]):\n        slicers = []\n        if len(self.configuration_manager.patch_size) < len(image_size):\n            assert len(self.configuration_manager.patch_size) == len(\n                image_size) - 1, 'if tile_size has less entries than image_size, ' \\\n                                 'len(tile_size) ' \\\n                                 'must be one shorter than len(image_size) ' \\\n                                 '(only dimension ' \\\n                                 'discrepancy of 1 allowed).'\n            steps = compute_steps_for_sliding_window(image_size[1:], self.configuration_manager.patch_size,\n                                                     self.tile_step_size)\n            if self.verbose: print(f'n_steps {image_size[0] * len(steps[0]) * len(steps[1])}, image size is'\n                                   f' {image_size}, tile_size {self.configuration_manager.patch_size}, '\n                                   f'tile_step_size {self.tile_step_size}\\nsteps:\\n{steps}')\n            for d in range(image_size[0]):\n                for sx in steps[0]:\n                    for sy in steps[1]:\n                        slicers.append(\n                            tuple([slice(None), d, *[slice(si, si + ti) for si, ti in\n                                                     zip((sx, sy), self.configuration_manager.patch_size)]]))\n        else:\n            steps = compute_steps_for_sliding_window(image_size, self.configuration_manager.patch_size,\n                                                     self.tile_step_size)\n            if self.verbose: print(\n                f'n_steps {np.prod([len(i) for i in steps])}, image size is {image_size}, tile_size {self.configuration_manager.patch_size}, '\n                f'tile_step_size {self.tile_step_size}\\nsteps:\\n{steps}')\n            for sx in steps[0]:\n                for sy in steps[1]:\n                    for sz in steps[2]:\n                        slicers.append(\n                            tuple([slice(None), *[slice(si, si + ti) for si, ti in\n                                                  zip((sx, sy, sz), self.configuration_manager.patch_size)]]))\n        return slicers\n\n    @torch.inference_mode()\n    def _internal_predict_sliding_window_return_logits(\n        self,\n        data: torch.Tensor,\n        slicers,\n        do_on_device: bool = True,\n    ):\n        assert do_on_device, \"CPU execution disabled by design.\"\n\n        predicted_logits = None\n        results_device = self.device  # must be cuda\n        queue = Queue(maxsize=8)\n\n        def producer_cpu(sl_subset, q):\n            try:\n                for s in sl_subset:\n                    q.put(s)\n                q.put(\"end\")\n            except Exception:\n                # If consumer dies early, producer may throw; it's fine to exit.\n                pass\n\n        t = Thread(target=producer_cpu, args=(slicers, queue), daemon=True)\n        t.start()\n\n        try:\n            empty_cache(results_device)\n            # data remains on CPU; tiles move to GPU on demand\n\n            if self.verbose and not self.allow_tqdm:\n                print(f'running prediction: {len(slicers)} steps on {results_device}')\n\n            processed = 0\n            while True:\n                item = queue.get()\n                if item == \"end\":\n                    queue.task_done()\n                    break\n\n                sl = item\n                # Move just this tile to GPU\n                workon = torch.clone(\n                    data[sl][None], memory_format=torch.contiguous_format\n                ).to(results_device, non_blocking=True)\n\n                with torch.autocast(results_device.type, enabled=True):\n                    pred = self.network(workon)[0]  # GPU tensor\n\n                spatial_max = torch.amax(pred, dim=(1, 2, 3))  # [B, C], GPU\n                if predicted_logits is None:\n                    predicted_logits = spatial_max\n                else:\n                    predicted_logits = torch.maximum(predicted_logits, spatial_max)\n\n                # Aggressive cleanup per tile\n                del workon, pred, spatial_max, sl\n                queue.task_done()\n                processed += 1\n                if (processed & 0xF) == 0:\n                    empty_cache(results_device)\n\n            # Ensure queue fully drained\n            queue.join()\n\n            if predicted_logits is None:\n                raise RuntimeError(\"No tiles were processed; predicted_logits is None.\")\n\n            # Sanity check for infs\n            if torch.any(torch.isinf(predicted_logits)):\n                raise RuntimeError(\"Encountered inf in predicted logits.\")\n\n            return predicted_logits\n\n        except Exception as e:\n            # Mark queue so producer can exit if it's blocked\n            try:\n                queue.put_nowait(\"end\")\n            except Exception:\n                pass\n            raise  # propagate; no CPU fallback\n\n        finally:\n            # Always try to join the producer thread\n            try:\n                t.join(timeout=2.0)\n            except Exception:\n                pass\n\n            # Drain any leftover items (they are CPU slicers, not GPU tensors)\n            try:\n                while True:\n                    queue.get_nowait()\n                    queue.task_done()\n            except Empty:\n                pass\n            try:\n                queue.join()\n            except Exception:\n                pass\n\n            # Delete locals that may hold GPU refs\n            try:\n                del predicted_logits\n            except Exception:\n                pass\n\n            empty_cache(results_device)\n\n    @torch.inference_mode()\n    def predict_sliding_window_return_logits(self, input_image: torch.Tensor):\n        assert isinstance(input_image, torch.Tensor)\n        assert torch.cuda.is_available(), \"CUDA required.\"\n        assert self.device.type == 'cuda', \"CUDA required.\"\n    \n        n_cuda = torch.cuda.device_count()\n        # IMPORTANT:\n        # - Single-GPU: keep self.network on self.device (cuda:0)\n        # - Multi-GPU: keep self.network on CPU so we don't have 3 copies\n        if n_cuda >= 2:\n            # move any stray copy off GPU and free memory\n            try:\n                if next(self.network.parameters()).is_cuda:\n                    self.network.to(\"cpu\")\n                    torch.cuda.empty_cache()\n                    torch.cuda.ipc_collect()\n            except StopIteration:\n                pass\n            base_net = getattr(self, \"_net_uncompiled\", self.network)\n            if next(base_net.parameters()).is_cuda:\n                base_net.to(\"cpu\")\n            net_ready = self.network  # lives on CPU\n        else:\n            self.network = self.network.to(self.device).eval()\n            net_ready = self.network\n    \n        empty_cache(self.device)\n    \n        with torch.autocast(\"cuda\", enabled=True):\n            assert input_image.ndim == 4, 'input_image must be (c, x, y, z)'\n            data, _ = pad_nd_image(\n                input_image, self.configuration_manager.patch_size,\n                'constant', {'value': 0}, True, None\n            )\n            slicers = self._internal_get_sliding_window_slicers(data.shape[1:])\n    \n            if n_cuda >= 2:\n                if self.verbose:\n                    print('Using multi-GPU (2 devices).')\n                out = self._predict_sliding_window_multi_device(data, slicers)  # raises on error\n            else:\n                out = self._internal_predict_sliding_window_return_logits(data, slicers, True)  # raises on error\n    \n            return out\n\n    @torch.inference_mode()\n    def _internal_maybe_mirror_and_predict(self, x: torch.Tensor, net) -> torch.Tensor:\n        \"\"\"\n        Vectorized TTA: stack all flip variants in the batch, single forward pass,\n        then invert flips on outputs and average.\n        \"\"\"\n        # plain forward if TTA off\n        if not self.use_mirroring or not self.allowed_mirroring_axes:\n            return net(x)\n    \n        # Validate and convert spatial axes to tensor dims (offset by 2 for [B,C,...])\n        assert max(self.allowed_mirroring_axes) <= x.ndim - 3, \\\n            'mirror_axes does not match the dimension of the input!'\n        flip_dims = [a + 2 for a in self.allowed_mirroring_axes]\n    \n        # All non-empty combinations of flip dims (e.g., (d,), (e,), (d,e), ...)\n        combos = [c for k in range(1, len(flip_dims) + 1)\n                    for c in itertools.combinations(flip_dims, k)]\n    \n        B = x.shape[0]\n        device = x.device\n    \n        # Build augmented batch: original + all flips\n        variants = [x] + [torch.flip(x, dims=axes) for axes in combos]\n        aug = torch.cat(variants, dim=0)  # shape: [(1+N)*B, C, *spatial]\n    \n        # Single forward pass\n        y = net(aug)  # shape: [(1+N)*B, C_out, *spatial]\n    \n        # Reshape to [num_variants, B, C_out, *spatial]\n        num_variants = 1 + len(combos)\n        y = y.view(num_variants, B, *y.shape[1:])\n    \n        # Unflip outputs for flipped inputs, then average\n        pred = y[0]  # original orientation\n        for i, axes in enumerate(combos, start=1):\n            pred = pred + torch.flip(y[i], dims=axes)\n    \n        pred = pred / num_variants\n        return pred\n\n    \n    @torch.inference_mode()  # applies to the caller thread only; threads will set it again.\n    def _predict_sliding_window_multi_device(self, data: torch.Tensor, slicers):\n        dev_ids = [0, 1]\n        devices = [torch.device(f'cuda:{i}') for i in dev_ids]\n        agg_device = devices[0]\n    \n        mid = max(1, min(len(slicers) - 1, len(slicers) // 2))\n        slices_per_dev = [slicers[:mid], slicers[mid:]]\n\n        results, errors = {}, {}\n    \n        def worker(dev: torch.device, sl_subset, key: str):\n            net = None\n            local_preds_cpu = []  # list of [B,C] CPU tensors\n\n            if self.use_gaussian:\n                predicted_logits = torch.zeros((self.label_manager.num_segmentation_heads, *data.shape[1:]),\n                                           dtype=torch.half, device=dev)\n                n_predictions = torch.zeros(data.shape[1:], dtype=torch.half, device=dev)\n\n                gaussian = compute_gaussian(tuple(self.configuration_manager.patch_size), sigma_scale=1. / 8,\n                                                    value_scaling_factor=10, device=dev)\n            try:\n                torch.set_grad_enabled(False)\n                with torch.inference_mode():\n                    empty_cache(dev)\n    \n                    net = self._get_or_make_device_net(dev)\n    \n                    for idx, sl in enumerate(sl_subset):\n                        tile_cpu = torch.clone(\n                            data[sl][None], memory_format=torch.contiguous_format\n                        ).pin_memory()                         # page-lock -> faster DMA\n                        workon = tile_cpu.to(dev, non_blocking=True, dtype=torch.half)\n                        del tile_cpu\n    \n                        with torch.autocast(\"cuda\", enabled=True):\n                            #pred = net(workon)[0]          # [B,C,D,H,W] -> your net returns [B,C,...]\n                            pred = self._internal_maybe_mirror_and_predict(workon, net)[0]\n    \n                        if self.use_gaussian:\n                            pred *= gaussian\n                            predicted_logits[sl] += pred\n                            n_predictions[sl[1:]] += gaussian\n                            del workon, pred\n                        else:\n                            # Max val of each channel\n                            spatial_max = torch.amax(pred, dim=(1, 2, 3))  # [B, C] on dev\n                            # move tiny result to CPU and drop graph refs deterministically\n                            local_preds_cpu.append(spatial_max.detach().to(\"cpu\"))\n                            del workon, pred, spatial_max\n                            \n                        if (idx & 0xF) == 0:\n                            empty_cache(dev)\n    \n                if self.use_gaussian:\n                    results[key] = {}\n                    if (predicted_logits == 0).all():\n                        results[key][\"pred\"] = None\n                        results[key][\"n\"] = None\n                    else:\n                        results[key] = {}\n                        results[key][\"pred\"] = predicted_logits.detach().cpu()\n                        results[key][\"n\"] = n_predictions.detach().cpu()\n                else:\n                    # Reduce on CPU with torch (no numpy) to avoid grad issues\n                    if len(local_preds_cpu) == 0:\n                        results[key] = None\n                    elif len(local_preds_cpu) == 1:\n                        results[key] = local_preds_cpu[0]          # [B,C] CPU\n                    else:\n                        results[key] = torch.stack(local_preds_cpu, dim=0).amax(dim=0)  # [B,C] CPU\n    \n            except Exception as ex:\n                # Only strings; don't keep exception objects that can capture frames\n                import traceback\n                errors[key] = (str(ex), traceback.format_exc())\n            finally:\n                # free all refs\n                try: del net\n                except: pass\n                try:\n                    del local_preds_cpu\n                except: pass\n                try:\n                    torch.cuda.synchronize(dev)\n                    torch.cuda.set_device(dev)\n                    torch.cuda.empty_cache()\n                    torch.cuda.ipc_collect()\n                except: pass\n    \n        # launch & join\n        threads = []\n        for dev, subset, key in zip(devices, slices_per_dev, ['dev0', 'dev1']):\n            t = Thread(target=worker, args=(dev, subset, key), daemon=True)\n            t.start()\n            threads.append(t)\n        for t in threads:\n            t.join()\n    \n        if errors:\n            # proactively drop per-device nets to free VRAM before raising\n            self._clear_device_nets()\n            key, (msg, tb) = next(iter(errors.items()))\n            raise RuntimeError(f\"Worker {key} failed: {msg}\\n{tb}\")\n        \n        r0, r1 = results.get('dev0', None), results.get('dev1', None)\n        if r0 is None and r1 is None:\n            raise RuntimeError('No predictions produced.')\n    \n        if self.use_gaussian:\n            predicted_logits = r0[\"pred\"] if r1[\"pred\"] is None else (r1[\"pred\"] if r0[\"pred\"] is None else r0[\"pred\"] + r1[\"pred\"] )\n            n_predictions = r0[\"n\"] if r1[\"n\"] is None else (r1[\"n\"] if r0[\"n\"] is None else r0[\"n\"] + r1[\"n\"] )\n            torch.div(predicted_logits, n_predictions, out=predicted_logits)\n            out_cpu = predicted_logits\n            del predicted_logits, n_predictions\n        else:\n            # CPU combine (tiny [B,C]), then move once to agg_device\n            out_cpu = r0 if r1 is None else (r1 if r0 is None else torch.maximum(r0, r1))\n\n        out = out_cpu.to(agg_device, non_blocking=True)\n    \n        try:\n            del results, r0, r1, out_cpu\n        except Exception:\n            pass\n        try:\n            torch.cuda.set_device(agg_device)\n            torch.cuda.empty_cache()\n            torch.cuda.ipc_collect()\n        except Exception:\n            pass\n        return out","metadata":{"_uuid":"772c29ea-7aff-49cb-9283-386731b314a7","_cell_guid":"5a814085-d630-4bf8-81f3-11af79cc833f","trusted":true,"execution":{"iopub.status.busy":"2025-10-20T15:18:05.916435Z","iopub.execute_input":"2025-10-20T15:18:05.917262Z","iopub.status.idle":"2025-10-20T15:18:05.930726Z","shell.execute_reply.started":"2025-10-20T15:18:05.917228Z","shell.execute_reply":"2025-10-20T15:18:05.929814Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Using a worker with which we communicate because we cannot overwrite python packages which were already provided within the kernel. \nHowever, launching withing a subprocess allows us using the overwritten packages.","metadata":{}},{"cell_type":"code","source":"%%writefile /kaggle/working/worker.py\nimport sys, os, json, traceback\nfrom pathlib import Path\n\n# Lazy imports inside worker process\nimport numpy as np\nimport pandas as pd\nimport torch\nimport SimpleITK as sitk\n\nfrom argparse import Namespace\n\nfrom nnUNetPredictor import nnUNetPredictor\nfrom nnunetv2.dataset_conversion.kaggle_2025_rsna.official_data_to_nnunet import load_and_crop\n# ----------------- Globals (cached across requests) -----------------\npredictor = None\npreprocessor = None\nargs = None\n\ndef log(msg):\n    sys.stderr.write(str(msg) + \"\\n\"); sys.stderr.flush()\n\ndef jwrite(obj):\n    sys.stdout.write(json.dumps(obj) + \"\\n\")\n    sys.stdout.flush()\n\n@torch.no_grad()\ndef _lazy_init(local_args: dict):\n    global predictor, preprocessor, args\n    if args is None:\n        args = Namespace(**local_args)\n        # misc speed envs\n        os.environ['nnUNet_compile'] = 'True'\n        os.environ['torch.backends.cudnn.benchmark'] = 'True'\n    if predictor is None:\n        log(\"Initializing predictor...\")\n        f = [i if i == 'all' else int(i) for i in args.f]\n        device = torch.device(args.device if torch.cuda.is_available() else \"cpu\")\n        pred = nnUNetPredictor(\n            tile_step_size=args.step_size,\n            use_gaussian=args.use_gaussian,\n            use_mirroring=not args.disable_tta,\n            device=device,\n            verbose=args.verbose,\n            verbose_preprocessing=args.verbose,\n            allow_tqdm=not args.disable_progress_bar\n        )\n        pred.initialize_from_trained_model_folder(args.m, f, checkpoint_name=args.chk)\n        predictor = pred\n    if preprocessor is None:\n        log(\"Initializing preprocessor...\")\n        preprocessor = predictor.configuration_manager.preprocessor_class()\n\nimport time\nimport torch\n    \n@torch.no_grad()\ndef run_inference(series_dir: Path) -> dict:\n    print(series_dir.name)\n    timings = {}\n    t_total = time.perf_counter()\n    \n    # 1) crop with regressor\n    t = time.perf_counter()\n    img, properties = load_and_crop(series_dir)\n    timings[\"crop_bbox\"] = time.perf_counter() - t\n    \n    # 2) preprocessing\n    t = time.perf_counter()\n    data, _, _ = preprocessor.run_case_npy(\n        np.array([img]),\n        None,\n        properties,\n        predictor.plans_manager,\n        predictor.configuration_manager,\n        predictor.dataset_json\n    )\n\n    timings[\"preprocess\"] = time.perf_counter() - t\n    \n    # 3) logits (inference)\n    t = time.perf_counter()\n    logits = predictor.predict_logits_from_preprocessed_data(torch.from_numpy(data)).cpu()\n    timings[\"inference_logits\"] = time.perf_counter() - t\n    \n    # 4) probs (postproc)\n    probs = torch.sigmoid(logits)\n    \n    timings[\"total\"] = time.perf_counter() - t_total\n    \n    print(\n        \"Timing (s): \"\n        f\"crop_bbox={timings['crop_bbox']:.3f}, \"\n        f\"preprocess={timings['preprocess']:.3f}, \"\n        f\"inference_logits={timings['inference_logits']:.3f}, \"\n        f\"TOTAL={timings['total']:.3f}\"\n    )\n\n    if predictor.use_gaussian:\n        max_per_c = torch.amax(probs, dim=(1, 2, 3)).to(dtype=torch.float32, device=\"cpu\")\n    else:\n        max_per_c = probs\n\n    vals = max_per_c.numpy().tolist()                 \n    labels = list(predictor.dataset_json['labels'].keys())[1:] + ['Aneurysm Present']\n    if len(vals) == 13:\n        aneurysm_present = float(max(vals))                \n        data_row = vals + [aneurysm_present]               \n        return {\"columns\": [str(x) for x in labels], \"data\": data_row}\n    else: \n        return {\"columns\": [str(x) for x in labels], \"data\": vals}\n\ndef main():\n    # handshake\n    jwrite({\"ok\": True, \"event\": \"worker_started\"})\n    for line in sys.stdin:\n        line = line.strip()\n        if not line:\n            continue\n        try:\n            req = json.loads(line)\n            cmd = req.get(\"cmd\")\n            if cmd == \"init\":\n                _lazy_init(req[\"args\"])\n                jwrite({\"ok\": True})\n            elif cmd == \"predict\":\n                res = run_inference(Path(req[\"series_path\"]))\n                jwrite({\"ok\": True, \"result\": res})\n            else:\n                jwrite({\"ok\": False, \"error\": f\"unknown cmd: {cmd}\"})\n        except Exception as e:\n            jwrite({\"ok\": False, \"error\": str(e), \"trace\": traceback.format_exc()})\n\nif __name__ == \"__main__\":\n    main()","metadata":{"_uuid":"1f7f9d8f-024c-430b-9ca7-eedb21a3cdd4","_cell_guid":"6da4cb98-a5ab-4f73-bce3-77601599f096","trusted":true,"execution":{"iopub.status.busy":"2025-10-20T15:18:06.017329Z","iopub.execute_input":"2025-10-20T15:18:06.017526Z","iopub.status.idle":"2025-10-20T15:18:06.023624Z","shell.execute_reply.started":"2025-10-20T15:18:06.017512Z","shell.execute_reply":"2025-10-20T15:18:06.022822Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"proxy = None\n\ndef _ensure_worker_started(init_args: dict):\n    \"\"\"\n    Lazily create the worker venv, install offline wheels, write worker.py,\n    and start the long-lived worker process. Cache predictor in the worker.\n    \"\"\"\n    global proxy\n    if proxy is not None:\n        return\n\n    # --- Parent<->Worker bridge (JSON-L protocol) ---\n    import subprocess, json, threading, queue, sys, os\n    from pathlib import Path\n    \n    class WorkerProxy:\n        def __init__(self):\n            # text=True -> str lines; bufsize=1 -> line-buffered\n            self.p = subprocess.Popen(\n                ['python', '/kaggle/working/worker.py'],\n                stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE,\n                text=True, bufsize=1\n            )\n            self.q = queue.Queue()\n    \n            # stdout reader: only enqueue parsed dicts\n            self._t_out = threading.Thread(target=self._reader_stdout, daemon=True)\n            self._t_out.start()\n    \n            # stderr reader: just mirror logs to parent's stderr so we see them\n            self._t_err = threading.Thread(target=self._reader_stderr, daemon=True)\n            self._t_err.start()\n    \n            hello = self.recv(timeout=600)\n            if not isinstance(hello, dict) or not hello.get(\"ok\"):\n                raise RuntimeError(f\"Worker failed to start: {hello}\")\n\n        def shutdown(self):\n            try:\n                self.p.terminate()\n                self.p.wait(timeout=5)\n            except Exception:\n                try:\n                    self.p.kill()\n                except Exception:\n                    pass\n        def _reader_stdout(self):\n            for line in self.p.stdout:\n                line = line.strip()\n                if not line:\n                    continue\n                try:\n                    obj = json.loads(line)\n                    if isinstance(obj, dict):\n                        self.q.put(obj)\n                    else:\n                        # ignore non-dict JSON (e.g., \"null\", numbers, strings)\n                        sys.stderr.write(f\"[worker-nonjson-stdout] {line}\\n\"); sys.stderr.flush()\n                except Exception:\n                    sys.stderr.write(f\"[worker-nonjson-stdout] {line}\\n\"); sys.stderr.flush()\n    \n        def _reader_stderr(self):\n            # forward worker's logs\n            for line in self.p.stderr:\n                sys.stderr.write(f\"[worker] {line}\")\n            sys.stderr.flush()\n    \n        def call(self, obj, timeout=600):\n            try:\n                self.p.stdin.write(json.dumps(obj) + \"\\n\")\n                self.p.stdin.flush()\n            except Exception as e:\n                raise RuntimeError(f\"Failed to send to worker: {e}\")\n            return self.recv(timeout)\n    \n        def recv(self, timeout=600):\n            try:\n                msg = self.q.get(timeout=timeout)\n                if not isinstance(msg, dict):\n                    raise RuntimeError(f\"Bad worker message: {msg!r}\")\n                return msg\n            except queue.Empty:\n                raise TimeoutError(\"Worker timed out\")\n\n\n\n    proxy = WorkerProxy()\n    # send init with your nnU-Net paths/args\n    init_ok = proxy.call({\"cmd\": \"init\", \"args\": init_args}, timeout=1800)\n    if not init_ok.get(\"ok\"):\n        raise RuntimeError(f\"worker init failed: {init_ok}\")","metadata":{"_uuid":"2d6b852d-a1ff-4a76-a8fe-b6d464740888","_cell_guid":"ddd2edc4-9df5-4b66-b675-c905abd5bf28","trusted":true,"execution":{"iopub.status.busy":"2025-10-20T15:18:06.363388Z","iopub.execute_input":"2025-10-20T15:18:06.363638Z","iopub.status.idle":"2025-10-20T15:18:06.374329Z","shell.execute_reply.started":"2025-10-20T15:18:06.363619Z","shell.execute_reply":"2025-10-20T15:18:06.373523Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Inference","metadata":{}},{"cell_type":"markdown","source":"Definition of the model to use and implementing the predict function.","metadata":{}},{"cell_type":"code","source":"INIT_ARGS = {\n    'm': '/kaggle/input/rsna-2025-7th-place-checkpoint/Dataset004_iarsna_crop/Kaggle2025RSNATrainer__nnUNetResEncUNetMPlans__3d_fullres_bs32',\n    'f': ['all'],                 \n    'step_size': 0.5,\n    'disable_tta': True,\n    'use_gaussian': False,\n    'chk': 'checkpoint_epoch_1500.pth',\n    'verbose': False,\n    'save_probabilities': False,\n    'continue_prediction': False,\n    'device': 'cuda',\n    'disable_progress_bar': False,\n}\nproxy = None\nproxy_failures = 0\n\nimport shutil\ndef predict(series_path: str) -> pd.DataFrame:\n    global proxy, LABEL_COLS, proxy_failures\n    \n    _ensure_worker_started(INIT_ARGS)  \n    out_dir = \"/kaggle/working/out\"   \n    os.makedirs(out_dir, exist_ok=True)\n    df = pl.DataFrame([[0.5] * len(LABEL_COLS)], schema=LABEL_COLS,  orient=\"row\")\n\n    try: \n        r = proxy.call({\"cmd\": \"predict\", \"series_path\": series_path, \"out_dir\": out_dir}, timeout=1200)\n    \n        if not isinstance(r, dict):\n            raise RuntimeError(f\"Bad response from worker (not a dict): {r!r}\")\n        \n        if not r.get(\"ok\"):\n            raise RuntimeError(f\"worker error: {r.get('error')}\\n{r.get('trace','')}\")\n\n        payload = r[\"result\"]\n        df = pl.DataFrame([payload[\"data\"]], schema=payload[\"columns\"],  orient=\"row\")\n        label_cols_present = [c for c in LABEL_COLS if c in df.columns]\n        df = df.select(\n            [*[pl.col(c) for c in label_cols_present], pl.all().exclude(LABEL_COLS)]\n        )\n        proxy_failures = 0\n    except Exception as e:\n        if proxy_failures >= 2:\n            # hard reset the worker process -> frees ALL its CUDA memory\n            try:\n                print('Restarting proxy, too many consecutive errors.')\n                proxy.shutdown()\n            except Exection as e: \n                print('Even that didn`t work... thing look bad...')\n            finally:\n                proxy_failures = 0\n                proxy = None\n        else:\n            print('Something went wrong, returning default prediction.')\n            print(e)\n            proxy_failures += 1\n        \n    # Required disk cleanup (do NOT remove the venv)\n    shutil.rmtree('/kaggle/shared', ignore_errors=True)\n    # Optionally clean prediction outputs if they get large:\n    shutil.rmtree(out_dir, ignore_errors=True)\n    return df","metadata":{"_uuid":"34f8d86a-ae30-4735-8f53-457ab9ade215","_cell_guid":"b8d2326d-6670-4908-852b-904eb2bd3e77","trusted":true,"collapsed":false,"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"shutil.rmtree('/kaggle/shared', ignore_errors=True)\ninference_server = kaggle_evaluation.rsna_inference_server.RSNAInferenceServer(predict)\n\nif os.getenv('KAGGLE_IS_COMPETITION_RERUN'):\n    inference_server.serve()\nelse:\n    inference_server.run_local_gateway()\n    display(pl.read_parquet('/kaggle/working/submission.parquet'))","metadata":{"_uuid":"199e2c60-6f94-425f-a06e-0dc823dde0aa","_cell_guid":"6bf9effd-735d-4a57-953e-fa966df3694b","trusted":true,"collapsed":false,"execution":{"iopub.status.busy":"2025-10-20T15:19:01.434116Z","iopub.execute_input":"2025-10-20T15:19:01.434955Z","iopub.status.idle":"2025-10-20T15:21:19.77437Z","shell.execute_reply.started":"2025-10-20T15:19:01.434923Z","shell.execute_reply":"2025-10-20T15:21:19.773724Z"},"jupyter":{"outputs_hidden":false}},"outputs":[],"execution_count":null}]}