{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":84969,"databundleVersionId":10033515,"sourceType":"competition"},{"sourceId":9862305,"sourceType":"datasetVersion","datasetId":6052780},{"sourceId":9867543,"sourceType":"datasetVersion","datasetId":6040935},{"sourceId":9869730,"sourceType":"datasetVersion","datasetId":6058495},{"sourceId":10106143,"sourceType":"datasetVersion","datasetId":6233973},{"sourceId":10364147,"sourceType":"datasetVersion","datasetId":6383310},{"sourceId":10476214,"sourceType":"datasetVersion","datasetId":6486959},{"sourceId":10522003,"sourceType":"datasetVersion","datasetId":6512054},{"sourceId":10524633,"sourceType":"datasetVersion","datasetId":6513752},{"sourceId":10551753,"sourceType":"datasetVersion","datasetId":6528712},{"sourceId":10573055,"sourceType":"datasetVersion","datasetId":6509390},{"sourceId":10587481,"sourceType":"datasetVersion","datasetId":6552409},{"sourceId":10591029,"sourceType":"datasetVersion","datasetId":6554845},{"sourceId":10595206,"sourceType":"datasetVersion","datasetId":6557851},{"sourceId":10595259,"sourceType":"datasetVersion","datasetId":6557884},{"sourceId":206640467,"sourceType":"kernelVersion"},{"sourceId":213055282,"sourceType":"kernelVersion"},{"sourceId":215061124,"sourceType":"kernelVersion"},{"sourceId":219703660,"sourceType":"kernelVersion"},{"sourceId":220601967,"sourceType":"kernelVersion"},{"sourceId":220802047,"sourceType":"kernelVersion"},{"sourceId":220905028,"sourceType":"kernelVersion"},{"sourceId":208471,"sourceType":"modelInstanceVersion","modelInstanceId":176271,"modelId":198599},{"sourceId":208598,"sourceType":"modelInstanceVersion","modelInstanceId":175864,"modelId":198200},{"sourceId":208866,"sourceType":"modelInstanceVersion","modelInstanceId":163890,"modelId":186240},{"sourceId":209184,"sourceType":"modelInstanceVersion","modelInstanceId":175864,"modelId":198200},{"sourceId":209261,"sourceType":"modelInstanceVersion","modelInstanceId":176271,"modelId":198599},{"sourceId":211515,"sourceType":"modelInstanceVersion","modelInstanceId":180333,"modelId":202595},{"sourceId":221414,"sourceType":"modelInstanceVersion","modelInstanceId":180111,"modelId":186240},{"sourceId":237845,"sourceType":"modelInstanceVersion","modelInstanceId":203128,"modelId":224862},{"sourceId":238613,"sourceType":"modelInstanceVersion","modelInstanceId":203787,"modelId":225518},{"sourceId":240586,"sourceType":"modelInstanceVersion","modelInstanceId":205581,"modelId":227328},{"sourceId":241798,"sourceType":"modelInstanceVersion","modelInstanceId":206545,"modelId":228293},{"sourceId":242194,"sourceType":"modelInstanceVersion","modelInstanceId":206545,"modelId":228293},{"sourceId":242790,"sourceType":"modelInstanceVersion","modelInstanceId":207362,"modelId":229084},{"sourceId":243128,"sourceType":"modelInstanceVersion","modelInstanceId":207362,"modelId":229084}],"dockerImageVersionId":30787,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"deps_path = '/kaggle/input/czii-cryoet-dependencies'","metadata":{"execution":{"iopub.status.busy":"2025-02-05T08:27:57.360864Z","iopub.execute_input":"2025-02-05T08:27:57.361170Z","iopub.status.idle":"2025-02-05T08:27:57.368837Z","shell.execute_reply.started":"2025-02-05T08:27:57.361142Z","shell.execute_reply":"2025-02-05T08:27:57.368100Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"! cp -r /kaggle/input/czii-cryoet-dependencies/asciitree-0.3.3/ asciitree-0.3.3/","metadata":{"execution":{"iopub.status.busy":"2025-02-05T08:27:57.370586Z","iopub.execute_input":"2025-02-05T08:27:57.370833Z","iopub.status.idle":"2025-02-05T08:27:58.486834Z","shell.execute_reply.started":"2025-02-05T08:27:57.370810Z","shell.execute_reply":"2025-02-05T08:27:58.485731Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"! pip wheel asciitree-0.3.3/asciitree-0.3.3/\n! cp /kaggle/input/best-weights-code/* ./ -r\n! cp /kaggle/input/best-weights ./ -r\n! cp /kaggle/input/voxhrnet-v2/voxhrnetV2.py ./ -r\n! cp /kaggle/input/load-vox-net/load_model.py ./ -r","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2025-02-05T08:27:58.488111Z","iopub.execute_input":"2025-02-05T08:27:58.488403Z","iopub.status.idle":"2025-02-05T08:28:41.758357Z","shell.execute_reply.started":"2025-02-05T08:27:58.488375Z","shell.execute_reply":"2025-02-05T08:28:41.757003Z"},"trusted":true},"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\n!cp -r /kaggle/input/tensorrt-10-1-0/torch2trt-master /kaggle/working/torch2trt\n!pip install /kaggle/working/torch2trt","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:28:41.759933Z","iopub.execute_input":"2025-02-05T08:28:41.760260Z","iopub.status.idle":"2025-02-05T08:34:03.501882Z","shell.execute_reply.started":"2025-02-05T08:28:41.760231Z","shell.execute_reply":"2025-02-05T08:34:03.500297Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"! cp /kaggle/input/d/luoziqian/unet2e3d-6c/* ./\n! cp /kaggle/input/voxhrnet/voxhrnet.py ./ -r\n! cp /kaggle/input/vox-networks-dataset/voxhrnet.py ./\n! cp /kaggle/input/vox-networks-dataset/config_small.yaml ./\n! cp /kaggle/input/vox-networks-dataset/model10.py ./","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:34:03.505608Z","iopub.execute_input":"2025-02-05T08:34:03.506067Z","iopub.status.idle":"2025-02-05T08:34:08.574736Z","shell.execute_reply.started":"2025-02-05T08:34:03.506026Z","shell.execute_reply":"2025-02-05T08:34:08.573751Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install asciitree-0.3.3-py3-none-any.whl\n! pip install /kaggle/input/einops-0-8-none-any/einops-0.8.0-py3-none-any.whl","metadata":{"execution":{"iopub.status.busy":"2025-02-05T08:34:08.576076Z","iopub.execute_input":"2025-02-05T08:34:08.576395Z","iopub.status.idle":"2025-02-05T08:35:29.106850Z","shell.execute_reply.started":"2025-02-05T08:34:08.576368Z","shell.execute_reply":"2025-02-05T08:35:29.105794Z"},"trusted":true},"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,"execution":{"iopub.status.busy":"2025-02-05T08:35:29.108295Z","iopub.execute_input":"2025-02-05T08:35:29.108676Z","iopub.status.idle":"2025-02-05T08:35:48.740206Z","shell.execute_reply.started":"2025-02-05T08:35:29.108632Z","shell.execute_reply":"2025-02-05T08:35:48.739313Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"! pip install /kaggle/input/vox-networks-dataset/yacs-0.1.8-py3-none-any.whl","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:35:48.741515Z","iopub.execute_input":"2025-02-05T08:35:48.741836Z","iopub.status.idle":"2025-02-05T08:36:29.535316Z","shell.execute_reply.started":"2025-02-05T08:35:48.741808Z","shell.execute_reply":"2025-02-05T08:36:29.534480Z"}},"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-05T08:36:29.536780Z","iopub.execute_input":"2025-02-05T08:36:29.537189Z","iopub.status.idle":"2025-02-05T08:37:01.945913Z","shell.execute_reply.started":"2025-02-05T08:36:29.537145Z","shell.execute_reply":"2025-02-05T08:37:01.945208Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport shutil\nimport tempfile\nfrom tqdm import tqdm\nimport matplotlib.pyplot as plt\n%matplotlib inline\n\nfrom monai.transforms import (\n    Compose,\n    Orientationd,\n    RandFlipd,\n    RandShiftIntensityd,\n    RandRotate90d,\n)\nfrom monai.data import (\n    ThreadDataLoader,\n    CacheDataset,\n    load_decathlon_datalist,\n    decollate_batch,\n    set_track_meta,\n)\nfrom monai.inferers import sliding_window_inference\nfrom monai.networks.nets import SwinUNETR\nfrom monai.metrics import DiceMetric\nfrom monai.losses import DiceCELoss\nimport torch\nimport einops\nimport warnings\n\n\nwarnings.filterwarnings(\"ignore\")\nos.environ[\"CUDA_DEVICE_ORDER\"] = \"PCI_BUS_ID\"\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:37:01.947018Z","iopub.execute_input":"2025-02-05T08:37:01.948043Z","iopub.status.idle":"2025-02-05T08:37:02.003513Z","shell.execute_reply.started":"2025-02-05T08:37:01.947997Z","shell.execute_reply":"2025-02-05T08:37:02.002602Z"}},"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_with_overlap(\n    dimension_size: int, patch_size: int, overlap: int\n) -> List[int]:\n    if dimension_size <= patch_size:\n        return [0]\n\n    num_patches = np.ceil(\n        (dimension_size - overlap) / (patch_size - overlap) + 1\n    ).astype(int)\n    patch_starts = []\n    for i in range(num_patches):\n        pos = int(i * (patch_size - overlap))\n        if pos + patch_size > dimension_size:\n            pos = dimension_size - patch_size\n        if pos not in patch_starts:\n            patch_starts.append(pos)\n    return patch_starts\n\n\ndef extract_3d_patches_overlap(\n    arrays: List[np.ndarray],\n    patch_sizes: Tuple[int, int, int],\n    overlap_sizes: Tuple[int, int, int],\n) -> Tuple[List[np.ndarray], List[Tuple[int, int, int]]]:\n\n    patch_starts_x = calculate_patch_starts_with_overlap(\n        arrays[0].shape[0], patch_sizes[0], overlap_sizes[0]\n    )\n    patch_starts_y = calculate_patch_starts_with_overlap(\n        arrays[0].shape[1], patch_sizes[1], overlap_sizes[1]\n    )\n    patch_starts_z = calculate_patch_starts_with_overlap(\n        arrays[0].shape[2], patch_sizes[2], overlap_sizes[2]\n    )\n    patch_size_d, patch_size_h, patch_size_w = patch_sizes\n    patches = []\n    coordinates = []\n    for arr in arrays:\n        for x in patch_starts_x:\n            for y in patch_starts_y:\n                for z in patch_starts_z:\n                    patch = arr[\n                        x : x + patch_size_d, y : y + patch_size_h, z : z + patch_size_w\n                    ]\n                    patches.append(patch)\n                    coordinates.append((x, y, z))\n\n    return patches, coordinates","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:37:02.004682Z","iopub.execute_input":"2025-02-05T08:37:02.005003Z","iopub.status.idle":"2025-02-05T08:37:02.023139Z","shell.execute_reply.started":"2025-02-05T08:37:02.004956Z","shell.execute_reply":"2025-02-05T08:37:02.022413Z"}},"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-exp-name\"\nTEST_DATA_DIR = \"/kaggle/input/czii-cryo-et-object-identification\"","metadata":{"execution":{"iopub.status.busy":"2025-02-05T08:37:02.024304Z","iopub.execute_input":"2025-02-05T08:37:02.025145Z","iopub.status.idle":"2025-02-05T08:37:02.034387Z","shell.execute_reply.started":"2025-02-05T08:37:02.025112Z","shell.execute_reply":"2025-02-05T08:37:02.033634Z"},"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":"from torch.nn.modules import Module\nfrom torch import nn","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:37:02.035265Z","iopub.execute_input":"2025-02-05T08:37:02.035497Z","iopub.status.idle":"2025-02-05T08:37:02.044823Z","shell.execute_reply.started":"2025-02-05T08:37:02.035473Z","shell.execute_reply":"2025-02-05T08:37:02.044166Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!cp -r /kaggle/input/unet2d3e/* ./","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:37:02.049144Z","iopub.execute_input":"2025-02-05T08:37:02.049403Z","iopub.status.idle":"2025-02-05T08:37:03.095694Z","shell.execute_reply.started":"2025-02-05T08:37:02.049380Z","shell.execute_reply":"2025-02-05T08:37:03.094465Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import lightning.pytorch as pl\n\nfrom monai.networks.nets import UNet\nfrom monai.losses import TverskyLoss\nfrom monai.metrics import DiceMetric\n\n# use warmup lr scheduler\nfrom torch.optim.lr_scheduler import (\n    CosineAnnealingLR,\n    CosineAnnealingWarmRestarts,\n    StepLR,\n)\nimport lightning.pytorch as pl\n\nfrom monai.networks.nets import UNet, AttentionUnet\nfrom monai.losses import TverskyLoss\nfrom monai.metrics import DiceMetric\n\n# use warmup lr scheduler\nfrom torch.optim.lr_scheduler import (\n    CosineAnnealingLR,\n    CosineAnnealingWarmRestarts,\n    StepLR,\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:37:03.097171Z","iopub.execute_input":"2025-02-05T08:37:03.097461Z","iopub.status.idle":"2025-02-05T08:37:03.597212Z","shell.execute_reply.started":"2025-02-05T08:37:03.097434Z","shell.execute_reply":"2025-02-05T08:37:03.596491Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"config_v2 = {\n    \"MODEL\": {\n        \"NAME\": \"voxhrnet\",\n        \"EXTRA\": {\n            \"STAGE2\": {\n                \"NUM_MODULES\": 1,\n                \"NUM_BRANCHES\": 2,\n                \"BLOCK\": \"BASIC\",\n                \"NUM_BLOCKS\": [3, 3],\n                \"NUM_CHANNELS\": [16, 32],\n            },\n            \"STAGE3\": {\n                \"NUM_MODULES\": 1,\n                \"NUM_BRANCHES\": 3,\n                \"BLOCK\": \"BASIC\",\n                \"NUM_BLOCKS\": [3, 3, 3],\n                \"NUM_CHANNELS\": [16, 32, 64],\n            },\n        },\n    }\n}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:37:03.598197Z","iopub.execute_input":"2025-02-05T08:37:03.598457Z","iopub.status.idle":"2025-02-05T08:37:03.603314Z","shell.execute_reply.started":"2025-02-05T08:37:03.598431Z","shell.execute_reply":"2025-02-05T08:37:03.602348Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# from monai.networks.nets import UNet,SegResNet,DynUNet\n# from custom_vnet import CustomVNet\n# from model2_6c import Net\n# # from voxhrnet import build_model\n# from voxhrnetV2 import build_model as build_model_v2\n# from load_model import build_model as build_model_v3\n\n# basic_unet = UNet(\n#     spatial_dims=3,\n#     in_channels=1,\n#     out_channels=7,\n#     channels=(48, 64, 80, 80),\n#     strides=(2, 2, 1),\n#     num_res_units=1,\n# )\n# basic_unet_1 = UNet(\n#     spatial_dims=3,\n#     in_channels=1,\n#     out_channels=7,\n#     channels=(48, 64, 80, 80),\n#     strides=(2, 2, 1),\n#     num_res_units=2,\n# )\n# basic_unet_6c = UNet(\n#     spatial_dims=3,\n#     in_channels=1,\n#     out_channels=6,\n#     channels=(48, 64, 80, 80),\n#     strides=(2, 2, 1),\n#     num_res_units=1,\n# )\n\n# segresnet_6c_v1 = SegResNet(\n#     in_channels=1,\n#     out_channels=6,\n#     dropout_prob=0.1,\n#     upsample_mode=\"deconv\",\n# )\n\n# basic_unet2e3d = Net(\n#     out_channels=6,\n#     arch=\"resnet18d\",\n#     decoder_dim=[80, 80, 64, 32, 16],\n#     pretrained=False,\n# )\n\n# basic_dynunet_v1 = DynUNet(\n#     spatial_dims=3,\n#     in_channels=1,\n#     out_channels=6,\n#     kernel_size=(3, 3, 3, 3),\n#     strides=((1, 1, 1), 2, 2, 1),\n#     upsample_kernel_size=(2, 2, 1),\n#     filters=[16, 24, 48, 80, 80],\n#     norm_name=\"instance\",\n#     act_name=\"PRELU\",\n#     deep_supervision=True,\n# )\n\n# basic_dynunet_v1 = DynUNet(\n#     spatial_dims=3,\n#     in_channels=1,\n#     out_channels=6,\n#     kernel_size=(3, 3, 3, 3),\n#     strides=((1, 1, 1), 2, 2, 1),\n#     upsample_kernel_size=(2, 2, 1),\n#     filters=[16, 24, 48, 80, 80],\n#     norm_name=\"instance\",\n#     act_name=\"PRELU\",\n#     deep_supervision=True,\n# )\n\n# basic_dynunet_v2 = DynUNet(\n#     spatial_dims=3,\n#     in_channels=1,\n#     out_channels=6,\n#     kernel_size=(3, 3, 3, 3, 3),\n#     strides=((1, 1, 1), 2, 2, 2, 1),\n#     upsample_kernel_size=(2, 2, 2, 1),\n#     filters=[16, 24, 48, 80, 80],\n#     norm_name=\"instance\",\n#     act_name=\"PRELU\",\n#     deep_supervision=True,\n# )\n\n# # basic_voxhrnet_v0 = build_model(1,6)\n# basic_voxhrnet_v2 = build_model_v2(1,6,config_v2)\n# basic_voxhrnet_v3 = build_model_v3()\n\n# basic_dense_vnet = CustomVNet(in_channels=1, classes=7)","metadata":{"execution":{"iopub.status.busy":"2025-02-05T08:37:03.604576Z","iopub.execute_input":"2025-02-05T08:37:03.604893Z","iopub.status.idle":"2025-02-05T08:37:03.614015Z","shell.execute_reply.started":"2025-02-05T08:37:03.604858Z","shell.execute_reply":"2025-02-05T08:37:03.613147Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\n\nbest_weights_dir = \"best-weights\"\nos.listdir(best_weights_dir)\nbest_weights_path_list = [\n 'best-weights/epoch122-step2952-valid_loss0.3625-val_metric0.8367.ckpt',\n 'best-weights/epoch148-step3576-valid_loss1.1154-val_metric0.7722.ckpt',\n 'best-weights/epoch153-step3696-valid_loss0.3021-val_metric0.8900.ckpt',\n 'best-weights/epoch194-step4680-valid_loss1.0213-val_metric0.8788.ckpt',\n 'best-weights/epoch152-step3672-valid_loss0.4333-val_metric0.7929.ckpt',\n] \\\n + ['/kaggle/input/unet2e3d-6c/pytorch/default/1/unet2E3D-v1-epoch114-val_loss0.55-val_metric0.53-step2760.ckpt'] \n # + ['/kaggle/input/vox-networks-dataset/epoch133-step3216-valid_loss0.4223-val_metric0.8612.ckpt']\n#  + ['/kaggle/input/dynunet-6c/pytorch/default/1/DynUnet-epoch189-val_loss0.51-val_metric0.55-step4560.ckpt'] \\\n#  +['/kaggle/input/segresnet-6c/pytorch/default/1/epoch314-val_loss0.54-val_metric0.54-step7560.ckpt']\n\n#  +['/kaggle/input/my-weights/unet3D-epoch173-val_loss0.53-val_metric0.54-step4176.ckpt']\n # +['/kaggle/input/unet2e3d-6c/pytorch/default/1/unet2E3D-v1-epoch114-val_loss0.55-val_metric0.53-step2760.ckpt'] \\\n# best_weights_path_list=['/kaggle/input/dynunet-6c/pytorch/default/1/DynUnet-epoch189-val_loss0.51-val_metric0.55-step4560.ckpt']\n# best_weights_path_list=['/kaggle/input/dynunet-6c/pytorch/default/2/DynUnet-v2-epoch164-val_loss0.53-val_metric0.53-step3960.ckpt']\n# best_weights_path_list=['/kaggle/input/vox-networks-dataset/epoch133-step3216-valid_loss0.4223-val_metric0.8612.ckpt']\nbest_weights_path_list","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:37:03.615142Z","iopub.execute_input":"2025-02-05T08:37:03.615470Z","iopub.status.idle":"2025-02-05T08:37:03.628132Z","shell.execute_reply.started":"2025-02-05T08:37:03.615434Z","shell.execute_reply":"2025-02-05T08:37:03.627423Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"trt_model_path_list=[\n    \"/kaggle/input/fork-of-fork-of-model2trt-d93817/0_unet.engine\",\n    \"/kaggle/input/fork-of-fork-of-model2trt-d93817/1_unet.engine\",\n    \"/kaggle/input/fork-of-fork-of-model2trt-d93817/2_unet_1.engine\",\n    \"/kaggle/input/fork-of-fork-of-model2trt-d93817/3_unet.engine\",\n    \"/kaggle/input/fork-of-fork-of-model2trt-d93817/4_basic_voxresnet_v0.engine\",\n    \"/kaggle/input/fork-of-fork-of-fork-of-model2trt-d93817/0_bisic_voxhrnet_v4.engine\",\n    \"/kaggle/input/fork-of-fork-of-fork-of-fork-of-model2trt-d93817/0_unet2e3d_7c.engine\",\n    \"/kaggle/input/fork-of-fork-of-model2trt-d93817/7_unet_6c.engine\",\n    \"/kaggle/input/fork-of-fork-of-model2trt-d93817/8_vnet.engine\",\n    \"/kaggle/input/fork-of-fork-of-model2trt-d93817/9_segresnet_6c_v1.engine\",\n]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:37:03.629272Z","iopub.execute_input":"2025-02-05T08:37:03.629830Z","iopub.status.idle":"2025-02-05T08:37:03.639946Z","shell.execute_reply.started":"2025-02-05T08:37:03.629791Z","shell.execute_reply":"2025-02-05T08:37:03.639099Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import copy\nimport torch\n\n# model_list = []\n# for path in best_weights_path_list:\n#     ckpt = torch.load(path)\n#     state_dict_ = ckpt[\"state_dict\"]\n#     state_dict = {}\n#     for k in state_dict_.keys():\n#         if \"model.\" in k:\n#             state_dict[k[6:]] = state_dict_[k]\n#     try:\n#         model = copy.deepcopy(basic_unet)\n#         model.load_state_dict(state_dict)\n#         model_list.append(model.to(\"cpu\"))\n#         print(\"load unet\")\n#         continue\n#     except:\n#         pass\n#     try:\n#         model = copy.deepcopy(basic_unet_1)\n#         model.load_state_dict(state_dict)\n#         model_list.append(model.to(\"cpu\"))\n#         print(\"load unet_1\")\n#         continue\n#     except:\n#         # print(\"load failed\")\n#         pass\n#     try:\n#         model = copy.deepcopy(basic_dense_vnet)\n#         model.load_state_dict(state_dict)\n#         model_list.append(model.to(\"cpu\"))\n#         print(\"load vnet\")\n#         continue\n#     except:\n#         # print(\"load failed\")\n    #     pass\n    # try:\n    #     model = copy.deepcopy(basic_unet_6c)\n    #     model.load_state_dict(state_dict)\n    #     model_list.append(model.to(\"cpu\"))\n    #     print(\"load unet_6c\")\n    #     continue\n    # except:\n    #     # print(\"load failed\")\n    #     pass\n    # try:\n    #     model = copy.deepcopy(basic_unet2e3d)\n    #     model.load_state_dict(state_dict)\n    #     model_list.append(model.to(\"cpu\"))\n    #     print(\"load unet2e3d\")\n    #     continue\n    # except:\n    #     # print(\"load failed\")\n    #     pass\n    # try:\n    #     model = copy.deepcopy(segresnet_6c_v1)\n    #     model.load_state_dict(state_dict)\n    #     model_list.append(model.to(\"cpu\"))\n    #     print(\"load segresnet_6c_v1\")\n    #     continue\n    # except:\n    #     # print(\"load failed\")\n    #     pass\n    # try:\n    #     model = copy.deepcopy(basic_dynunet_v1)\n    #     model.load_state_dict(state_dict)\n    #     model_list.append(model.to(\"cpu\"))\n    #     print(\"load dynunet\")\n    #     continue\n    # except:\n    #     # print(\"load failed\")\n    #     pass\n    # try:\n    #     model = copy.deepcopy(basic_dynunet_v2)\n    #     model.load_state_dict(state_dict)\n    #     model_list.append(model.to(\"cpu\"))\n    #     print(\"load basic_dynunet_v2\")  \n    #     continue\n    # except:\n    #     # print(\"load failed\")\n    #     pass\n    # try:\n    #     model = copy.deepcopy(basic_voxhrnet_v0)\n    #     model.load_state_dict(state_dict)\n    #     model_list.append(model.to(\"cpu\"))\n    #     print(\"load basic_voxhrnet_v0\")\n    #     continue\n    # except:\n    #     pass\n    # try:\n    #     model = copy.deepcopy(basic_voxhrnet_v3)\n    #     model.load_state_dict(state_dict)\n    #     model_list.append(model.to(\"cpu\"))\n    #     print(\"load basic_voxhrnet_v3\")\n    #     continue\n    # except:\n    #     pass\n    # try:\n    #     model = copy.deepcopy(basic_voxhrnet_v2)\n    #     model.load_state_dict(state_dict)\n    #     model_list.append(model.to(\"cpu\"))\n    #     print(\"load basic_voxhrnet_v2\")\n    #     continue\n    # except:\n    #     # print(\"load failed\")\n    #     pass\n    ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:37:03.640969Z","iopub.execute_input":"2025-02-05T08:37:03.641239Z","iopub.status.idle":"2025-02-05T08:37:03.654782Z","shell.execute_reply.started":"2025-02-05T08:37:03.641213Z","shell.execute_reply":"2025-02-05T08:37:03.654105Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# model_list=trt_model_path_list+model_list[-1:]\nmodel_list=trt_model_path_list","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:37:03.655700Z","iopub.execute_input":"2025-02-05T08:37:03.655947Z","iopub.status.idle":"2025-02-05T08:37:03.665986Z","shell.execute_reply.started":"2025-02-05T08:37:03.655902Z","shell.execute_reply":"2025-02-05T08:37:03.665221Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"torch._dynamo.config.cache_size_limit = 64\ntorch.backends.cudnn.benchmark = True\ntorch.backends.cudnn.enabled = True\ntorch.jit.enable_onednn_fusion(True)\n# use cudnn.benchmark for faster training\ntorch.backends.cudnn.version(), torch.backends.cudnn.is_available()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:37:03.666877Z","iopub.execute_input":"2025-02-05T08:37:03.667252Z","iopub.status.idle":"2025-02-05T08:37:03.685195Z","shell.execute_reply.started":"2025-02-05T08:37:03.667223Z","shell.execute_reply":"2025-02-05T08:37:03.684474Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Train the model\n\n","metadata":{}},{"cell_type":"code","source":"torch.set_float32_matmul_precision('medium')\n\n# Check if CUDA is available and then count the GPUs\nif torch.cuda.is_available():\n    num_gpus = torch.cuda.device_count()\n    print(f\"Number of GPUs available: {num_gpus}\")\nelse:\n    print(\"No GPU available. Running on CPU.\")\ndevices = list(range(num_gpus))\nprint(devices)","metadata":{"execution":{"iopub.status.busy":"2025-02-05T08:37:03.686299Z","iopub.execute_input":"2025-02-05T08:37:03.686867Z","iopub.status.idle":"2025-02-05T08:37:03.725363Z","shell.execute_reply.started":"2025-02-05T08:37:03.686819Z","shell.execute_reply":"2025-02-05T08:37:03.724714Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Let there be gradients!\n\nLocally this config seems to train for about 1000 steps before the model starts overfitting. ","metadata":{}},{"cell_type":"markdown","source":"## Predict on the test set\n\n","metadata":{}},{"cell_type":"code","source":"device_0 = \"cuda:0\"\ndevice_1 = \"cuda:1\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:37:03.726184Z","iopub.execute_input":"2025-02-05T08:37:03.726391Z","iopub.status.idle":"2025-02-05T08:37:03.729733Z","shell.execute_reply.started":"2025-02-05T08:37:03.726369Z","shell.execute_reply":"2025-02-05T08:37:03.728964Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# for idx, model in enumerate(model_list):\n#     # model_list[idx] = model_list[idx].model\n#     if isinstance(model, str):\n#         continue\n#     model_list[idx] = model_list[idx].eval().half()\n#     model_list[idx] = model_list[idx].cpu()\n#     # model_list[idx] = torch.compile(model_list[idx], mode=\"reduce-overhead\")\n\n# torch.cuda.empty_cache()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:37:03.730863Z","iopub.execute_input":"2025-02-05T08:37:03.731324Z","iopub.status.idle":"2025-02-05T08:37:03.737121Z","shell.execute_reply.started":"2025-02-05T08:37:03.731285Z","shell.execute_reply":"2025-02-05T08:37:03.736498Z"}},"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)","metadata":{"execution":{"iopub.status.busy":"2025-02-05T08:37:03.737973Z","iopub.execute_input":"2025-02-05T08:37:03.738226Z","iopub.status.idle":"2025-02-05T08:37:03.753533Z","shell.execute_reply.started":"2025-02-05T08:37:03.738203Z","shell.execute_reply":"2025-02-05T08:37:03.752956Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import 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\"","metadata":{"execution":{"iopub.status.busy":"2025-02-05T08:37:03.754356Z","iopub.execute_input":"2025-02-05T08:37:03.755094Z","iopub.status.idle":"2025-02-05T08:37:05.052016Z","shell.execute_reply.started":"2025-02-05T08:37:03.755057Z","shell.execute_reply":"2025-02-05T08:37:05.051322Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Non-random transforms to be cached\ninference_transforms = Compose([\n    EnsureChannelFirstd(keys=[\"image\"], channel_dim=\"no_channel\"),\n    NormalizeIntensityd(keys=\"image\"),\n    Orientationd(keys=[\"image\"], axcodes=\"RAS\")\n])","metadata":{"execution":{"iopub.status.busy":"2025-02-05T08:37:05.052986Z","iopub.execute_input":"2025-02-05T08:37:05.053778Z","iopub.status.idle":"2025-02-05T08:37:05.058295Z","shell.execute_reply.started":"2025-02-05T08:37:05.053749Z","shell.execute_reply":"2025-02-05T08:37:05.057455Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import cc3d\n\nid_to_name = {1: \"apo-ferritin\", \n              2: \"beta-amylase\",\n              3: \"beta-galactosidase\", \n              4: \"ribosome\", \n              5: \"thyroglobulin\", \n              6: \"virus-like-particle\"}","metadata":{"execution":{"iopub.status.busy":"2025-02-05T08:37:05.059350Z","iopub.execute_input":"2025-02-05T08:37:05.059614Z","iopub.status.idle":"2025-02-05T08:37:05.076390Z","shell.execute_reply.started":"2025-02-05T08:37:05.059589Z","shell.execute_reply":"2025-02-05T08:37:05.075625Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Iterate over test set\n\n\nBelow we will: \n1. Read in a run\n2. Split it into patches of size (96, 96, 96)\n3. Create a dataset from the patches\n4. Predict the segmentation mask\n5. Glue the mask back together\n6. Find the connected components for each class\n7. Find the centroids of the connected components\n8. Add to the dataframe\n\nThen do this for all runs. \n\nThis can probably be optimized quite a bit. ","metadata":{}},{"cell_type":"code","source":"# Non-random transforms to be cached\ninference_transforms = Compose(\n    [\n        EnsureChannelFirstd(keys=[\"image\"], channel_dim=\"no_channel\"),\n        NormalizeIntensityd(keys=\"image\"),\n        Orientationd(keys=[\"image\"], axcodes=\"RAS\"),\n    ]\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:37:05.077424Z","iopub.execute_input":"2025-02-05T08:37:05.078144Z","iopub.status.idle":"2025-02-05T08:37:05.085831Z","shell.execute_reply.started":"2025-02-05T08:37:05.078114Z","shell.execute_reply":"2025-02-05T08:37:05.085257Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from monai.networks.nets import UNet\nfrom torch2trt import TRTModule\nimport tensorrt as trt\nfrom cuda import cudart","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:37:05.086744Z","iopub.execute_input":"2025-02-05T08:37:05.087008Z","iopub.status.idle":"2025-02-05T08:37:05.194361Z","shell.execute_reply.started":"2025-02-05T08:37:05.086983Z","shell.execute_reply":"2025-02-05T08:37:05.193763Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class Net:\n    def __init__(self, weights, device=0):\n        cudart.cudaSetDevice(device)\n        self.output_type = ['infer']\n        logger = trt.Logger(trt.Logger.INFO)\n        trt.init_libnvinfer_plugins(logger, \"\")\n        runtime = trt.Runtime(logger)\n        with open(weights, \"rb\") as f:\n            engine = runtime.deserialize_cuda_engine(f.read())\n        \n        self.trt_model = TRTModule(input_names=['images'],\n                                   output_names=['output'],\n                                   engine=engine)\n        self.device = device\n\n    def __del__(self):\n        del self.trt_model\n\n\n    def __call__(self, image):\n        # image=image.float()\n        # image =image.cuda(self.device)\n        output = {}\n        if 'infer' in self.output_type:\n            output['particle'] = self.trt_model(image)\n        return output['particle']","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:37:05.195188Z","iopub.execute_input":"2025-02-05T08:37:05.195397Z","iopub.status.idle":"2025-02-05T08:37:05.201856Z","shell.execute_reply.started":"2025-02-05T08:37:05.195375Z","shell.execute_reply":"2025-02-05T08:37:05.200966Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"threshold = {\n    \"apo-ferritin\": 0.05,\n    \"beta-amylase\": 0.05,\n    \"beta-galactosidase\": 0.05,\n    \"ribosome\": 0.05,\n    \"thyroglobulin\": 0.05,\n    \"virus-like-particle\": 0.05,\n}\nthreshold = list(threshold.values())\nthreshold = torch.tensor(threshold, device=device).reshape(6, 1, 1, 1)","metadata":{"execution":{"iopub.status.busy":"2025-02-05T08:37:05.203072Z","iopub.execute_input":"2025-02-05T08:37:05.203637Z","iopub.status.idle":"2025-02-05T08:37:05.385069Z","shell.execute_reply.started":"2025-02-05T08:37:05.203600Z","shell.execute_reply":"2025-02-05T08:37:05.384328Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def extract_3d_patches_overlap_wo_data(\n    arrays_shape: Tuple[int,int,int],\n    patch_sizes: Tuple[int, int, int],\n    overlap_sizes: Tuple[int, int, int],\n) -> List[Tuple[int, int, int]]:\n\n    patch_starts_x = calculate_patch_starts_with_overlap(\n        arrays_shape[0], patch_sizes[0], overlap_sizes[0])\n    patch_starts_y = calculate_patch_starts_with_overlap(\n        arrays_shape[1], patch_sizes[1], overlap_sizes[1])\n    patch_starts_z = calculate_patch_starts_with_overlap(\n        arrays_shape[2], patch_sizes[2], overlap_sizes[2])\n    patch_size_d, patch_size_h, patch_size_w = patch_sizes\n    coordinates = []\n    for x in patch_starts_x:\n        for y in patch_starts_y:\n            for z in patch_starts_z:\n                coordinates.append((x, y, z))\n\n    return coordinates","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:37:05.386359Z","iopub.execute_input":"2025-02-05T08:37:05.386746Z","iopub.status.idle":"2025-02-05T08:37:05.393820Z","shell.execute_reply.started":"2025-02-05T08:37:05.386706Z","shell.execute_reply":"2025-02-05T08:37:05.392858Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import threading\nfrom copy import deepcopy\nimport gc\n\nlock = threading.Lock()\n\n\ndef build_weights(\n    patch_sizes, overlap_sizes, EDGE_WEIGHT, mask_hard_cases=False, device=\"cuda:0\"\n):\n    # overlap_sizes = (32, 32, 32)\n    weights = torch.zeros(\n        *patch_sizes, dtype=torch.half, device=device, requires_grad=False\n    )\n    weights += EDGE_WEIGHT\n    weights[\n        overlap_sizes[0] // 4 : patch_sizes[0] - overlap_sizes[0] // 4,\n        overlap_sizes[1] // 4 : patch_sizes[1] - overlap_sizes[1] // 4,\n        overlap_sizes[2] // 4 : patch_sizes[2] - overlap_sizes[2] // 4,\n    ] += (\n        1 - EDGE_WEIGHT\n    )\n    return weights\n\n\ndef infer(\n    run,\n    device_id,\n    CERTAINTY_THRESHOLD=0.25,\n    EDGE_WEIGHT=0.1,\n    submit_df=None,\n    model_list=None,\n    patch_sizes_list=None,\n    overlap_sizes=None,\n    tomo_threshold=None,\n    act_mode_list=None,\n    probability=None,\n    count=None,\n):\n    cudart.cudaSetDevice(device_id)\n    torch_device = torch.device(f\"cuda:{device_id}\")\n    with torch.no_grad():\n        device = torch.device(f\"cuda:{device_id}\")\n        tomo = run.get_voxel_spacing(10)\n        tomo = tomo.get_tomogram(tomo_type).numpy()\n\n        # tomo = inference_transforms({\"image\": tomo})[\"image\"]\n\n        # tomo_patches, coordinates = extract_3d_patches_overlap(\n        #     [tomo], patch_sizes, overlap_sizes\n        # )\n        # image = torch.tensor(tomo, device=device, dtype=torch.half)\n\n        D, H, W = (184, 630, 630)\n        probability.zero_()\n        count.zero_()\n        coordinates_list = [\n            extract_3d_patches_overlap_wo_data((D, H, W), patch_sizes, overlap_sizes)\n            for patch_sizes in patch_sizes_list\n        ]\n        weights_list = [\n            build_weights(patch_sizes, overlap_sizes, EDGE_WEIGHT, False, device)\n            for patch_sizes in patch_sizes_list\n        ]\n        # coordinates = extract_3d_patches_overlap_wo_data(\n        #     (D, H, W), patch_sizes, overlap_sizes\n        # )\n\n        for idx, (model_, patch_sizes, coordinates, weights) in enumerate(\n            zip(model_list, patch_sizes_list, coordinates_list, weights_list)\n        ):\n            if isinstance(model_, str):\n                print(model_)\n                model = Net(weights=deepcopy(model_), device=device_id)\n            else:\n                model = model_.to(device)\n            # weights = weights.to(device)\n            for i in range(len(coordinates)):\n                coor = coordinates[i]\n                input_tensor_0 = tomo[\n                    coor[0] : coor[0] + patch_sizes[0],\n                    coor[1] : coor[1] + patch_sizes[1],\n                    coor[2] : coor[2] + patch_sizes[2],\n                ]\n                input_tensor_0 = inference_transforms({\"image\": input_tensor_0})[\n                    \"image\"\n                ]  # (1, D, H, W)\n                input_tensor_0 = torch.tensor(\n                    input_tensor_0,\n                    dtype=torch.half,\n                    device=device,\n                    requires_grad=False,\n                ).unsqueeze(\n                    0\n                )  # input_tensor shape B,C,D,H,W\n                input_tensor_90 = torch.rot90(input_tensor_0, 1, [3, 4])\n                input_tensor_180 = torch.rot90(input_tensor_0, 2, [3, 4])\n                input_tensor_270 = torch.rot90(input_tensor_0, 3, [3, 4])\n                input_tensor_2 = torch.flip(input_tensor_0, [2])\n                input_tensor_3 = torch.flip(input_tensor_0, [3])\n                input_tensor_4 = torch.flip(input_tensor_0, [4])\n                # input_tensor_270_2 = torch.flip(input_tensor_270, [2])\n                _, _, patch_d, patch_h, patch_w = input_tensor_0.shape\n                # coor = tomo_ds[i][\"coord\"]\n\n                model_output_0 = model(input_tensor_0)\n                torch.cuda.synchronize(torch_device)\n                model_output_90 = model(input_tensor_90)\n                torch.cuda.synchronize(torch_device)\n                model_output_180 = model(input_tensor_180)\n                torch.cuda.synchronize(torch_device)\n                model_output_270 = model(input_tensor_270)\n                torch.cuda.synchronize(torch_device)\n                model_output_2 = model(input_tensor_2)\n                torch.cuda.synchronize(torch_device)\n                model_output_3 = model(input_tensor_3)\n                torch.cuda.synchronize(torch_device)\n                model_output_4 = model(input_tensor_4)\n                torch.cuda.synchronize(torch_device)\n\n                model_output_90 = torch.rot90(model_output_90, 3, [3, 4])\n                model_output_180 = torch.rot90(model_output_180, 2, [3, 4])\n                model_output_270 = torch.rot90(model_output_270, 1, [3, 4])\n                model_output_2 = torch.flip(model_output_2, [2])\n                model_output_3 = torch.flip(model_output_3, [3])\n                model_output_4 = torch.flip(model_output_4, [4])\n\n                if act_mode_list and act_mode_list[idx] == \"sigmoid\":\n                    probs = (\n                        torch.nn.functional.sigmoid(model_output_0[0])\n                        + torch.nn.functional.sigmoid(model_output_90[0])\n                        + torch.nn.functional.sigmoid(model_output_180[0])\n                        + torch.nn.functional.sigmoid(model_output_270[0])\n                        + torch.nn.functional.sigmoid(model_output_2[0])\n                        + torch.nn.functional.sigmoid(model_output_3[0])\n                        + torch.nn.functional.sigmoid(model_output_4[0])\n                    ) / 7\n                else:\n                    probs = (\n                        torch.softmax(model_output_0[0], dim=0)\n                        + torch.softmax(model_output_90[0], dim=0)\n                        + torch.softmax(model_output_180[0], dim=0)\n                        + torch.softmax(model_output_270[0], dim=0)\n                        + torch.softmax(model_output_2[0], dim=0)\n                        + torch.softmax(model_output_3[0], dim=0)\n                        + torch.softmax(model_output_4[0], dim=0)\n                    ) / 7\n                # print(\"Max min mean\")\n                # print(probs.max())\n                # print(probs.min())\n                # print(probs.mean())\n                probability[\n                    0,\n                    coor[0] : coor[0] + patch_sizes[0],\n                    coor[1] : coor[1] + patch_sizes[1],\n                    coor[2] : coor[2] + patch_sizes[2],\n                ] += (\n                    probs[1] * weights\n                )\n                probability[\n                    -4:,\n                    coor[0] : coor[0] + patch_sizes[0],\n                    coor[1] : coor[1] + patch_sizes[1],\n                    coor[2] : coor[2] + patch_sizes[2],\n                ] += (\n                    probs[-4:] * weights\n                )\n                count[\n                    coor[0] : coor[0] + patch_sizes[0],\n                    coor[1] : coor[1] + patch_sizes[1],\n                    coor[2] : coor[2] + patch_sizes[2],\n                ] += weights\n\n            # weights = weights.to(\"cpu\")\n            if isinstance(model_, str):\n                # del model.trt_model\n                # del model\n                # gc.collect()\n                torch.cuda.empty_cache()\n            else:\n                model.to(\"cpu\")\n                torch.cuda.empty_cache()\n\n        probability = probability / count\n\n        probability0 = probability\n        probability1 = F.interpolate(\n            probability0[1:], scale_factor=0.5, mode=\"bilinear\", align_corners=False\n        )\n        binary0 = (probability0 > CERTAINTY_THRESHOLD).data.cpu().numpy()\n        binary1 = (probability1 > CERTAINTY_THRESHOLD).data.cpu().numpy()\n\n        location = [np.empty((0, 3)) for i in range(6)]\n\n        for c in [0]:\n            componet = cc3d.connected_components(binary0[c])\n            stats = cc3d.statistics(componet)\n            zyx = stats[\"centroids\"][1:] * 10.012444\n            # zyx_large = zyx[stats[\"voxel_counts\"][1:] > BLOB_THRESHOLD]\n            if tomo_threshold:\n                zyx = zyx[stats[\"voxel_counts\"][1:] > tomo_threshold[c]]\n            xyz = np.ascontiguousarray(zyx[:, ::-1])\n            location[c] = xyz\n\n        for c in [2, 3, 4, 5]:\n            componet = cc3d.connected_components(binary1[c - 2])\n            stats = cc3d.statistics(componet)\n            zyx = stats[\"centroids\"][1:] * 10.012444 * [[1, 2, 2]]\n            # zyx_large = zyx[stats[\"voxel_counts\"][1:] > int(BLOB_THRESHOLD // 4)]\n            if tomo_threshold:\n                zyx = zyx[stats[\"voxel_counts\"][1:] > tomo_threshold[c]]\n            xyz = np.ascontiguousarray(zyx[:, ::-1])\n            location[c] = xyz\n        print(\"location\", np.concatenate(location).shape)\n        for class_id, name in id_to_name.items():\n            class_id = int(class_id) - 1\n            \"\"\"pd.DataFrame(\n            {\n                \"experiment\": id,\n                \"particle_type\": name,\n                \"x\": xyz[:, 0],\n                \"y\": xyz[:, 1],\n                \"z\": xyz[:, 2],\n            }\"\"\"\n            submit_df.append(\n                pd.DataFrame(\n                    {\n                        \"experiment\": run.name,\n                        \"particle_type\": name,\n                        \"x\": location[class_id][:, 0],\n                        \"y\": location[class_id][:, 1],\n                        \"z\": location[class_id][:, 2],\n                    }\n                )\n            )\n\n\ndef thread_infer(\n    run_list,\n    device_id,\n    CERTAINTY_THRESHOLD=0.25,\n    EDGE_WEIGHT=0.1,\n    submit_df=None,\n    model_list=None,\n    patch_sizes_list=None,\n    overlap_sizes=None,\n    tomo_threshold=None,\n    act_mode_list=None,\n):\n    D, H, W = (184, 630, 630)\n    device = torch.device(f\"cuda:{device_id}\")\n    probability = torch.zeros(\n        (5, D, H, W), device=device, dtype=torch.half, requires_grad=False\n    )\n    count = torch.zeros((D, H, W), device=device, dtype=torch.half, requires_grad=False)\n\n    while True:\n        with lock:\n            if len(run_list) == 0:\n                break\n            run = run_list.pop()\n        infer(\n            run,\n            device_id,\n            CERTAINTY_THRESHOLD,\n            EDGE_WEIGHT,\n            submit_df,\n            model_list,\n            patch_sizes_list,\n            overlap_sizes,\n            tomo_threshold,\n            act_mode_list,\n            probability,\n            count,\n        )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:37:05.395128Z","iopub.execute_input":"2025-02-05T08:37:05.395377Z","iopub.status.idle":"2025-02-05T08:37:05.429719Z","shell.execute_reply.started":"2025-02-05T08:37:05.395353Z","shell.execute_reply":"2025-02-05T08:37:05.428694Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch.nn.functional as F\nimport pandas as pd\nimport threading\nimport copy\n\nCERTAINTY_THRESHOLD = 0.15\nEDGE_WEIGHT = 0.1\npatch_sizes_list = [\n    [128, 384, 384],\n    [128, 384, 384],\n    [128, 384, 384],\n    [128, 384, 384],\n    [128, 384, 384],\n    [128, 384, 384],\n    [128, 384, 384],\n    [128, 384, 384],\n    [128, 384, 384],\n    [128, 384, 384],\n]\nact_mode_list = [\n    \"softmax\",\n    \"softmax\",\n    \"softmax\",\n    \"softmax\",\n    \"softmax\",\n    \"softmax\",\n    \"softmax\",\n    \"softmax\",\n    'softmax',\n    'softmax',\n]\noverlap_sizes = [32, 32, 32]\ntomo_threshold = [2, 2, 2, 5, 5, 5]\n\nclasses = [1, 2, 3, 4, 5, 6]\n\ntask_runs = root.runs # * 150\n# num_tasks = len(task_runs)\n# task_device_0 = task_runs[: num_tasks // 2]\n# task_device_1 = task_runs[num_tasks // 2 :]\n\nthread_list = []\nsubmit_device_0 = []\nsubmit_device_1 = []\n# model_list str & nn.model\n# model_device_0 = [copy.deepcopy(model).to(\"cuda:0\") for model in model_list if isinstance(model, nn.Module) else model]\n# model_device_1 = [copy.deepcopy(model).to(\"cuda:1\") for model in model_list if isinstance(model, nn.Module) else model]\nmodel_device_0 = []\nmodel_device_1 = []\nprint(\"start\")\nfor model in model_list:\n    if isinstance(model, nn.Module):\n        model_device_0.append(copy.deepcopy(model).to(\"cuda:0\"))\n        model_device_1.append(copy.deepcopy(model).to(\"cuda:1\"))\n    else:\n        model_device_0.append(model)\n        model_device_1.append(model)\nthread_list.append(\n    threading.Thread(\n        target=thread_infer,\n        args=(\n            task_runs,\n            0,\n            CERTAINTY_THRESHOLD,\n            EDGE_WEIGHT,\n            submit_device_0,\n            model_device_0,\n            patch_sizes_list,\n            overlap_sizes,\n            tomo_threshold,\n            act_mode_list,\n        ),\n    )\n)\nthread_list.append(\n    threading.Thread(\n        target=thread_infer,\n        args=(\n            task_runs,\n            1,\n            CERTAINTY_THRESHOLD,\n            EDGE_WEIGHT,\n            submit_device_1,\n            model_device_1,\n            patch_sizes_list,\n            overlap_sizes,\n            tomo_threshold,\n            act_mode_list,\n        ),\n    )\n)\nfor thread in thread_list:\n    thread.start()\nfor thread in thread_list:\n    thread.join()\n\nif len(task_runs)>0:\n    # throw error\n    print(\"task_runs is not empty\")\n    raise ValueError(\"task_runs is not empty\")\n\nsubmit_df = pd.concat(submit_device_0 + submit_device_1)\nsubmit_df.insert(loc=0, column=\"id\", value=np.arange(len(submit_df)))\n\nsubmit_df.to_csv(\"submission.csv\", index=False)\nsubmit_df\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:37:05.430867Z","iopub.execute_input":"2025-02-05T08:37:05.431225Z","iopub.status.idle":"2025-02-05T08:41:47.400824Z","shell.execute_reply.started":"2025-02-05T08:37:05.431184Z","shell.execute_reply":"2025-02-05T08:41:47.399865Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!ls","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:41:47.402023Z","iopub.execute_input":"2025-02-05T08:41:47.402308Z","iopub.status.idle":"2025-02-05T08:41:48.473336Z","shell.execute_reply.started":"2025-02-05T08:41:47.402283Z","shell.execute_reply":"2025-02-05T08:41:48.472445Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!cp -r /kaggle/input/hengck-czii-cryo-et-01/* .","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:41:48.474660Z","iopub.execute_input":"2025-02-05T08:41:48.474997Z","iopub.status.idle":"2025-02-05T08:41:49.946295Z","shell.execute_reply.started":"2025-02-05T08:41:48.474962Z","shell.execute_reply":"2025-02-05T08:41:49.945250Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from czii_helper import *\nfrom dataset import *\nfrom scipy.optimize import linear_sum_assignment\nimport matplotlib.pyplot as plt","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:41:49.947701Z","iopub.execute_input":"2025-02-05T08:41:49.948054Z","iopub.status.idle":"2025-02-05T08:41:49.956129Z","shell.execute_reply.started":"2025-02-05T08:41:49.948024Z","shell.execute_reply":"2025-02-05T08:41:49.955274Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nif os.getenv('KAGGLE_IS_COMPETITION_RERUN'):\n    MODE = 'submit'\nelse:\n    MODE = 'local'\n\n\n\n\n\n\n\nvalid_dir ='/kaggle/input/czii-cryo-et-object-identification/train'\nvalid_id = ['TS_6_4', ]\n\ndef do_one_eval(truth, predict, threshold):\n    P=len(predict)\n    T=len(truth)\n\n    if P==0:\n        hit=[[],[]]\n        miss=np.arange(T).tolist()\n        fp=[]\n        metric = [P,T,len(hit[0]),len(miss),len(fp)]\n        return hit, fp, miss, metric\n\n    if T==0:\n        hit=[[],[]]\n        fp=np.arange(P).tolist()\n        miss=[]\n        metric = [P,T,len(hit[0]),len(miss),len(fp)]\n        return hit, fp, miss, metric\n\n    #---\n    distance = predict.reshape(P,1,3)-truth.reshape(1,T,3)\n    distance = distance**2\n    distance = distance.sum(axis=2)\n    distance = np.sqrt(distance)\n    p_index, t_index = linear_sum_assignment(distance)\n\n    valid = distance[p_index, t_index] <= threshold\n    p_index = p_index[valid]\n    t_index = t_index[valid]\n    hit = [p_index.tolist(), t_index.tolist()]\n    miss = np.arange(T)\n    miss = miss[~np.isin(miss,t_index)].tolist()\n    fp = np.arange(P)\n    fp = fp[~np.isin(fp,p_index)].tolist()\n\n    metric = [P,T,len(hit[0]),len(miss),len(fp)] #for lb metric F-beta copmutation\n    return hit, fp, miss, metric\n\n\ndef compute_lb(submit_df, overlay_dir):\n    valid_id = valid_id = ['TS_6_4', ]\n    print(valid_id)\n\n    eval_df = []\n    for id in valid_id:\n        truth = read_one_truth(id, overlay_dir) #=f'{valid_dir}/overlay/ExperimentRuns')\n        id_df = submit_df[submit_df['experiment'] == id]\n        for p in PARTICLE:\n            p = dotdict(p)\n            print('\\r', id, p.name, end='', flush=True)\n            xyz_truth = truth[p.name]\n            xyz_predict = id_df[id_df['particle_type'] == p.name][['x', 'y', 'z']].values\n            hit, fp, miss, metric = do_one_eval(xyz_truth, xyz_predict, p.radius* 0.5)\n            eval_df.append(dotdict(\n                id=id, particle_type=p.name,\n                P=metric[0], T=metric[1], hit=metric[2], miss=metric[3], fp=metric[4],\n            ))\n    print('')\n    eval_df = pd.DataFrame(eval_df)\n    gb = eval_df.groupby('particle_type').agg('sum').drop(columns=['id'])\n    gb.loc[:, 'precision'] = gb['hit'] / gb['P']\n    gb.loc[:, 'precision'] = gb['precision'].fillna(0)\n    gb.loc[:, 'recall'] = gb['hit'] / gb['T']\n    gb.loc[:, 'recall'] = gb['recall'].fillna(0)\n    gb.loc[:, 'f-beta4'] = 17 * gb['precision'] * gb['recall'] / (16 * gb['precision'] + gb['recall'])\n    gb.loc[:, 'f-beta4'] = gb['f-beta4'].fillna(0)\n\n    gb = gb.sort_values('particle_type').reset_index(drop=False)\n    # https://www.kaggle.com/competitions/czii-cryo-et-object-identification/discussion/544895\n    gb.loc[:, 'weight'] = [1, 0, 2, 1, 2, 1]\n    lb_score = (gb['f-beta4'] * gb['weight']).sum() / gb['weight'].sum()\n    return gb, lb_score\n\n\n#debug\nif 1:\n    if MODE=='local':\n    #if 1:\n        submit_df=pd.read_csv(\n           'submission.csv'\n            # '/kaggle/input/hengck-czii-cryo-et-weights-01/submission.csv'\n        )\n        gb, lb_score = compute_lb(submit_df, f'{valid_dir}/overlay/ExperimentRuns')\n        print(gb)\n        print('lb_score:',lb_score)\n        print('')\n\n\n        #show one ----------------------------------\n        fig = plt.figure(figsize=(18, 8))\n\n        id = valid_id[0]\n        truth = read_one_truth(id,overlay_dir=f'{valid_dir}/overlay/ExperimentRuns')\n\n        submit_df = submit_df[submit_df['experiment']==id]\n        for p in PARTICLE:\n            p = dotdict(p)\n            xyz_truth = truth[p.name]\n            xyz_predict = submit_df[submit_df['particle_type']==p.name][['x','y','z']].values\n            hit, fp, miss, _ = do_one_eval(xyz_truth, xyz_predict, p.radius)\n            print(id, p.name)\n            print('\\t num truth   :',len(xyz_truth) )\n            print('\\t num predict :',len(xyz_predict) )\n            print('\\t num hit  :',len(hit[0]) )\n            print('\\t num fp   :',len(fp) )\n            print('\\t num miss :',len(miss) )\n\n            ax = fig.add_subplot(2, 3, p.label, projection='3d')\n            if hit[0]:\n                pt = xyz_predict[hit[0]]\n                ax.scatter(pt[:, 0], pt[:, 1], pt[:, 2], alpha=0.5, color='r')\n                pt = xyz_truth[hit[1]]\n                ax.scatter(pt[:,0], pt[:,1], pt[:,2], s=80, facecolors='none', edgecolors='r')\n            if fp:\n                pt = xyz_predict[fp]\n                ax.scatter(pt[:, 0], pt[:, 1], pt[:, 2], alpha=1, color='k')\n            if miss:\n                pt = xyz_truth[miss]\n                ax.scatter(pt[:, 0], pt[:, 1], pt[:, 2], s=160, alpha=1, facecolors='none', edgecolors='k')\n\n            ax.set_title(f'{p.name} ({p.difficulty})')\n\n        plt.tight_layout()\n        plt.show()\n        \n        #--- \n        zz=0","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:41:49.957335Z","iopub.execute_input":"2025-02-05T08:41:49.957656Z","iopub.status.idle":"2025-02-05T08:41:51.213317Z","shell.execute_reply.started":"2025-02-05T08:41:49.957629Z","shell.execute_reply":"2025-02-05T08:41:51.212496Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nif os.getenv('KAGGLE_IS_COMPETITION_RERUN'):\n    MODE = 'submit'\nelse:\n    MODE = 'local'\n\n\n\n\n\n\n\nvalid_dir ='/kaggle/input/czii-cryo-et-object-identification/train'\nvalid_id = ['TS_6_4', ]\n\ndef do_one_eval(truth, predict, threshold):\n    P=len(predict)\n    T=len(truth)\n\n    if P==0:\n        hit=[[],[]]\n        miss=np.arange(T).tolist()\n        fp=[]\n        metric = [P,T,len(hit[0]),len(miss),len(fp)]\n        return hit, fp, miss, metric\n\n    if T==0:\n        hit=[[],[]]\n        fp=np.arange(P).tolist()\n        miss=[]\n        metric = [P,T,len(hit[0]),len(miss),len(fp)]\n        return hit, fp, miss, metric\n\n    #---\n    distance = predict.reshape(P,1,3)-truth.reshape(1,T,3)\n    distance = distance**2\n    distance = distance.sum(axis=2)\n    distance = np.sqrt(distance)\n    p_index, t_index = linear_sum_assignment(distance)\n\n    valid = distance[p_index, t_index] <= threshold\n    p_index = p_index[valid]\n    t_index = t_index[valid]\n    hit = [p_index.tolist(), t_index.tolist()]\n    miss = np.arange(T)\n    miss = miss[~np.isin(miss,t_index)].tolist()\n    fp = np.arange(P)\n    fp = fp[~np.isin(fp,p_index)].tolist()\n\n    metric = [P,T,len(hit[0]),len(miss),len(fp)] #for lb metric F-beta copmutation\n    return hit, fp, miss, metric\n\n\ndef compute_lb(submit_df, overlay_dir):\n    valid_id = list(submit_df[\"experiment\"].unique())\n    print(valid_id)\n\n    eval_df = []\n    for id in valid_id:\n        truth = read_one_truth(id, overlay_dir) #=f'{valid_dir}/overlay/ExperimentRuns')\n        id_df = submit_df[submit_df['experiment'] == id]\n        for p in PARTICLE:\n            p = dotdict(p)\n            print('\\r', id, p.name, end='', flush=True)\n            xyz_truth = truth[p.name]\n            xyz_predict = id_df[id_df['particle_type'] == p.name][['x', 'y', 'z']].values\n            hit, fp, miss, metric = do_one_eval(xyz_truth, xyz_predict, p.radius* 0.5)\n            eval_df.append(dotdict(\n                id=id, particle_type=p.name,\n                P=metric[0], T=metric[1], hit=metric[2], miss=metric[3], fp=metric[4],\n            ))\n    print('')\n    eval_df = pd.DataFrame(eval_df)\n    gb = eval_df.groupby('particle_type').agg('sum').drop(columns=['id'])\n    gb.loc[:, 'precision'] = gb['hit'] / gb['P']\n    gb.loc[:, 'precision'] = gb['precision'].fillna(0)\n    gb.loc[:, 'recall'] = gb['hit'] / gb['T']\n    gb.loc[:, 'recall'] = gb['recall'].fillna(0)\n    gb.loc[:, 'f-beta4'] = 17 * gb['precision'] * gb['recall'] / (16 * gb['precision'] + gb['recall'])\n    gb.loc[:, 'f-beta4'] = gb['f-beta4'].fillna(0)\n\n    gb = gb.sort_values('particle_type').reset_index(drop=False)\n    # https://www.kaggle.com/competitions/czii-cryo-et-object-identification/discussion/544895\n    gb.loc[:, 'weight'] = [1, 0, 2, 1, 2, 1]\n    lb_score = (gb['f-beta4'] * gb['weight']).sum() / gb['weight'].sum()\n    return gb, lb_score\n\n\n#debug\nif 1:\n    if MODE=='local':\n    #if 1:\n        submit_df=pd.read_csv(\n           'submission.csv'\n            # '/kaggle/input/hengck-czii-cryo-et-weights-01/submission.csv'\n        )\n        gb, lb_score = compute_lb(submit_df, f'{valid_dir}/overlay/ExperimentRuns')\n        print(gb)\n        print('lb_score:',lb_score)\n        print('')\n\n\n        #show one ----------------------------------\n        fig = plt.figure(figsize=(18, 8))\n\n        id = valid_id[0]\n        truth = read_one_truth(id,overlay_dir=f'{valid_dir}/overlay/ExperimentRuns')\n\n        submit_df = submit_df[submit_df['experiment']==id]\n        for p in PARTICLE:\n            p = dotdict(p)\n            xyz_truth = truth[p.name]\n            xyz_predict = submit_df[submit_df['particle_type']==p.name][['x','y','z']].values\n            hit, fp, miss, _ = do_one_eval(xyz_truth, xyz_predict, p.radius)\n            print(id, p.name)\n            print('\\t num truth   :',len(xyz_truth) )\n            print('\\t num predict :',len(xyz_predict) )\n            print('\\t num hit  :',len(hit[0]) )\n            print('\\t num fp   :',len(fp) )\n            print('\\t num miss :',len(miss) )\n\n            ax = fig.add_subplot(2, 3, p.label, projection='3d')\n            if hit[0]:\n                pt = xyz_predict[hit[0]]\n                ax.scatter(pt[:, 0], pt[:, 1], pt[:, 2], alpha=0.5, color='r')\n                pt = xyz_truth[hit[1]]\n                ax.scatter(pt[:,0], pt[:,1], pt[:,2], s=80, facecolors='none', edgecolors='r')\n            if fp:\n                pt = xyz_predict[fp]\n                ax.scatter(pt[:, 0], pt[:, 1], pt[:, 2], alpha=1, color='k')\n            if miss:\n                pt = xyz_truth[miss]\n                ax.scatter(pt[:, 0], pt[:, 1], pt[:, 2], s=160, alpha=1, facecolors='none', edgecolors='k')\n\n            ax.set_title(f'{p.name} ({p.difficulty})')\n\n        plt.tight_layout()\n        plt.show()\n        \n        #--- \n        zz=0","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-05T08:41:51.218431Z","iopub.execute_input":"2025-02-05T08:41:51.218722Z","iopub.status.idle":"2025-02-05T08:41:52.482975Z","shell.execute_reply.started":"2025-02-05T08:41:51.218696Z","shell.execute_reply":"2025-02-05T08:41:52.482114Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}