{"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":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":10616474,"sourceType":"datasetVersion","datasetId":6572779},{"sourceId":10622345,"sourceType":"datasetVersion","datasetId":6576941},{"sourceId":206640467,"sourceType":"kernelVersion"},{"sourceId":213055282,"sourceType":"kernelVersion"},{"sourceId":215061124,"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-02T15:13:22.566626Z","iopub.execute_input":"2025-02-02T15:13:22.567056Z","iopub.status.idle":"2025-02-02T15:13:22.578005Z","shell.execute_reply.started":"2025-02-02T15:13:22.566997Z","shell.execute_reply":"2025-02-02T15:13:22.577146Z"},"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-02T15:13:22.579996Z","iopub.execute_input":"2025-02-02T15:13:22.580605Z","iopub.status.idle":"2025-02-02T15:13:23.750492Z","shell.execute_reply.started":"2025-02-02T15:13:22.580556Z","shell.execute_reply":"2025-02-02T15:13:23.749248Z"},"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\n! cp /kaggle/input/hrnet-v1/* ./ -r","metadata":{"_kg_hide-output":true,"execution":{"iopub.status.busy":"2025-02-02T15:13:23.753233Z","iopub.execute_input":"2025-02-02T15:13:23.753721Z","iopub.status.idle":"2025-02-02T15:14:07.102954Z","shell.execute_reply.started":"2025-02-02T15:13:23.753690Z","shell.execute_reply":"2025-02-02T15:14:07.101906Z"},"trusted":true},"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 ./\n! cp /kaggle/input/voxresnet-v0/voxresnetV0.py ./","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-02-02T15:14:07.106966Z","iopub.execute_input":"2025-02-02T15:14:07.107203Z","iopub.status.idle":"2025-02-02T15:14:13.206468Z","shell.execute_reply.started":"2025-02-02T15:14:07.107178Z","shell.execute_reply":"2025-02-02T15:14:13.205479Z"}},"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-02T15:14:13.207712Z","iopub.execute_input":"2025-02-02T15:14:13.207960Z","iopub.status.idle":"2025-02-02T15:15:34.359997Z","shell.execute_reply.started":"2025-02-02T15:14:13.207937Z","shell.execute_reply":"2025-02-02T15:15:34.358920Z"},"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-02T15:15:34.361449Z","iopub.execute_input":"2025-02-02T15:15:34.361809Z","iopub.status.idle":"2025-02-02T15:15:53.441267Z","shell.execute_reply.started":"2025-02-02T15:15:34.361772Z","shell.execute_reply":"2025-02-02T15:15:53.440439Z"},"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-02T15:15:53.442672Z","iopub.execute_input":"2025-02-02T15:15:53.443355Z","execution_failed":"2025-02-02T15:16:11.343Z"}},"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":{"execution_failed":"2025-02-02T15:16:11.343Z"},"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":{"execution_failed":"2025-02-02T15:16:11.343Z"}},"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":{"execution_failed":"2025-02-02T15:16:11.343Z"}},"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":{"execution_failed":"2025-02-02T15:16:11.343Z"},"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":{"execution_failed":"2025-02-02T15:16:11.343Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!cp -r /kaggle/input/unet2d3e/* ./","metadata":{"trusted":true,"execution":{"execution_failed":"2025-02-02T15:16:11.343Z"}},"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":{"execution_failed":"2025-02-02T15:16:11.343Z"}},"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":{"execution_failed":"2025-02-02T15:16:11.343Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from monai.networks.nets import UNet,SegResNet,DynUNet\nfrom custom_vnet import CustomVNet\nfrom model2_6c import Net\n# from voxhrnet import build_model\nfrom voxhrnetV2 import build_model as build_model_v2\nfrom load_model import build_model as build_model_v3\nfrom model10_v1 import build_model as build_model_v4\nfrom voxresnetV0 import VoxResNet as VoxResNetV0\n\nbasic_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)\nbasic_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)\nbasic_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\nsegresnet_6c_v1 = SegResNet(\n    in_channels=1,\n    out_channels=6,\n    dropout_prob=0.1,\n    upsample_mode=\"deconv\",\n)\n\nbasic_unet2e3d = Net(\n    out_channels=6,\n    arch=\"resnet18d\",\n    decoder_dim=[80, 80, 64, 32, 16],\n    pretrained=False,\n)\n\nbasic_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\nbasic_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\nbasic_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)\nbasic_voxhrnet_v2 = build_model_v2(1,6,config_v2)\nbasic_voxhrnet_v3 = build_model_v3()\nbisic_voxhrnet_v4 = build_model_v4()\n\nbasic_voxresnet_v0 = VoxResNetV0(in_channels=1, n_classes=7)\n\nbasic_dense_vnet = CustomVNet(in_channels=1, classes=7)","metadata":{"execution":{"execution_failed":"2025-02-02T15:16:11.343Z"},"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 '/kaggle/input/voxresnet-v0/epoch195-step4704-valid_loss0.4258-val_metric0.7914.ckpt',\n] \\\n + ['/kaggle/input/hrnet-v1/epoch188-step4536-valid_loss0.4231-val_metric0.8659.ckpt'] \\\n +['/kaggle/input/unet2e3d-6c/pytorch/default/1/unet2E3D-v1-epoch114-val_loss0.55-val_metric0.53-step2760.ckpt'] \n # + ['/kaggle/input/my-weights/unet3D-epoch173-val_loss0.53-val_metric0.54-step4176.ckpt'] # \\\n # +['best-weights/epoch152-step3672-valid_loss0.4333-val_metric0.7929.ckpt'] \n\n # + ['/kaggle/input/hrnet-v1/epoch188-step4536-valid_loss0.4231-val_metric0.8659.ckpt'] \n\n\n # +['/kaggle/input/my-weights/unet3D-epoch173-val_loss0.53-val_metric0.54-step4176.ckpt']\n\n\n\n # +['/kaggle/input/segresnet-6c/pytorch/default/1/epoch314-val_loss0.54-val_metric0.54-step7560.ckpt']\n\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 # = ['/kaggle/input/voxresnet-v0/epoch195-step4704-valid_loss0.4258-val_metric0.7914.ckpt']","metadata":{"trusted":true,"execution":{"execution_failed":"2025-02-02T15:16:11.344Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import copy\nimport torch\n\nmodel_list = []\nfor 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(bisic_voxhrnet_v4)\n        model.load_state_dict(state_dict)\n        model_list.append(model.to(\"cpu\"))\n        print(\"load bisic_voxhrnet_v4\")\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        pass\n    try:\n        model = copy.deepcopy(basic_voxresnet_v0)\n        model.load_state_dict(state_dict)\n        model_list.append(model.to(\"cpu\"))\n        print(\"load basic_voxresnet_v0\")\n        continue\n    except:\n        pass\n    ","metadata":{"trusted":true,"execution":{"execution_failed":"2025-02-02T15:16:11.344Z"}},"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":{"execution_failed":"2025-02-02T15:16:11.344Z"}},"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":{"execution_failed":"2025-02-02T15:16:11.344Z"},"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":{"execution_failed":"2025-02-02T15:16:11.344Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for idx, model in enumerate(model_list):\n    # model_list[idx] = model_list[idx].model\n    model_list[idx] = model_list[idx].eval().half()\n    model_list[idx] = model_list[idx].cpu()\n    torch.save(model_list[idx].state_dict(),f\"model_{idx}.pt\")\n    # model_list[idx] = torch.compile(model_list[idx], mode=\"reduce-overhead\")\n\ntorch.cuda.empty_cache()","metadata":{"trusted":true,"execution":{"execution_failed":"2025-02-02T15:16:11.344Z"}},"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":{"execution_failed":"2025-02-02T15:16:11.344Z"},"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":{"execution_failed":"2025-02-02T15:16:11.344Z"},"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":{"execution_failed":"2025-02-02T15:16:11.344Z"},"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":{"execution_failed":"2025-02-02T15:16:11.344Z"},"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":{"execution_failed":"2025-02-02T15:16:11.344Z"}},"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":{"execution_failed":"2025-02-02T15:16:11.344Z"},"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":{"execution_failed":"2025-02-02T15:16:11.344Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import threading\n\nlock = threading.Lock()\n\n\ndef build_weights(\n    patch_sizes,\n    overlap_sizes,\n    EDGE_WEIGHT,\n    mask_hard_cases=False,\n    device=\"cuda:0\",\n    weight=1.0,\n):\n    # overlap_sizes = (32, 32, 32)\n    weights = torch.zeros(*patch_sizes, dtype=torch.half, device=device)\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 * weight\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    weight_list=None,\n):\n    name = run.name\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        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(\n                patch_sizes, overlap_sizes, EDGE_WEIGHT, False, device, weight\n            )\n            for patch_sizes, weight in zip(patch_sizes_list, weight_list)\n        ]\n        # coordinates = extract_3d_patches_overlap_wo_data(\n        #     (D, H, W), patch_sizes, overlap_sizes\n        # )\n        # probability = torch.zeros((5, D, H, W), device=device).half()\n        # count = torch.zeros((D, H, W), device=device).half()\n        probability.zero_()\n        count.zero_()\n        for idx, (model, patch_sizes, coordinates, weights) in enumerate(\n            zip(model_list, patch_sizes_list, coordinates_list, weights_list)\n        ):\n            # weights = weights.to(device)\n            for i in range(len(coordinates)):\n                coor = coordinates[i]\n\n                input_tensor_0 = image[\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                    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                model_output_90 = model(input_tensor_90)\n                model_output_180 = model(input_tensor_180)\n                model_output_270 = model(input_tensor_270)\n                model_output_2 = model(input_tensor_2)\n                model_output_3 = model(input_tensor_3)\n                model_output_4 = model(input_tensor_4)\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                # torch.save(probs, f\"probs_{name}_{idx}_{i}.pt\")\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\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\n# def 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# ):\n#     for run in run_list:\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#         )\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    weight_list=None,\n):\n    if weight_list is None:\n        weight_list = [1.0] * len(model_list)\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            weight_list,\n        )","metadata":{"trusted":true,"execution":{"execution_failed":"2025-02-02T15:16:11.344Z"}},"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]\nact_mode_list=[\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\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 = []\nmodel_device_0 = [copy.deepcopy(model).to(\"cuda:0\") for model in model_list]\nmodel_device_1 = [copy.deepcopy(model).to(\"cuda:1\") for model in model_list]\nweight_list = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.5, 0.3]\n\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            weight_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            weight_list,\n        ),\n    )\n)\nfor thread in thread_list:\n    thread.start()\nfor thread in thread_list:\n    thread.join()\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","metadata":{"trusted":true,"execution":{"execution_failed":"2025-02-02T15:16:11.344Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!ls","metadata":{"trusted":true,"execution":{"execution_failed":"2025-02-02T15:16:11.344Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!cp -r /kaggle/input/hengck-czii-cryo-et-01/* .","metadata":{"trusted":true,"execution":{"execution_failed":"2025-02-02T15:16:11.344Z"}},"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":{"execution_failed":"2025-02-02T15:16:11.344Z"}},"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":{"execution_failed":"2025-02-02T15:16:11.344Z"}},"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":{"execution_failed":"2025-02-02T15:16:11.344Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}