{"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":10225818,"sourceType":"datasetVersion","datasetId":6321984},{"sourceId":206640467,"sourceType":"kernelVersion"}],"dockerImageVersionId":30787,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"try:\n    import zarr\nexcept: \n    !cp -r '/kaggle/input/hengck-czii-cryo-et-01/wheel_file' '/kaggle/working/'\n    !pip install /kaggle/working/wheel_file/asciitree-0.3.3/asciitree-0.3.3\n    !pip install --no-index --find-links=/kaggle/working/wheel_file zarr\n    !pip install --no-index --find-links=/kaggle/working/wheel_file connected-components-3d\nfrom datetime import datetime\nimport pytz\nimport sys\nsys.path.append('/kaggle/input/hengck-czii-cryo-et-01')\nfrom czii_helper import *\nfrom dataset import *\nfrom model2 import *","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-20T13:11:13.553212Z","iopub.execute_input":"2024-12-20T13:11:13.553603Z","iopub.status.idle":"2024-12-20T13:12:32.379963Z","shell.execute_reply.started":"2024-12-20T13:11:13.553559Z","shell.execute_reply":"2024-12-20T13:12:32.379097Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from typing import List, Tuple, Union\ndeps_path = '/kaggle/input/czii-cryoet-dependencies'\n! pip install -q --no-index --find-links {deps_path} --requirement {deps_path}/requirements.txt\nimport lightning.pytorch as pl\nclass Model(pl.LightningModule):\n    def __init__(\n        self, \n        spatial_dims: int = 3,\n        in_channels: int = 1,\n        out_channels: int = 7,\n        channels: Union[Tuple[int, ...], List[int]] = (48, 64, 80, 80),\n        strides: Union[Tuple[int, ...], List[int]] = (2, 2, 1),\n        num_res_units: int = 1,\n    ):\n        super().__init__()\n        self.save_hyperparameters()\n        self.model = UNet(\n            spatial_dims=self.hparams.spatial_dims,\n            in_channels=self.hparams.in_channels,\n            out_channels=self.hparams.out_channels,\n            channels=self.hparams.channels,\n            strides=self.hparams.strides,\n            num_res_units=self.hparams.num_res_units,\n        )\n    def forward(self, x):\n        return self.model(x)\n# 设置训练超参数\nchannels = (48, 64, 80, 80)\nstrides_pattern = (2, 2, 1)\nnum_res_units = 1","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-20T13:15:40.273507Z","iopub.execute_input":"2024-12-20T13:15:40.274374Z","iopub.status.idle":"2024-12-20T13:16:08.291785Z","shell.execute_reply.started":"2024-12-20T13:15:40.274333Z","shell.execute_reply":"2024-12-20T13:16:08.290876Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def extract_3d_patches_minimal_overlap(arrays: List[np.ndarray], patch_size: int) -> Tuple[List[np.ndarray], List[Tuple[int, int, int]]]:\n    if not arrays or not isinstance(arrays, list):\n        raise ValueError(\"Input must be a non-empty list of arrays\")\n    \n    # Verify all arrays have the same shape\n    shape = arrays[0].shape\n    if not all(arr.shape == shape for arr in arrays):\n        raise ValueError(\"All input arrays must have the same shape\")\n    \n    if patch_size > min(shape):\n        raise ValueError(f\"patch_size ({patch_size}) must be smaller than smallest dimension {min(shape)}\")\n    \n    m, n, l = shape\n    patches = []\n    coordinates = []\n    \n    # Calculate starting positions for each dimension\n    x_starts = calculate_patch_starts(m, patch_size)\n    y_starts = calculate_patch_starts(n, patch_size)\n    z_starts = calculate_patch_starts(l, patch_size)\n    \n    # Extract patches from each array\n    for arr in arrays:\n        for x in x_starts:\n            for y in y_starts:\n                for z in z_starts:\n                    patch = arr[\n                        x:x + patch_size,\n                        y:y + patch_size,\n                        z:z + patch_size\n                    ]\n                    patches.append(patch)\n                    coordinates.append((x, y, z))\n    \n    return patches, coordinates\ndef reconstruct_array(patches: List[np.ndarray], \n                     coordinates: List[Tuple[int, int, int]], \n                     original_shape: Tuple[int, int, int]) -> np.ndarray:\n    reconstructed = np.zeros(original_shape, dtype=np.int64)  # To track overlapping regions\n    \n    patch_size = patches[0].shape[0]\n    \n    for patch, (x, y, z) in zip(patches, coordinates):\n        reconstructed[\n            x:x + patch_size,\n            y:y + patch_size,\n            z:z + patch_size\n        ] = patch\n        \n    \n    return reconstructed\ndef calculate_patch_starts(dimension_size: int, patch_size: int) -> List[int]:\n    if dimension_size <= patch_size:\n        return [0]\n        \n    # Calculate number of patches needed\n    n_patches = np.ceil(dimension_size / patch_size)\n    \n    if n_patches == 1:\n        return [0]\n    \n    # Calculate overlap\n    total_overlap = (n_patches * patch_size - dimension_size) / (n_patches - 1)\n    \n    # Generate starting positions\n    positions = []\n    for i in range(int(n_patches)):\n        pos = int(i * (patch_size - total_overlap))\n        if pos + patch_size > dimension_size:\n            pos = dimension_size - patch_size\n        if pos not in positions:  # Avoid duplicates\n            positions.append(pos)\n    \n    return positions\nimport pandas as pd\n\ndef dict_to_df(coord_dict, experiment_name):\n    # Create lists to store data\n    all_coords = []\n    all_labels = []\n    \n    # Process each label and its coordinates\n    for label, coords in coord_dict.items():\n        all_coords.append(coords)\n        all_labels.extend([label] * len(coords))\n    \n    # Concatenate all coordinates\n    all_coords = np.vstack(all_coords)\n    \n    df = pd.DataFrame({\n        'experiment': experiment_name,\n        'particle_type': all_labels,\n        'x': all_coords[:, 0],\n        'y': all_coords[:, 1],\n        'z': all_coords[:, 2]\n    })\n\n    \n    return df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-20T13:16:15.380918Z","iopub.execute_input":"2024-12-20T13:16:15.381586Z","iopub.status.idle":"2024-12-20T13:16:15.396909Z","shell.execute_reply.started":"2024-12-20T13:16:15.381552Z","shell.execute_reply":"2024-12-20T13:16:15.395849Z"}},"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)\nTRAIN_DATA_DIR = \"/kaggle/input/create-numpy-dataset-exp-name\"\nimport json\ncopick_config_path = TRAIN_DATA_DIR + \"/copick.config\"\n\nwith open(copick_config_path) as f:\n    copick_config = json.load(f)\n\ncopick_config['static_root'] = '/kaggle/input/czii-cryo-et-object-identification/test/static'\n\ncopick_test_config_path = 'copick_test.config'\n\nwith open(copick_test_config_path, 'w') as outfile:\n    json.dump(copick_config, outfile)\nimport copick\n\nroot = copick.from_file(copick_test_config_path)\n\ncopick_user_name = \"copickUtils\"\ncopick_segmentation_name = \"paintedPicks\"\nvoxel_size = 10\ntomo_type = \"denoised\"\ninference_transforms = Compose([\n    EnsureChannelFirstd(keys=[\"image\"], channel_dim=\"no_channel\"),\n    NormalizeIntensityd(keys=\"image\"),\n    Orientationd(keys=[\"image\"], axcodes=\"RAS\")\n])\nimport cc3d\n\nid_to_name = {1: \"apo-ferritin\", \n              2: \"beta-amylase\",\n              3: \"beta-galactosidase\", \n              4: \"ribosome\", \n              5: \"thyroglobulin\", \n              6: \"virus-like-particle\"}\nBLOB_THRESHOLD = 200\nCERTAINTY_THRESHOLD = 0.05\n\nclasses = [1, 2, 3, 4, 5, 6]\nimport torch\nimport numpy as np\nimport pandas as pd\nimport cc3d\nfrom monai.data import CacheDataset\nfrom monai.transforms import Compose, EnsureType\nfrom torch import nn\nfrom tqdm import tqdm\n\n# 假设你有一个函数来加载模型权重\ndef load_models(model_paths):\n    models = []\n    for model_path in model_paths:\n        model = torch.load(model_path)\n        model.to('cuda')\n        model.eval()\n        models.append(model)\n    return models\n\n# 假设你有3个模型的路径\nmodel_paths = [\n    '/kaggle/input/weight-all/weight/val_on_TS_6_6_best_model.pt',\n    '/kaggle/input/weight-all/weight/val_on_TS_73_6_best_model.pt',\n    '/kaggle/input/weight-all/weight/val_on_TS_69_2_best_model.pt'\n]\n\n# 加载模型\nmodels = load_models(model_paths)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-20T13:16:19.562733Z","iopub.execute_input":"2024-12-20T13:16:19.563101Z","iopub.status.idle":"2024-12-20T13:17:10.674989Z","shell.execute_reply.started":"2024-12-20T13:16:19.563070Z","shell.execute_reply":"2024-12-20T13:17:10.674127Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def ensemble_prediction_tta(models, input_tensor, threshold=0.5):\n    probs_list = []\n    data_copy0 = input_tensor.clone()\n    data_copy0=torch.flip(data_copy0, dims=[2])\n    data_copy1 = input_tensor.clone()\n    data_copy1=torch.flip(data_copy1, dims=[3])\n    data_copy2 = input_tensor.clone()\n    data_copy2=torch.flip(data_copy2, dims=[4])\n    data_copy3 = input_tensor.clone()\n    data_copy3 = data_copy3.rot90(1, dims=[3, 4])\n    with torch.no_grad():\n        model_output0 = model(input_tensor)\n        model_output1 = model(data_copy0)\n        model_output1=torch.flip(model_output1, dims=[2])\n        model_output2 = model(data_copy1)\n        model_output2=torch.flip(model_output2, dims=[3])\n        model_output3 = model(data_copy2)\n        model_output3=torch.flip(model_output3, dims=[4])\n        probs0 = torch.softmax(model_output0[0], dim=0)\n        probs1 = torch.softmax(model_output1[0], dim=0)\n        probs2 = torch.softmax(model_output2[0], dim=0)\n        probs3 = torch.softmax(model_output3[0], dim=0)\n        probs_list.append(probs0)\n        probs_list.append(probs1)\n        probs_list.append(probs2)\n        probs_list.append(probs3)\n    avg_probs = torch.mean(torch.stack(probs_list), dim=0)\n    thresh_probs = avg_probs > threshold\n    _, max_classes = thresh_probs.max(dim=0)\n    return max_classes","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-20T13:17:14.643023Z","iopub.execute_input":"2024-12-20T13:17:14.644553Z","iopub.status.idle":"2024-12-20T13:17:14.653029Z","shell.execute_reply.started":"2024-12-20T13:17:14.644512Z","shell.execute_reply":"2024-12-20T13:17:14.652100Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"sub=[]\nfor model in models:\n    with torch.no_grad():\n        location_df = []\n        for run in root.runs:\n            tomo = run.get_voxel_spacing(10)\n            tomo = tomo.get_tomogram(tomo_type).numpy()\n            tomo_patches, coordinates = extract_3d_patches_minimal_overlap([tomo], 96)\n            tomo_patched_data = [{\"image\": img} for img in tomo_patches]\n            tomo_ds = CacheDataset(data=tomo_patched_data, transform=inference_transforms, cache_rate=1.0)\n            pred_masks = []\n            for i in tqdm(range(len(tomo_ds))):\n                input_tensor = tomo_ds[i]['image'].unsqueeze(0).to(\"cuda\")\n                max_classes = ensemble_prediction_tta(models, input_tensor, threshold=CERTAINTY_THRESHOLD)\n                pred_masks.append(max_classes.cpu().numpy())\n            reconstructed_mask = reconstruct_array(pred_masks, coordinates, tomo.shape)\n            location = {}\n            for c in classes:\n                cc = cc3d.connected_components(reconstructed_mask == c)\n                stats = cc3d.statistics(cc)\n                zyx = stats['centroids'][1:] * 10.012444  # 转换单位\n                zyx_large = zyx[stats['voxel_counts'][1:] > BLOB_THRESHOLD]\n                xyz = np.ascontiguousarray(zyx_large[:, ::-1])\n                location[id_to_name[c]] = xyz\n            df = dict_to_df(location, run.name)\n            location_df.append(df)\n        location_df = pd.concat(location_df)\n        location_df.insert(loc=0, column='id', value=np.arange(len(location_df)))\n        sub.append(location_df)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-20T13:17:17.423116Z","iopub.execute_input":"2024-12-20T13:17:17.423528Z","iopub.status.idle":"2024-12-20T13:20:27.192052Z","shell.execute_reply.started":"2024-12-20T13:17:17.423490Z","shell.execute_reply":"2024-12-20T13:20:27.191166Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nfrom sklearn.cluster import DBSCAN\n\n# 假设sub已经给定，拼接DataFrame\ndf = pd.concat([sub[0], sub[1], sub[2]], ignore_index=True)\n# 粒子半径映射\nparticle_names = ['apo-ferritin', 'beta-amylase', 'beta-galactosidase', 'ribosome', 'thyroglobulin', 'virus-like-particle']\nparticle_radius = {\n    'apo-ferritin': 60,\n    'beta-amylase': 65,\n    'beta-galactosidase': 90,\n    'ribosome': 150,\n    'thyroglobulin': 130,\n    'virus-like-particle': 135,\n}\n\nfinal = []  # 用于存储最终的结果\nfor pidx, p in enumerate(particle_names):\n    # 筛选出该粒子类型的所有点\n    pdf = df[df['particle_type'] == p].reset_index(drop=True)\n    p_rad = particle_radius[p]\n    \n    # 根据 experiment 分组\n    grouped = pdf.groupby(['experiment'])\n    \n    for exp, group in grouped:\n        group = group.reset_index(drop=True)\n        \n        # 使用DBSCAN进行聚类\n        coords = group[['x', 'y', 'z']].values\n        db = DBSCAN(eps=p_rad, min_samples=2, metric='euclidean').fit(coords)\n        labels = db.labels_\n        \n        # 将聚类结果添加到DataFrame中\n        group['cluster'] = labels\n        \n        # 对每个簇进行处理\n        for cluster_id in np.unique(labels):\n            if cluster_id == -1:\n                continue  # 跳过噪声点\n            \n            cluster_points = group[group['cluster'] == cluster_id]\n            \n            # 计算簇的中心（平均位置）\n            avg_x = cluster_points['x'].mean()\n            avg_y = cluster_points['y'].mean()\n            avg_z = cluster_points['z'].mean()\n            \n            # 更新簇内点的位置\n            group.loc[group['cluster'] == cluster_id, ['x', 'y', 'z']] = avg_x, avg_y, avg_z\n            group = group.drop_duplicates(subset=['x', 'y', 'z'])\n        # 将处理后的数据添加到 final 列表\n        final.append(group)\n\n# 合并处理后的数据\ndf_save = pd.concat(final, ignore_index=True)\ndf_save = df_save.drop(columns=['cluster'])\n\n# 排序按 'experiment' 和 'particle_type' 两列\ndf_save = df_save.sort_values(by=['experiment', 'particle_type']).reset_index(drop=True)\n\n# 重新生成 'id' 列，从 1 开始\ndf_save['id'] = np.arange(0, len(df_save))\n\n# 输出结果到 CSV\ndf_save.to_csv('submission.csv', index=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2024-12-20T13:22:00.437583Z","iopub.execute_input":"2024-12-20T13:22:00.438509Z","iopub.status.idle":"2024-12-20T13:22:02.958326Z","shell.execute_reply.started":"2024-12-20T13:22:00.438464Z","shell.execute_reply":"2024-12-20T13:22:02.957533Z"}},"outputs":[],"execution_count":null}]}