{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":84969,"databundleVersionId":10033515},{"sourceType":"datasetVersion","sourceId":9869730,"datasetId":6058495,"databundleVersionId":10122655},{"sourceType":"datasetVersion","sourceId":10250551,"datasetId":6340217,"databundleVersionId":10547451},{"sourceType":"datasetVersion","sourceId":10221432,"datasetId":6318762,"databundleVersionId":10514794},{"sourceType":"datasetVersion","sourceId":9862305,"datasetId":6052780,"databundleVersionId":10114338},{"sourceType":"datasetVersion","sourceId":9979164,"datasetId":6126229,"databundleVersionId":10244954},{"sourceType":"datasetVersion","sourceId":10642759,"datasetId":6338199,"databundleVersionId":10984436},{"sourceType":"datasetVersion","sourceId":10256507,"datasetId":6344718,"databundleVersionId":10554034},{"sourceType":"datasetVersion","sourceId":10229415,"datasetId":6293204,"databundleVersionId":10523661},{"sourceType":"datasetVersion","sourceId":10225912,"datasetId":6155493,"databundleVersionId":10519728},{"sourceType":"datasetVersion","sourceId":10467917,"datasetId":6430595,"databundleVersionId":10791228},{"sourceType":"datasetVersion","sourceId":10476214,"datasetId":6486959,"databundleVersionId":10800706},{"sourceType":"datasetVersion","sourceId":10624471,"datasetId":6578264,"databundleVersionId":10964198},{"sourceType":"kernelVersion","sourceId":206165222},{"sourceType":"kernelVersion","sourceId":220221133}],"dockerImageVersionId":30840,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## Installing offline deps\n\nAs this is a code comp, there is no internet. \nSo we have to do some silly things to get dependencies in here. \nWhy is asciitree such a PITA? ","metadata":{}},{"cell_type":"code","source":"deps_path = '/kaggle/input/czii-cryoet-dependencies'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-02T17:47:17.511532Z","iopub.execute_input":"2025-02-02T17:47:17.512032Z","iopub.status.idle":"2025-02-02T17:47:17.516241Z","shell.execute_reply.started":"2025-02-02T17:47:17.511990Z","shell.execute_reply":"2025-02-02T17:47:17.515292Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"! cp -r /kaggle/input/czii-cryoet-dependencies/asciitree-0.3.3/ asciitree-0.3.3/","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-02T17:47:17.517521Z","iopub.execute_input":"2025-02-02T17:47:17.517881Z","iopub.status.idle":"2025-02-02T17:47:17.723113Z","shell.execute_reply.started":"2025-02-02T17:47:17.517849Z","shell.execute_reply":"2025-02-02T17:47:17.721630Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"! pip wheel asciitree-0.3.3/asciitree-0.3.3/\n","metadata":{"_kg_hide-output":true,"trusted":true,"execution":{"iopub.status.busy":"2025-02-02T17:47:17.725447Z","iopub.execute_input":"2025-02-02T17:47:17.725881Z","iopub.status.idle":"2025-02-02T17:47:21.184438Z","shell.execute_reply.started":"2025-02-02T17:47:17.725841Z","shell.execute_reply":"2025-02-02T17:47:21.183150Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install asciitree-0.3.3-py3-none-any.whl","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-02T17:47:21.186071Z","iopub.execute_input":"2025-02-02T17:47:21.186349Z","iopub.status.idle":"2025-02-02T17:47:25.749011Z","shell.execute_reply.started":"2025-02-02T17:47:21.186323Z","shell.execute_reply":"2025-02-02T17:47:25.747885Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"! pip install -q --no-index --find-links {deps_path} --requirement {deps_path}/requirements.txt","metadata":{"_kg_hide-output":true,"trusted":true,"execution":{"iopub.status.busy":"2025-02-02T17:47:25.750301Z","iopub.execute_input":"2025-02-02T17:47:25.750558Z","iopub.status.idle":"2025-02-02T17:47:34.235014Z","shell.execute_reply.started":"2025-02-02T17:47:25.750534Z","shell.execute_reply":"2025-02-02T17:47:34.233931Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install /kaggle/input/tensorrt-10-1-0/nvidia_cuda_runtime_cu12-12.2.140-py3-none-manylinux1_x86_64.whl\n!pip install /kaggle/input/tensorrt-10-1-0/tensorrt_cu12_bindings-10.1.0-cp310-none-manylinux_2_17_x86_64.whl\n!pip install /kaggle/input/tensorrt-10-1-0/tensorrt_cu12_libs-10.1.0-py2.py3-none-manylinux_2_17_x86_64.whl\n!pip install /kaggle/input/tensorrt-10-1-0/tensorrt_cu12-10.1.0-py2.py3-none-any.whl\n!pip install /kaggle/input/tensorrt-10-1-0/tensorrt-10.1.0-py2.py3-none-any.whl\n!pip install /kaggle/input/tensorrt-10-1-0/polygraphy-0.49.14-py2.py3-none-any.whl","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-02T17:47:34.236111Z","iopub.execute_input":"2025-02-02T17:47:34.236499Z","iopub.status.idle":"2025-02-02T17:48:33.169197Z","shell.execute_reply.started":"2025-02-02T17:47:34.236463Z","shell.execute_reply":"2025-02-02T17:48:33.168048Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from typing import List, Tuple, Union\nimport numpy as np\nimport torch\nfrom monai.data import DataLoader, Dataset, CacheDataset, decollate_batch\nfrom monai.transforms import (\n    Compose, \n    EnsureChannelFirstd, \n    Orientationd,  \n    AsDiscrete,  \n    RandFlipd, \n    RandRotate90d, \n    NormalizeIntensityd,\n    RandCropByLabelClassesd,\n)","metadata":{"execution":{"iopub.status.busy":"2025-02-02T17:48:33.170349Z","iopub.execute_input":"2025-02-02T17:48:33.170584Z","iopub.status.idle":"2025-02-02T17:49:02.037711Z","shell.execute_reply.started":"2025-02-02T17:48:33.170562Z","shell.execute_reply":"2025-02-02T17:49:02.036627Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import sys\nsys.path.append('/kaggle/input/czii-ckpt/')\nfrom UNetWithMem import UNetWithMem\nsys.path.append('/kaggle/input/czii-weight-by-fbeta')\nfrom UNet3d import ResUNet3D\nfrom DLinkNet3D import DLinkNet3D\nsys.path.append('/kaggle/input/czii-radius-0-5')\nfrom UNetWithMemRes import UNetWithMemRes","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-02T17:49:02.040311Z","iopub.execute_input":"2025-02-02T17:49:02.041141Z","iopub.status.idle":"2025-02-02T17:49:02.109328Z","shell.execute_reply.started":"2025-02-02T17:49:02.041108Z","shell.execute_reply":"2025-02-02T17:49:02.108341Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Define some helper functions\n\n\n### Patching helper functions\n\nThese are mostly used to split large volumes into smaller ones and stitch them back together. ","metadata":{}},{"cell_type":"code","source":"def calculate_patch_starts(dimension_size: int, patch_size: int) -> List[int]:\n    \"\"\"\n    Calculate the starting positions of patches along a single dimension\n    with minimal overlap to cover the entire dimension.\n    \n    Parameters:\n    -----------\n    dimension_size : int\n        Size of the dimension\n    patch_size : int\n        Size of the patch in this dimension\n        \n    Returns:\n    --------\n    List[int]\n        List of starting positions for patches\n    \"\"\"\n    if dimension_size <= patch_size:\n        return [0]\n        \n    # Calculate number of patches needed\n    n_patches = np.ceil(dimension_size / patch_size)\n    \n    if n_patches == 1:\n        return [0]\n    \n    # Calculate overlap\n    total_overlap = (n_patches * patch_size - dimension_size) / (n_patches - 1)\n    \n    # Generate starting positions\n    positions = []\n    for i in range(int(n_patches)):\n        pos = int(i * (patch_size - total_overlap))\n        if pos + patch_size > dimension_size:\n            pos = dimension_size - patch_size\n        if pos not in positions:  # Avoid duplicates\n            positions.append(pos)\n    \n    return positions\n\ndef extract_3d_patches_minimal_overlap(arrays: List[np.ndarray], patch_size: int) -> Tuple[List[np.ndarray], List[Tuple[int, int, int]]]:\n    \"\"\"\n    Extract 3D patches from multiple arrays with minimal overlap to cover the entire array.\n    \n    Parameters:\n    -----------\n    arrays : List[np.ndarray]\n        List of input arrays, each with shape (m, n, l)\n    patch_size : int\n        Size of cubic patches (a x a x a)\n        \n    Returns:\n    --------\n    patches : List[np.ndarray]\n        List of all patches from all input arrays\n    coordinates : List[Tuple[int, int, int]]\n        List of starting coordinates (x, y, z) for each patch\n    \"\"\"\n    if not arrays or not isinstance(arrays, list):\n        raise ValueError(\"Input must be a non-empty list of arrays\")\n    \n    # Verify all arrays have the same shape\n    shape = arrays[0].shape\n    if not all(arr.shape == shape for arr in arrays):\n        raise ValueError(\"All input arrays must have the same shape\")\n    \n    if patch_size > min(shape):\n        raise ValueError(f\"patch_size ({patch_size}) must be smaller than smallest dimension {min(shape)}\")\n    \n    m, n, l = shape\n    patches = []\n    coordinates = []\n    \n    # Calculate starting positions for each dimension\n    x_starts = calculate_patch_starts(m, patch_size)\n    y_starts = calculate_patch_starts(n, patch_size)\n    z_starts = calculate_patch_starts(l, patch_size)\n    \n    # Extract patches from each array\n    for arr in arrays:\n        for x in x_starts:\n            for y in y_starts:\n                for z in z_starts:\n                    patch = arr[\n                        x:x + patch_size,\n                        y:y + patch_size,\n                        z:z + patch_size\n                    ]\n                    patches.append(patch)\n                    coordinates.append((x, y, z))\n    \n    return patches, coordinates\n\n# Note: I should probably averge the overlapping areas, \n# but here they are just overwritten by the most recent one. \n\ndef reconstruct_array(patches: List[np.ndarray], \n                     coordinates: List[Tuple[int, int, int]], \n                     original_shape: Tuple[int, int, int]) -> np.ndarray:\n    \"\"\"\n    Reconstruct array from patches.\n    \n    Parameters:\n    -----------\n    patches : List[np.ndarray]\n        List of patches to reconstruct from\n    coordinates : List[Tuple[int, int, int]]\n        Starting coordinates for each patch\n    original_shape : Tuple[int, int, int]\n        Shape of the original array\n        \n    Returns:\n    --------\n    np.ndarray\n        Reconstructed array\n    \"\"\"\n    reconstructed = np.zeros(original_shape, dtype=np.int64)  # To track overlapping regions\n    \n    patch_size = patches[0].shape[0]\n    \n    for patch, (x, y, z) in zip(patches, coordinates):\n        reconstructed[\n            x:x + patch_size,\n            y:y + patch_size,\n            z:z + patch_size\n        ] = patch\n        \n    \n    return reconstructed","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2025-02-02T17:49:02.111218Z","iopub.execute_input":"2025-02-02T17:49:02.111642Z","iopub.status.idle":"2025-02-02T17:49:02.123699Z","shell.execute_reply.started":"2025-02-02T17:49:02.111596Z","shell.execute_reply":"2025-02-02T17:49:02.122573Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Submission helper functions\n\nThese help with getting the submission in the correct format","metadata":{}},{"cell_type":"code","source":"import pandas as pd\n\ndef dict_to_df(coord_dict, experiment_name):\n    \"\"\"\n    Convert dictionary of coordinates to pandas DataFrame.\n    \n    Parameters:\n    -----------\n    coord_dict : dict\n        Dictionary where keys are labels and values are Nx3 coordinate arrays\n        \n    Returns:\n    --------\n    pd.DataFrame\n        DataFrame with columns ['x', 'y', 'z', 'label']\n    \"\"\"\n    # Create lists to store data\n    all_coords = []\n    all_labels = []\n    \n    # Process each label and its coordinates\n    for label, coords in coord_dict.items():\n        all_coords.append(coords)\n        all_labels.extend([label] * len(coords))\n    \n    # Concatenate all coordinates\n    all_coords = np.vstack(all_coords)\n    \n    df = pd.DataFrame({\n        'experiment': experiment_name,\n        'particle_type': all_labels,\n        'x': all_coords[:, 0],\n        'y': all_coords[:, 1],\n        'z': all_coords[:, 2]\n    })\n\n    \n    return df","metadata":{"_kg_hide-input":true,"execution":{"iopub.status.busy":"2025-02-02T17:49:02.124737Z","iopub.execute_input":"2025-02-02T17:49:02.125066Z","iopub.status.idle":"2025-02-02T17:49:02.141682Z","shell.execute_reply.started":"2025-02-02T17:49:02.125040Z","shell.execute_reply":"2025-02-02T17:49:02.140662Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Reading in the data","metadata":{}},{"cell_type":"code","source":"TRAIN_DATA_DIR = \"/kaggle/input/create-numpy-dataset\"\nTEST_DATA_DIR = \"/kaggle/input/czii-cryo-et-object-identification\"","metadata":{"execution":{"iopub.status.busy":"2025-02-02T17:49:02.142670Z","iopub.execute_input":"2025-02-02T17:49:02.143063Z","iopub.status.idle":"2025-02-02T17:49:02.158057Z","shell.execute_reply.started":"2025-02-02T17:49:02.143026Z","shell.execute_reply":"2025-02-02T17:49:02.157109Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Initialize the model\n\nThis model is pretty much directly copied from [3D U-Net PyTorch Lightning distributed training](https://www.kaggle.com/code/zhuowenzhao11/3d-u-net-pytorch-lightning-distributed-training)","metadata":{}},{"cell_type":"code","source":"cp -r /kaggle/input/tensorrt-10-1-0/torch2trt-master /kaggle/working/torch2trt","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-02T17:49:02.158840Z","iopub.execute_input":"2025-02-02T17:49:02.159119Z","iopub.status.idle":"2025-02-02T17:49:03.024521Z","shell.execute_reply.started":"2025-02-02T17:49:02.159095Z","shell.execute_reply":"2025-02-02T17:49:03.023239Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install /kaggle/working/torch2trt","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-02T17:49:03.025850Z","iopub.execute_input":"2025-02-02T17:49:03.026173Z","iopub.status.idle":"2025-02-02T17:49:16.039306Z","shell.execute_reply.started":"2025-02-02T17:49:03.026142Z","shell.execute_reply":"2025-02-02T17:49:16.038239Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from torch2trt import TRTModule\nimport tensorrt as trt\nfrom cuda import cudart\n\nLOGGER = trt.Logger(trt.Logger.INFO)\ntrt.init_libnvinfer_plugins(LOGGER, \"\")\n\nclass Net:\n    def __init__(self, weights, device=0):\n        self.device = device\n        \n        torch.cuda.set_device(self.device)\n        \n        self.runtime = trt.Runtime(LOGGER)\n        with open(weights, \"rb\") as f:\n            self.engine = self.runtime.deserialize_cuda_engine(f.read())\n        \n        self.trt_model = TRTModule(\n            input_names=['images'],\n            output_names=['output'],\n            engine=self.engine\n        )\n    \n    def __call__(self, img):\n        if img.device.index != self.device:\n            img = img.cuda(self.device)\n        \n        with torch.no_grad():\n            output = self.trt_model(img)\n        return output\n\n        \nclass EnsembleModel(torch.nn.Module):\n    def __init__(self, models):\n        super(EnsembleModel, self).__init__()\n        self.models  = torch.nn.ModuleList(models)\n    \n    def forward(self, x):\n        with torch.no_grad():\n            with torch.amp.autocast('cuda'):\n                outputs = [model(x) for model in self.models]\n        return sum(outputs) / len(outputs)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-02T17:49:16.040572Z","iopub.execute_input":"2025-02-02T17:49:16.041156Z","iopub.status.idle":"2025-02-02T17:49:16.171756Z","shell.execute_reply.started":"2025-02-02T17:49:16.041117Z","shell.execute_reply":"2025-02-02T17:49:16.170831Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport tensorrt as trt\nfrom cuda import cudart\nimport torch.nn as nn\n\nclass TRTEnsembleModel(nn.Module):\n    def __init__(self, models, weights):\n        super(TRTEnsembleModel, self).__init__()\n        self.models = []\n        self.weights = weights\n        for model_config in models:\n            self.models.append(Net(weights=model_config['engine_path'], device=model_config['device']))\n    \n    def forward(self, x):\n        with torch.no_grad():\n            outputs = [model(x) * weight for model, weight in zip(self.models, self.weights)]\n        return sum(outputs)\n\nclass TRTEnsembleModel_diff_channel(nn.Module):\n    def __init__(self, models_7_channel, models_8_channel, weights1, weights2, new_e_model, new_e_model_weight):\n        super(TRTEnsembleModel_diff_channel, self).__init__()\n        self.models_7_channel = []\n        self.models_8_channel = []\n        self.selected_channels = [0, 1, 3, 4, 5, 6]\n        self.weights1 = weights1\n        self.weights2 = weights2\n        self.new_e_model_weight = new_e_model_weight\n        \n        # 7-channel models\n        for model_config in models_7_channel:\n            self.models_7_channel.append(Net(weights=model_config['engine_path'], device=model_config['device']))\n            \n        # 8-channel models\n        for model_config in models_8_channel:\n            self.models_8_channel.append(Net(weights=model_config['engine_path'], device=model_config['device']))\n            \n        # New ensemble model\n        self.new_e_model = Net(weights=new_e_model['engine_path'], device=new_e_model['device'])\n    \n    def forward(self, x):\n        with torch.no_grad():\n            outputs1 = [model(x)[:, self.selected_channels, :, :, :] * weight \n                       for model, weight in zip(self.models_7_channel, self.weights1)]\n            outputs2 = [model(x)[:, self.selected_channels, :, :, :] * weight \n                       for model, weight in zip(self.models_8_channel, self.weights2)]\n            outputs3 = self.new_e_model(x) * self.new_e_model_weight\n\n        all_outputs = outputs1 + outputs2\n        return sum(all_outputs) + outputs3\n\ndef load_trt_model(weight1, weight2, new_e_model, new_e_model_weight):\n    \"\"\"\n    weight1, weight2: リストで、各要素は辞書で以下を含む:\n    {\n        'engine_path': str,  # .engineファイルのパス\n        'device': int        # GPUデバイス番号\n    }\n    \"\"\"\n    final_model = TRTEnsembleModel_diff_channel(\n        models_7_channel=weight1,\n        models_8_channel=weight2,\n        weights1=[0.05+0.0125],  # 元の実装の重みを維持\n        weights2=[0.05+0.0125, 0.05+0.0125, 0.05+0.0125, 0.1+0.0125],  # 元の実装の重みを維持\n        new_e_model=new_e_model,\n        new_e_model_weight=new_e_model_weight\n    )\n    return final_model\n\ndef load_trt_model_same_c(weight):\n    \"\"\"\n    weight: リストで、各要素は辞書で以下を含む:\n    {\n        'engine_path': str,  # .engineファイルのパス\n        'device': int        # GPUデバイス番号\n    }\n    \"\"\"\n    ens_weights = [0.05+0.0125, 0.1+0.0125, 0.05+0.0125]  # 元の実装の重みを維持\n    final_model = TRTEnsembleModel(weight, ens_weights)\n    return final_model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-02T17:49:16.172731Z","iopub.execute_input":"2025-02-02T17:49:16.173027Z","iopub.status.idle":"2025-02-02T17:49:16.185029Z","shell.execute_reply.started":"2025-02-02T17:49:16.173003Z","shell.execute_reply":"2025-02-02T17:49:16.184029Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 同じチャンネル数のモデルのアンサンブル\nweight = [\n    # {'engine_path': '/path/to/model1.engine', 'device': 0},\n    {'engine_path': '/kaggle/input/make-trt-main-ensemble-original/ensemble_model1/UNet3D.engine', 'device': 0},\n    {'engine_path': '/kaggle/input/make-trt-main-ensemble-original/ensemble_model1/UNetWithMem.engine', 'device': 0},\n    {'engine_path': '/kaggle/input/make-trt-main-ensemble-original/ensemble_model1/DLinkNet3D.engine', 'device': 0}\n]\n\nensemble_model1 = load_trt_model_same_c(weight)\n\n# 異なるチャンネル数のモデルのアンサンブル\nweight1 = [\n    {'engine_path': '/kaggle/input/make-trt-main-ensemble-original/ensemble_model2/UNet_0.engine', 'device': 1}\n]\n\nweight2 = [\n    {'engine_path': '/kaggle/input/make-trt-main-ensemble-original/ensemble_model2/UNet_1.engine', 'device': 1},\n    {'engine_path': '/kaggle/input/make-trt-main-ensemble-original/ensemble_model2/UNetWithMem_2.engine', 'device': 1},\n    {'engine_path': '/kaggle/input/make-trt-main-ensemble-original/ensemble_model2/UNet_3.engine', 'device': 1},\n    {'engine_path': '/kaggle/input/make-trt-main-ensemble-original/ensemble_model2/UNetWithMemRes_4.engine', 'device': 1}\n]\n\nnew_e_model = {'engine_path': '/kaggle/input/make-trt-main-ensemble-original/ensemble_model2/load_3_models_ensemble.engine', 'device': 1}\n\nensemble_model2 = load_trt_model(weight1, weight2, new_e_model, 0.4)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-02T17:49:16.186105Z","iopub.execute_input":"2025-02-02T17:49:16.186443Z","iopub.status.idle":"2025-02-02T17:49:18.016250Z","shell.execute_reply.started":"2025-02-02T17:49:16.186373Z","shell.execute_reply":"2025-02-02T17:49:18.015423Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"ensemble_models = [ensemble_model1,ensemble_model2]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-02T17:49:18.017137Z","iopub.execute_input":"2025-02-02T17:49:18.017373Z","iopub.status.idle":"2025-02-02T17:49:18.021288Z","shell.execute_reply.started":"2025-02-02T17:49:18.017352Z","shell.execute_reply":"2025-02-02T17:49:18.020455Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import json\ncopick_config_path = TRAIN_DATA_DIR + \"/copick.config\"\n\nwith open(copick_config_path) as f:\n    copick_config = json.load(f)\n\ncopick_config['static_root'] = '/kaggle/input/czii-cryo-et-object-identification/test/static'\n\ncopick_test_config_path = 'copick_test.config'\n\nwith open(copick_test_config_path, 'w') as outfile:\n    json.dump(copick_config, outfile)\nimport copick\n\nroot = copick.from_file(copick_test_config_path)\n\ncopick_user_name = \"copickUtils\"\ncopick_segmentation_name = \"paintedPicks\"\nvoxel_size = 10\ntomo_type = \"denoised\"\n\ninference_transforms = Compose([\n    EnsureChannelFirstd(keys=[\"image\"], channel_dim=\"no_channel\"),\n    NormalizeIntensityd(keys=\"image\"),\n    Orientationd(keys=[\"image\"], axcodes=\"RAS\")\n])\nimport cc3d\n\nid_to_name = {1: \"apo-ferritin\", \n              2: \"beta-galactosidase\", \n              3: \"ribosome\", \n              4: \"thyroglobulin\", \n              5: \"virus-like-particle\"}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-02T17:49:18.022279Z","iopub.execute_input":"2025-02-02T17:49:18.022586Z","iopub.status.idle":"2025-02-02T17:49:19.939529Z","shell.execute_reply.started":"2025-02-02T17:49:18.022561Z","shell.execute_reply":"2025-02-02T17:49:19.938588Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from monai.inferers import sliding_window_inference\nimport time\nimport concurrent.futures\n\nBLOB_THRESHOLD = {1: 2, 2: 33, 3: 78, 4: 42, 5: 400}\nCERTAINTY_THRESHOLD = 0.8\n    \nclasses = [1, 2, 3, 4, 5]\n\nis_sw = True\n\ndef run_inference(ensemble, device, input_tensor, size, sw_bs):\n    pred = 0\n    tta_ct = 0\n    FLIP_TTA = True\n    FLIP_TTA_LIST = [[2]]\n    weights = [1.0]\n    ROTATE_TTA = False\n    ROTATE_K_LIST = [1,2,3]\n    OVERLAP_T = 0.15\n    out_device = 'cpu'\n    \n    with torch.cuda.device(device):\n        with torch.amp.autocast(device_type='cuda'):  # Corrected device_type\n            pred += sliding_window_inference(\n                input_tensor,\n                roi_size=[size, size, size],\n                sw_batch_size=sw_bs,\n                sw_device=device,\n                device=out_device,  # Keep on GPU\n                predictor=ensemble,\n                overlap=OVERLAP_T,\n                progress=False\n            )\n            tta_ct += 1\n            if FLIP_TTA:\n                for dims in FLIP_TTA_LIST:\n                    t_pred = sliding_window_inference(\n                        inputs=torch.flip(input_tensor, dims=dims),\n                        roi_size=[size, size, size],\n                        sw_batch_size=sw_bs,\n                        predictor=ensemble,\n                        overlap=OVERLAP_T,\n                        device=out_device,  # Keep on GPU\n                        sw_device=device,\n                        progress=False\n                    )\n                    pred += torch.flip(t_pred, dims=dims)\n                    tta_ct += 1\n            if ROTATE_TTA:\n                for k in ROTATE_K_LIST:\n                    t_pred = sliding_window_inference(\n                        inputs=torch.rot90(input_tensor, k=k, dims=[3,4]),\n                        roi_size=[size, size, size],\n                        sw_batch_size=sw_bs,\n                        predictor=ensemble,\n                        overlap=OVERLAP_T,\n                        device=out_device,  # Keep on GPU\n                        sw_device=device,\n                        progress=False\n                    )\n                    pred += torch.rot90(t_pred, k=-k, dims=[3,4])\n                    tta_ct += 1\n    pred /= tta_ct\n    return pred.cpu() # Move to CPU once after processing\n\ndef run_inference_tta(ensemble, device, input_tensor, size, sw_bs):\n    pred = 0\n    tta_ct = 0\n    FLIP_TTA_LIST = [[2],[3]]\n    with torch.cuda.device(device):\n        with torch.amp.autocast(device_type='cuda'):\n            for dims in FLIP_TTA_LIST:\n                t_pred = sliding_window_inference(\n                    inputs=torch.flip(input_tensor, dims=dims),\n                    roi_size=[size, size, size],\n                    sw_batch_size=sw_bs,\n                    predictor=ensemble,\n                    overlap=0.15,\n                    device=device,  # Keep on GPU\n                    sw_device=device,\n                    progress=False\n                )\n                pred += torch.flip(t_pred, dims=dims)\n                tta_ct += 1\n    pred /= tta_ct\n    return pred.cpu()  # Move to CPU once after processing\n\n# Start the timer\nimport time\nstart_time = time.time()\n\ndevice_ids = [\"cuda:0\", \"cuda:1\"]\nsizes = [128, 180]\nsw_bs = [2, 2]\ninference_fns = [run_inference, run_inference]  # run_inference_tta\n\nstart_time = time.time()\n\nwith torch.no_grad():\n    location_df = []\n    for run in root.runs:\n        print(run)\n    \n        tomo = run.get_voxel_spacing(10)\n        tomo = tomo.get_tomogram(tomo_type).numpy()\n\n        tomo_patched_data = [{\"image\": tomo}]\n        tomo_ds = CacheDataset(data=tomo_patched_data, transform=inference_transforms, cache_rate=1.0, progress=False)\n        st = time.time()\n        \n        input_tensor = tomo_ds[0][\"image\"].unsqueeze(0).pin_memory()  # Pinned memory\n        \n        inputs = [input_tensor.to(device, non_blocking=True) for device in device_ids]  # Non-blocking transfer\n        \n        with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:\n            futures = [\n                executor.submit(inference_fns[i], ensemble_models[i], device_ids[i], inputs[i], sizes[i], sw_bs[i])\n                for i in range(2)\n            ]\n            outputs = [future.result() for future in concurrent.futures.as_completed(futures)]\n        \n        averaged_output = sum(outputs)  # Outputs are already on CPU\n        en = time.time()\n        print(f'use time: {en-st}s')\n        \n        reconstructed_mask = torch.softmax(averaged_output, dim=1)[0]\n        _, reconstructed_mask = (reconstructed_mask > CERTAINTY_THRESHOLD).max(0)\n        reconstructed_mask = reconstructed_mask.cpu().numpy()\n                \n\n        \n        location = {}\n        for c in classes:\n            cc = cc3d.connected_components(reconstructed_mask == c)\n            stats = cc3d.statistics(cc)\n            zyx = stats['centroids'][1:] * 10.012444\n            zyx_large = zyx[stats['voxel_counts'][1:] > BLOB_THRESHOLD[c]]\n            xyz = np.ascontiguousarray(zyx_large[:, ::-1])    \n            location[id_to_name[c]] = xyz\n    \n        df = dict_to_df(location, run.name)\n        location_df.append(df)\n        \n    location_df = pd.concat(location_df)\n\n\n\n\n# End the timer\nend_time = time.time()\n\n# Calculate and print the elapsed time\nelapsed_time = end_time - start_time\n\n# Calculate the processing time for 500 voxels\ntime_per_3_voxels = elapsed_time\ntime_per_500_voxels = (elapsed_time / 3) * 500\n\n# Convert the time for 500 voxels from seconds to hours\ntime_per_500_voxels_hours = time_per_500_voxels / 3600\n\n# Print the results\nprint(f\"The processing time for 3 voxels is {elapsed_time} seconds\")\nprint(f\"The processing time for 500 voxels is approximately {time_per_500_voxels_hours} hours\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-02T17:49:19.940389Z","iopub.execute_input":"2025-02-02T17:49:19.941438Z","execution_failed":"2025-02-02T17:51:39.413Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"location_df.insert(loc=0, column='id', value=np.arange(len(location_df)))","metadata":{"trusted":true,"execution":{"execution_failed":"2025-02-02T17:51:39.414Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nfrom sklearn.cluster import DBSCAN\n\n# 假设sub已经给定，拼接DataFrame\n# df = pd.concat([sub[0], sub[1], sub[2]], ignore_index=True)\n# 粒子半径映射\n# 'beta-amylase'\nparticle_names = ['apo-ferritin', 'beta-galactosidase', 'ribosome', 'thyroglobulin', 'virus-like-particle']\nparticle_radius = {\n    'apo-ferritin': 60,\n    # 'beta-amylase': 65,\n    'beta-galactosidase': 90,\n    'ribosome': 150,\n    'thyroglobulin': 130,\n    'virus-like-particle': 135,\n}\ndf = location_df.copy()\n\nfinal = []  # 用于存储最终的结果\nfor pidx, p in enumerate(particle_names):\n    # 筛选出该粒子类型的所有点\n    pdf = df[df['particle_type'] == p].reset_index(drop=True)\n    p_rad = particle_radius[p]\n    \n    # 根据 experiment 分组\n    grouped = pdf.groupby(['experiment'])\n    \n    for exp, group in grouped:\n        group = group.reset_index(drop=True)\n        \n        # 使用DBSCAN进行聚类\n        coords = group[['x', 'y', 'z']].values\n        db = DBSCAN(eps=p_rad, min_samples=2, metric='euclidean').fit(coords)\n        labels = db.labels_\n        \n        # 将聚类结果添加到DataFrame中\n        group['cluster'] = labels\n        \n        # 对每个簇进行处理\n        for cluster_id in np.unique(labels):\n            if cluster_id == -1:\n                continue  # 跳过噪声点\n            \n            cluster_points = group[group['cluster'] == cluster_id]\n            \n            # 计算簇的中心（平均位置）\n            avg_x = cluster_points['x'].mean()\n            avg_y = cluster_points['y'].mean()\n            avg_z = cluster_points['z'].mean()\n            \n            # 更新簇内点的位置\n            group.loc[group['cluster'] == cluster_id, ['x', 'y', 'z']] = avg_x, avg_y, avg_z\n            group = group.drop_duplicates(subset=['x', 'y', 'z'])\n        # 将处理后的数据添加到 final 列表\n        final.append(group)\n\n# 合并处理后的数据\ndf_save = pd.concat(final, ignore_index=True)\ndf_save = df_save.drop(columns=['cluster'])\n\n# 排序按 'experiment' 和 'particle_type' 两列\ndf_save = df_save.sort_values(by=['experiment', 'particle_type']).reset_index(drop=True)\n\n# 重新生成 'id' 列，从 1 开始\ndf_save['id'] = np.arange(0, len(df_save))\n\n# 输出结果到 CSV\ndf_save.to_csv('submission.csv', index=False)","metadata":{"trusted":true,"execution":{"execution_failed":"2025-02-02T17:51:39.414Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}