{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":84969,"databundleVersionId":10033515,"sourceType":"competition"},{"sourceId":9862305,"sourceType":"datasetVersion","datasetId":6052780},{"sourceId":9867543,"sourceType":"datasetVersion","datasetId":6040935},{"sourceId":10445850,"sourceType":"datasetVersion","datasetId":6465904},{"sourceId":10471985,"sourceType":"datasetVersion","datasetId":6484063},{"sourceId":206640467,"sourceType":"kernelVersion"},{"sourceId":211097053,"sourceType":"kernelVersion"}],"dockerImageVersionId":30823,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Import installation library","metadata":{}},{"cell_type":"code","source":"from IPython.display import clear_output\n!tar xfvz /kaggle/input/ultralytics-for-offline-install/archive.tar.gz\n!pip install --no-index --find-links=./packages ultralytics\n!rm -rf ./packages\ntry:\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 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\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 *\nclear_output()","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-01-24T03:39:53.338341Z","iopub.execute_input":"2025-01-24T03:39:53.338794Z","iopub.status.idle":"2025-01-24T03:40:29.578863Z","shell.execute_reply.started":"2025-01-24T03:39:53.338749Z","shell.execute_reply":"2025-01-24T03:40:29.578020Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport glob\nimport time\nimport sys\nimport warnings\nimport math\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport cv2\nimport torch\nfrom tqdm import tqdm\nfrom ultralytics import YOLO\nimport zarr\nfrom scipy.spatial import cKDTree\nfrom collections import defaultdict\nimport torch\nimport cc3d\nfrom monai.transforms import EnsureType\nfrom torch import nn\nfrom monai.networks.nets import UNet\nfrom monai.losses import TverskyLoss\nfrom monai.metrics import DiceMetric\nfrom sklearn.cluster import DBSCAN\n\nfrom typing import List, Tuple, Union\nfrom monai.data import DataLoader, Dataset, CacheDataset, decollate_batch\nimport json\nimport copick\n\nfrom monai.transforms import (\n    Compose, \n    EnsureChannelFirstd, \n    Orientationd,  \n    AsDiscrete,  \n    RandFlipd, \n    RandRotate90d, \n    NormalizeIntensityd,\n    RandCropByLabelClassesd,\n)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T03:40:29.580256Z","iopub.execute_input":"2025-01-24T03:40:29.580644Z","iopub.status.idle":"2025-01-24T03:40:36.073944Z","shell.execute_reply.started":"2025-01-24T03:40:29.580603Z","shell.execute_reply":"2025-01-24T03:40:36.073214Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# YOLO","metadata":{}},{"cell_type":"code","source":"model_path = '/kaggle/input/czii-yolo-l-trained-with-synthetic-data/best_synthetic.pt'\nmodel = YOLO(model_path)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T03:40:36.075711Z","iopub.execute_input":"2025-01-24T03:40:36.076549Z","iopub.status.idle":"2025-01-24T03:40:36.658044Z","shell.execute_reply.started":"2025-01-24T03:40:36.076522Z","shell.execute_reply":"2025-01-24T03:40:36.657246Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"runs_path = '/kaggle/input/czii-cryo-et-object-identification/test/static/ExperimentRuns/*'\nruns = sorted(glob.glob(runs_path))\nruns = [os.path.basename(run) for run in runs]\nsp = len(runs)//2\nruns1 = runs[:sp]\nruns1[:5]\n\nruns2 = runs[sp:]\nruns2[:5]\n\nassert torch.cuda.device_count() == 2","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T03:40:36.659648Z","iopub.execute_input":"2025-01-24T03:40:36.660008Z","iopub.status.idle":"2025-01-24T03:40:36.699086Z","shell.execute_reply.started":"2025-01-24T03:40:36.659974Z","shell.execute_reply":"2025-01-24T03:40:36.698431Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"particle_names = [\n    'apo-ferritin',\n    'beta-amylase',\n    'beta-galactosidase',\n    'ribosome',\n    'thyroglobulin',\n    'virus-like-particle'\n]\nparticle_to_index = {\n    'apo-ferritin': 0,\n    'beta-amylase': 1,\n    'beta-galactosidase': 2,\n    'ribosome': 3,\n    'thyroglobulin': 4,\n    'virus-like-particle': 5\n}\n\nindex_to_particle = {index: name for name, index in particle_to_index.items()}\n\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}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T03:40:36.699843Z","iopub.execute_input":"2025-01-24T03:40:36.700056Z","iopub.status.idle":"2025-01-24T03:40:36.704796Z","shell.execute_reply.started":"2025-01-24T03:40:36.700037Z","shell.execute_reply":"2025-01-24T03:40:36.704054Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## YOLO TTA","metadata":{}},{"cell_type":"code","source":"# Define predict_ensemble_tta in aggregator\nfrom ultralytics.utils.ops import non_max_suppression  \n\ndef predict_ensemble_tta(single_model, image_np, device, conf_thres=0.2):\n    \"\"\"\n    For a 640x640 numpy image:\n    - Multiple models (list of models)\n    - Multiple TTA (original, hflip, vflip, rot90)\n    Do the NMS for the last time\n    Return: [K,6] => x1,y1,x2,y2,conf,cls\n    \"\"\"\n    all_boxes = []\n    all_confs = []\n    all_clss = []\n\n    def do_infer(img_tta, invert_func):\n        #for m in models:\n            res = single_model.predict(img_tta, \n                            save=False, \n                            imgsz=640, \n                            conf=conf_thres, \n                            device=device, \n                            verbose=False)\n            for r in res:\n                boxes = r.boxes\n                if boxes is None or len(boxes)==0:\n                    continue\n                xyxy = boxes.xyxy.cpu().numpy()\n                confs = boxes.conf.cpu().numpy()\n                clss = boxes.cls.cpu().numpy().astype(int)\n                # Inverse transformation\n                xyxy_orig = invert_func(xyxy)\n                all_boxes.append(xyxy_orig)\n                all_confs.append(confs)\n                all_clss.append(clss)\n\n    # Master drawing\n    do_infer(image_np, invert_func=lambda x: x)\n\n    # Horizontal flip\n    img_hflip = cv2.flip(image_np, 1)\n    def invert_hflip(xyxy):\n        new_ = xyxy.copy()\n        x1 = 640 - xyxy[:,2]\n        x2 = 640 - xyxy[:,0]\n        new_[:,0] = x1\n        new_[:,2] = x2\n        return new_\n    do_infer(img_hflip, invert_func=invert_hflip)\n\n    # Vertical flip\n    img_vflip = cv2.flip(image_np, 0)\n    def invert_vflip(xyxy):\n        new_ = xyxy.copy()\n        y1 = 640 - xyxy[:,3]\n        y2 = 640 - xyxy[:,1]\n        new_[:,1] = y1\n        new_[:,3] = y2\n        return new_\n    do_infer(img_vflip, invert_func=invert_vflip)\n\n    # Rotate 90 degrees (clockwise)\n    img_rot90 = cv2.rotate(image_np, cv2.ROTATE_90_CLOCKWISE)\n    def invert_rot90(xyxy):\n        new_ = xyxy.copy()\n        # (x,y)->(y,640-x)\n        # Inverse transform (x',y')->(640-y',x')\n        x1_old,y1_old = xyxy[:,0], xyxy[:,1]\n        x2_old,y2_old = xyxy[:,2], xyxy[:,3]\n        X1 = 640 - y2_old\n        Y1 = x1_old\n        X2 = 640 - y1_old\n        Y2 = x2_old\n        new_[:,0] = X1\n        new_[:,1] = Y1\n        new_[:,2] = X2\n        new_[:,3] = Y2\n        return new_\n    do_infer(img_rot90, invert_func=invert_rot90)\n\n    if len(all_boxes)==0:\n        return None\n\n    boxes_cat = np.concatenate(all_boxes, axis=0)\n    confs_cat = np.concatenate(all_confs, axis=0)\n    clss_cat  = np.concatenate(all_clss, axis=0)\n    cat_data = np.column_stack([boxes_cat, confs_cat, clss_cat])  # shape [N,6]\n\n    # Need to add batch dimension => [1,N,6]\n    cat_tensor = torch.from_numpy(cat_data).float().unsqueeze(0).to(device)\n    # NMS \n    nms_out = non_max_suppression(cat_tensor, iou_thres=0.5, max_det=300)\n    # nms_out => list length =1(a graph), take nms_out[0]\n    if len(nms_out)==0 or nms_out[0] is None or len(nms_out[0])==0:\n        return None\n    final_nms = nms_out[0].cpu().numpy()  # shape [K,6]\n    return final_nms","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T03:40:36.705638Z","iopub.execute_input":"2025-01-24T03:40:36.705872Z","iopub.status.idle":"2025-01-24T03:40:36.722436Z","shell.execute_reply.started":"2025-01-24T03:40:36.705851Z","shell.execute_reply":"2025-01-24T03:40:36.721613Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class UnionFind:\n    def __init__(self, size):\n        self.parent = np.arange(size)\n        self.rank = np.zeros(size, dtype=int)\n\n    def find(self, u):\n        if self.parent[u] != u:\n            self.parent[u] = self.find(self.parent[u])  \n        return self.parent[u]\n\n    def union(self, u, v):\n        u_root = self.find(u)\n        v_root = self.find(v)\n        if u_root == v_root:\n            return\n            \n        if self.rank[u_root] < self.rank[v_root]:\n            self.parent[u_root] = v_root\n        else:\n            self.parent[v_root] = u_root\n            if self.rank[u_root] == self.rank[v_root]:\n                self.rank[u_root] += 1\n\nclass PredictionAggregator:\n    def __init__(self, first_conf=0.2, conf_coef=0.75):\n        self.first_conf = first_conf\n        self.conf_coef = conf_coef\n        self.particle_confs = np.array([0.5,0.0,0.2,0.5,0.2,0.5])\n\n    def convert_to_8bit(self, volume):\n        lower, upper = np.percentile(volume, (0.5, 99.5))\n        clipped = np.clip(volume, lower, upper)\n        scaled = ((clipped - lower) / (upper - lower + 1e-12) * 255).astype(np.uint8)\n        return scaled\n        \n    def make_predictions(self, run_id, models_yolo, device_no=\"cuda:0\"):\n        volume_path = f'/kaggle/input/czii-cryo-et-object-identification/test/static/ExperimentRuns/{run_id}/VoxelSpacing10.000/denoised.zarr'\n        vol = zarr.open(volume_path, mode='r')[0]\n        vol_8bit = self.convert_to_8bit(vol)\n        n_slices = vol_8bit.shape[0]\n\n        detections = {\n            'particle_type': [], 'confidence': [],\n            'x': [], 'y': [], 'z': []\n        }\n\n        for i in range(n_slices):\n            img = vol_8bit[i]\n            #3-channel\n            img_3ch = np.stack([img]*3, axis=-1)\n            img_640 = cv2.resize(img_3ch, (640,640))\n            # TTA\n            final_nms = predict_ensemble_tta(\n                single_model=model, \n                image_np=img_640,\n                device=device_no,\n                conf_thres=self.first_conf\n            )\n            if final_nms is None:\n                continue\n            # final_nms => shape [K,6] => x1,y1,x2,y2,conf,cls\n            for row in final_nms:\n                x1,y1,x2,y2, conf, cls_ = row\n                ptype = index_to_particle.get(int(cls_), \"unknown\")\n\n                # Computing center\n                xc = ((x1+x2)/2.) * 10 * (63/64)\n                yc = ((y1+y2)/2.) * 10 * (63/64)\n                zc = i*10 + 5\n\n                detections['particle_type'].append(ptype)\n                detections['confidence'].append(conf)\n                detections['x'].append(xc)\n                detections['y'].append(yc)\n                detections['z'].append(zc)\n\n        # cKDTree + unionFind + aggregated_confidences\n        if not detections['particle_type']:\n            return pd.DataFrame()\n        df = pd.DataFrame(detections)\n\n\n        \n        particle_types = np.array(detections['particle_type'])\n        confidences = np.array(detections['confidence'])\n        xs = np.array(detections['x'])\n        ys = np.array(detections['y'])\n        zs = np.array(detections['z'])\n\n        aggregated_data = []\n\n        for idx, particle in enumerate(particle_names):\n            if particle == 'beta-amylase':\n                continue \n\n            mask = (particle_types == particle)\n            if not np.any(mask):\n                continue  \n                \n            particle_confidences = confidences[mask]\n            particle_xs = xs[mask]\n            particle_ys = ys[mask]\n            particle_zs = zs[mask]\n            \n            coords = np.vstack((particle_xs, particle_ys, particle_zs)).T\n\n           \n            z_distance = 30 # How many slices can you \"jump\" to aggregate predictions 10 = 1, 20 = 2...\n            xy_distance = 20 # xy_tol_p2 in original code by ITK8191\n            \n            max_distance = math.sqrt(z_distance**2 + xy_distance**2)\n            tree = cKDTree(coords)            \n            pairs = tree.query_pairs(r=max_distance, p=2)\n\n            \n            uf = UnionFind(len(coords))\n            \n            coords_xy = coords[:, :2]\n            coords_z = coords[:, 2]\n            for u, v in pairs:\n                z_diff = abs(coords_z[u] - coords_z[v])\n                if z_diff > z_distance:\n                    continue  \n\n                xy_diff = np.linalg.norm(coords_xy[u] - coords_xy[v])\n                if xy_diff > xy_distance:\n                    continue  \n\n                uf.union(u, v)\n\n            roots = np.array([uf.find(i) for i in range(len(coords))])\n            unique_roots, inverse_indices, counts = np.unique(roots, return_inverse=True, return_counts=True)\n            conf_sums = np.bincount(inverse_indices, weights=particle_confidences)\n            \n            aggregated_confidences = conf_sums / (counts ** self.conf_coef)\n            cluster_per_particle = [4,1,2,9,4,8]\n            valid_clusters = (counts >= cluster_per_particle[idx]) & (aggregated_confidences > self.particle_confs[idx])\n\n            if not np.any(valid_clusters):\n                continue  \n\n            cluster_ids = unique_roots[valid_clusters]\n\n            centers_x = np.bincount(inverse_indices, weights=particle_xs) / counts\n            centers_y = np.bincount(inverse_indices, weights=particle_ys) / counts\n            centers_z = np.bincount(inverse_indices, weights=particle_zs) / counts\n\n            centers_x = centers_x[valid_clusters]\n            centers_y = centers_y[valid_clusters]\n            centers_z = centers_z[valid_clusters]\n\n            aggregated_df = pd.DataFrame({\n                'experiment': [run_id] * len(centers_x),\n                'particle_type': [particle] * len(centers_x),\n                'x': centers_x,\n                'y': centers_y,\n                'z': centers_z\n            })\n\n            aggregated_data.append(aggregated_df)\n\n        if aggregated_data:\n            return pd.concat(aggregated_data, axis=0)\n        else:\n            return pd.DataFrame()  \n    ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T03:40:36.744576Z","iopub.execute_input":"2025-01-24T03:40:36.744834Z","iopub.status.idle":"2025-01-24T03:40:36.765108Z","shell.execute_reply.started":"2025-01-24T03:40:36.744812Z","shell.execute_reply":"2025-01-24T03:40:36.764341Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# instance main class\naggregator = PredictionAggregator(first_conf=0.19,  conf_coef=0.34) #Update\naggregated_results = []\n\nfrom concurrent.futures import ProcessPoolExecutor \n\ndef inference(runs, model, device_no):\n    subs = []\n    for r in tqdm(runs, total=len(runs)):\n        df = aggregator.make_predictions(r, model, device_no=\"cuda:0\")\n        #df = aggregator.make_predictions(r, [model], device_no=\"cuda:0\")\n        subs.append(df)\n    \n    return subs\nstart_time = time.time()\n\nwith ProcessPoolExecutor(max_workers=2) as executor:\n    results = list(executor.map(inference, (runs1, runs2), (model, model), (\"0\", \"1\")))\n\n\nend_time = time.time()\n\nestimated_total_time = (end_time - start_time) / len(runs) * 500  \nprint(f'estimated total prediction time for 500 runs: {estimated_total_time:.4f} seconds')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T03:41:23.660130Z","iopub.execute_input":"2025-01-24T03:41:23.660550Z","iopub.status.idle":"2025-01-24T03:42:47.388699Z","shell.execute_reply.started":"2025-01-24T03:41:23.660518Z","shell.execute_reply":"2025-01-24T03:42:47.387712Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission0 = pd.concat(results[0])\nsubmission1 = pd.concat(results[1])\nsubmission_ = pd.concat([submission0, submission1]).reset_index(drop=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T03:44:57.695006Z","iopub.execute_input":"2025-01-24T03:44:57.695420Z","iopub.status.idle":"2025-01-24T03:44:57.703886Z","shell.execute_reply.started":"2025-01-24T03:44:57.695392Z","shell.execute_reply":"2025-01-24T03:44:57.702914Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission_.insert(0, 'id', range(len(submission_)))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T03:45:00.582953Z","iopub.execute_input":"2025-01-24T03:45:00.583363Z","iopub.status.idle":"2025-01-24T03:45:00.595698Z","shell.execute_reply.started":"2025-01-24T03:45:00.583330Z","shell.execute_reply":"2025-01-24T03:45:00.594660Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission_.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T03:45:12.100966Z","iopub.execute_input":"2025-01-24T03:45:12.101370Z","iopub.status.idle":"2025-01-24T03:45:12.120488Z","shell.execute_reply.started":"2025-01-24T03:45:12.101340Z","shell.execute_reply":"2025-01-24T03:45:12.119295Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission_.tail()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T03:48:20.834257Z","iopub.execute_input":"2025-01-24T03:48:20.834693Z","iopub.status.idle":"2025-01-24T03:48:20.846564Z","shell.execute_reply.started":"2025-01-24T03:48:20.834665Z","shell.execute_reply":"2025-01-24T03:48:20.845469Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Unet3D","metadata":{}},{"cell_type":"code","source":"class 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\ndef 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\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\n\nTRAIN_DATA_DIR = \"/kaggle/input/create-numpy-dataset-exp-name\"\n\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)\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])\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]\n\ndef load_models(model_paths):\n    models = []\n    for model_path in model_paths:\n        channels = (48, 64, 80, 80)\n        strides_pattern = (2, 2, 1)       \n        num_res_units = 1\n        learning_rate = 1e-3\n        num_epochs = 100\n        model = Model(channels=channels, strides=strides_pattern, num_res_units=num_res_units)\n        \n        weights =torch.load(model_path)['state_dict']\n        model.load_state_dict(weights)\n        model.to('cuda')\n        model.eval()\n        models.append(model)\n    return models\n\nmodel_paths = [\n    '/kaggle/input/cziials-a-230-unet/UNet-Model-val_metric0.450.ckpt',\n]\n\nmodels = load_models(model_paths)\ndef 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\nsub=[]\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)))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T03:45:16.520129Z","iopub.execute_input":"2025-01-24T03:45:16.520500Z","iopub.status.idle":"2025-01-24T03:46:15.721739Z","shell.execute_reply.started":"2025-01-24T03:45:16.520475Z","shell.execute_reply":"2025-01-24T03:46:15.720743Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Finaly Blend","metadata":{}},{"cell_type":"code","source":"df = pd.concat([submission_,location_df], 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    pdf = df[df['particle_type'] == p].reset_index(drop=True)\n    p_rad = particle_radius[p]\n    \n    grouped = pdf.groupby(['experiment'])\n    \n    for exp, group in grouped:\n        group = group.reset_index(drop=True)\n        \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        group['cluster'] = labels\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            avg_x = cluster_points['x'].mean()\n            avg_y = cluster_points['y'].mean()\n            avg_z = cluster_points['z'].mean()\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.append(group)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T03:46:54.506893Z","iopub.execute_input":"2025-01-24T03:46:54.507943Z","iopub.status.idle":"2025-01-24T03:46:56.072698Z","shell.execute_reply.started":"2025-01-24T03:46:54.507909Z","shell.execute_reply":"2025-01-24T03:46:56.071931Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# sumission","metadata":{}},{"cell_type":"code","source":"df_save = pd.concat(final, ignore_index=True)\ndf_save = df_save.drop(columns=['cluster'])\ndf_save = df_save.sort_values(by=['experiment', 'particle_type']).reset_index(drop=True)\ndf_save['id'] = np.arange(0, len(df_save))\ndf_save.to_csv('submission.csv', index=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T03:47:04.048332Z","iopub.execute_input":"2025-01-24T03:47:04.048704Z","iopub.status.idle":"2025-01-24T03:47:04.074448Z","shell.execute_reply.started":"2025-01-24T03:47:04.048662Z","shell.execute_reply":"2025-01-24T03:47:04.073685Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_save.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-24T03:47:06.856425Z","iopub.execute_input":"2025-01-24T03:47:06.856745Z","iopub.status.idle":"2025-01-24T03:47:06.867966Z","shell.execute_reply.started":"2025-01-24T03:47:06.856720Z","shell.execute_reply":"2025-01-24T03:47:06.867132Z"}},"outputs":[],"execution_count":null}]}