{"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":10494763,"sourceType":"datasetVersion","datasetId":6497826},{"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":"<h1 style=\"font-size: 28px; color: #D53A35; text-align: center; font-family: 'Arial Black', Gadget, sans-serif;\">\n  In the notebook here, I will use the synthetic data-trained YOLO and the Original data-trained YOLO model, and then add UNet to the stack\n</h1>\n","metadata":{}},{"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-20T07:09:00.709458Z","iopub.execute_input":"2025-01-20T07:09:00.709819Z","iopub.status.idle":"2025-01-20T07:09:32.706300Z","shell.execute_reply.started":"2025-01-20T07:09:00.709790Z","shell.execute_reply":"2025-01-20T07:09:32.705533Z"}},"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 json\nimport copick\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 cc3d\nfrom monai.data import CacheDataset\nfrom monai.transforms import Compose, 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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-20T07:09:32.708101Z","iopub.execute_input":"2025-01-20T07:09:32.708330Z","iopub.status.idle":"2025-01-20T07:09:34.817376Z","shell.execute_reply.started":"2025-01-20T07:09:32.708311Z","shell.execute_reply":"2025-01-20T07:09:34.816439Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# YOLO(Original data)","metadata":{}},{"cell_type":"code","source":"# Load a YOLO model (best_v11.pt)\nmodel_v11 = YOLO(\"/kaggle/input/czii-yolo11-baseline-weight/runs/detect/train/weights/best.pt\")\n\nruns_v11 = sorted(glob.glob('/kaggle/input/czii-cryo-et-object-identification/test/static/ExperimentRuns/*'))\nruns_v11 = [os.path.basename(x) for x in runs_v11]\n\nparticle_names_v11 = [\n    'apo-ferritin', \n    'beta-amylase', \n    'beta-galactosidase', \n    'ribosome', \n    'thyroglobulin', \n    'virus-like-particle'\n]\n\np2i_dict_v11 = {\n    'apo-ferritin': 0,\n    'beta-amylase': 1,\n    'beta-galactosidase': 2,\n    'ribosome': 3,\n    'thyroglobulin': 4,\n    'virus-like-particle': 5\n}\ni2p_v11 = {v: k for k, v in p2i_dict_v11.items()}\n\nparticle_radius_v11 = {\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\n# Here is the main class PredAggForYOLO in the first paragraph, \n# but the fixed variable inside is changed to *_v11 version to prevent conflicts\n\nclass PredAggForYOLO:\n    def __init__(self, first_conf=0.2, final_conf=0.3, conf_coef=0.75):\n        self.first_conf = first_conf\n        self.final_conf = final_conf\n        self.conf_coef = conf_coef\n        self.particle_confs = [0.5, 0.0, 0.2, 0.5, 0.2, 0.5]\n\n    def convert_to_8bit(self, x):\n        lower, upper = np.percentile(x, (0.5, 99.5))\n        x = np.clip(x, lower, upper)\n        x = (x - x.min()) / (x.max() - x.min() + 1e-12) * 255\n        return x.round().astype(\"uint8\")\n\n    def dfs(self, v):\n        self.passed[v] = True\n        self.conf_sum += self.pdf.iloc[v].confidence\n        self.cx += self.pdf.iloc[v].x\n        self.cy += self.pdf.iloc[v].y\n        self.cz += self.pdf.iloc[v].z\n        self.nv += 1\n        for next_v in self.adjacency_list[v]:\n            if (self.passed[next_v]):\n                continue\n            self.dfs(next_v)\n\n    def make_predict_yolo(self, r, model_in):\n        vol = zarr.open(\n            f'/kaggle/input/czii-cryo-et-object-identification/test/static/ExperimentRuns/{r}/VoxelSpacing10.000/denoised.zarr', \n            mode='r'\n        )\n        vol = vol[0]\n        vol2 = self.convert_to_8bit(vol)\n        n_imgs = vol2.shape[0]\n    \n        df = pd.DataFrame()\n        pts = []\n        confs = []\n        xs = []\n        ys = []\n        zs = []\n        \n        for i in range(n_imgs):\n            tmp_img = np.zeros((630, 630))\n            tmp_img[:] = vol2[i]\n    \n            inp_arr = np.stack([tmp_img]*3, axis=-1)\n            inp_arr = cv2.resize(inp_arr, (640,640))\n            \n            # Use the model_in passed in to do the reasoning\n            res = model_in.predict(\n                inp_arr, save=False, imgsz=640, \n                conf=self.first_conf, device=\"0\", \n                batch=1, verbose=False\n            )\n            \n            for j, result in enumerate(res):\n                boxes = result.boxes\n                for k in range(len(boxes.cls)):\n                    ptype = i2p_v11[boxes.cls.cpu().numpy()[k]]\n                    conf = boxes.conf.cpu().numpy()[k]\n                    xc = (boxes.xyxy[k,0] + boxes.xyxy[k,2]) / 2.0 * 10 * (63/64)\n                    yc = (boxes.xyxy[k,1] + boxes.xyxy[k,3]) / 2.0 * 10 * (63/64)\n                    zc = i * 10 + 5\n    \n                    pts.append(ptype)\n                    confs.append(conf)\n                    xs.append(xc.cpu().numpy())\n                    ys.append(yc.cpu().numpy())\n                    zs.append(zc)\n                \n        df['particle_type'] = pts\n        df['confidence'] = confs\n        df['x'] = xs\n        df['y'] = ys\n        df['z'] = zs\n\n        df = df.sort_values(['particle_type', 'z'], ascending=[True, True])\n    \n        agg_df = []\n        for pidx, p in enumerate(particle_names_v11):\n            if p == 'beta-amylase':\n                continue\n            pdf = df[df['particle_type'] == p].reset_index(drop=True)\n            self.pdf = pdf\n            p_rad = particle_radius_v11[p]\n\n            xy_tol = p_rad / 16.0\n            xy_tol_p2 = xy_tol ** 2\n\n            self.adjacency_list = [[] for _ in range(len(pdf))]\n            self.passed = [False for _ in range(len(pdf))]\n\n            for i in range(len(pdf)):\n                x1 = pdf['x'].iloc[i]\n                y1 = pdf['y'].iloc[i]\n                z1 = pdf['z'].iloc[i]\n                for j in range(i+1, len(pdf), 1):\n                    x2 = pdf['x'].iloc[j]\n                    y2 = pdf['y'].iloc[j]\n                    z2 = pdf['z'].iloc[j]\n                    if abs(z1-z2) > 20:\n                        break\n    \n                    dist_p2 = (x1 - x2)**2 + (y1 - y2)**2\n                    if dist_p2 < xy_tol_p2 and dist_p2 + (z1 - z2)**2 < p_rad**2 and abs(z1 - z2) <= 20:\n                        self.adjacency_list[i].append(j)\n                        self.adjacency_list[j].append(i)\n\n            rdf = pd.DataFrame()\n            cxs = []\n            cys = []\n            czs = []\n            for i in range(len(pdf)):\n                self.conf_sum = 0\n                self.nv = 0\n                self.cx = 0\n                self.cy = 0\n                self.cz = 0\n                if not self.passed[i]:\n                    self.dfs(i)\n\n                if self.nv >= 2 and (self.conf_sum / (self.nv**self.conf_coef) > self.particle_confs[pidx]):\n                    cxs.append(self.cx / self.nv)\n                    cys.append(self.cy / self.nv)\n                    czs.append(self.cz / self.nv)\n\n            rdf['experiment'] = [r] * len(cxs)\n            rdf['particle_type'] = [p] * len(cxs)\n            rdf['x'] = cxs\n            rdf['y'] = cys\n            rdf['z'] = czs\n\n            agg_df.append(rdf)\n\n        return pd.concat(agg_df, axis=0) if len(agg_df) > 0 else pd.DataFrame()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-20T07:09:34.818769Z","iopub.execute_input":"2025-01-20T07:09:34.819101Z","iopub.status.idle":"2025-01-20T07:10:36.692695Z","shell.execute_reply.started":"2025-01-20T07:09:34.819069Z","shell.execute_reply":"2025-01-20T07:10:36.691797Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"agent_v11 = PredAggForYOLO(first_conf=0.15, final_conf=0.2, conf_coef=0.5)\nsubs_v11 = []\n\ntick_v11 = time.time()\nfor r in tqdm(runs_v11, total=len(runs_v11)):\n    df = agent_v11.make_predict_yolo(r, model_v11)\n    subs_v11.append(df)\ntock_v11 = time.time()\n\nprint(f'[v11 YOLO] total inference time = {tock_v11 - tick_v11:.2f} seconds')\nsubmission_v11 = pd.concat(subs_v11).reset_index(drop=True)\nsubmission_v11.insert(0, 'id', range(len(submission_v11)))\nprint(submission_v11.head())\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# YOLO(synthetic data)","metadata":{}},{"cell_type":"code","source":"model_path_v2 = '/kaggle/input/czii-yolo-l-trained-with-synthetic-data/best_synthetic.pt'\nmodel_v2 = YOLO(model_path_v2)\n\nruns_path_v2 = '/kaggle/input/czii-cryo-et-object-identification/test/static/ExperimentRuns/*'\nruns_v2 = sorted(glob.glob(runs_path_v2))\nruns_v2 = [os.path.basename(run) for run in runs_v2]\n\nparticle_names = [\n    'apo-ferritin',\n    'beta-amylase',\n    'beta-galactosidase',\n    'ribosome',\n    'thyroglobulin',\n    'virus-like-particle'\n]\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}\nindex_to_particle = {index: name for name, index in particle_to_index.items()}\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\nclass 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        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, model_in, device_no):\n        volume_path = f'/kaggle/input/czii-cryo-et-object-identification/test/static/ExperimentRuns/{run_id}/VoxelSpacing10.000/denoised.zarr'\n        volume = zarr.open(volume_path, mode='r')[0]\n        volume_8bit = self.convert_to_8bit(volume)\n        num_slices = volume_8bit.shape[0]\n\n        detections = {\n            'particle_type': [],\n            'confidence': [],\n            'x': [],\n            'y': [],\n            'z': []\n        }\n\n        for slice_idx in range(num_slices):\n            img = volume_8bit[slice_idx]\n            input_image = cv2.resize(np.stack([img]*3, axis=-1), (640, 640))\n\n            results = model_in.predict(\n                input_image,\n                save=False,\n                imgsz=640,\n                conf=self.first_conf,\n                device=device_no,   \n                batch=1,\n                verbose=False,\n            )\n\n            for result in results:\n                boxes = result.boxes\n                if boxes is None:\n                    continue\n                cls = boxes.cls.cpu().numpy().astype(int)\n                conf = boxes.conf.cpu().numpy()\n                xyxy = boxes.xyxy.cpu().numpy()\n\n                xc = ((xyxy[:, 0] + xyxy[:, 2]) / 2.0) * 10 * (63/64)\n                yc = ((xyxy[:, 1] + xyxy[:, 3]) / 2.0) * 10 * (63/64)\n                zc = np.full(xc.shape, slice_idx * 10 + 5)\n\n                particle_types = [index_to_particle[c] for c in cls]\n\n                detections['particle_type'].extend(particle_types)\n                detections['confidence'].extend(conf)\n                detections['x'].extend(xc)\n                detections['y'].extend(yc)\n                detections['z'].extend(zc)\n\n        if not detections['particle_type']:\n            return pd.DataFrame()\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        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            z_distance = 30 \n            xy_distance = 20\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            uf = UnionFind(len(coords))\n            coords_xy = coords[:, :2]\n            coords_z = coords[:, 2]\n\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                xy_diff = np.linalg.norm(coords_xy[u] - coords_xy[v])\n                if xy_diff > xy_distance:\n                    continue\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]  # self-adjusting\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            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\naggregator = PredictionAggregator(first_conf=0.19, conf_coef=0.34)\n\ndef inference(runs_local, model_local, device_no):\n    subs = []\n    for r in tqdm(runs_local, total=len(runs_local)):\n        df = aggregator.make_predictions(r, model_local, device_no)\n        subs.append(df)\n    return subs\n    ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-20T07:10:36.693603Z","iopub.execute_input":"2025-01-20T07:10:36.693900Z","iopub.status.idle":"2025-01-20T07:11:05.131080Z","shell.execute_reply.started":"2025-01-20T07:10:36.693875Z","shell.execute_reply":"2025-01-20T07:11:05.130120Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"start_time = time.time()\n\n# Here unified serial directly on the same GPU\nresults_v2_single = inference(runs_v2, model_v2, device_no=\"0\")\n\nend_time = time.time()\nestimated_total_time = (end_time - start_time) / len(runs_v2) * 500\nprint(f\"[v2 YOLO single-gpu] estimated total prediction time for 500 runs: {estimated_total_time:.4f} seconds\")\n\nsubmission_v2 = pd.concat(results_v2_single).reset_index(drop=True)\nsubmission_v2.insert(0, 'id', range(len(submission_v2)))\n\nprint(\"submission_v2 shape:\", submission_v2.shape)\nprint(submission_v2.head())\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# YOLO fusion","metadata":{}},{"cell_type":"code","source":"# Stack the results of \"YOLO(Original data)\" and \"YOLO(synthetic data)\"\nfinal_submission_yolo = pd.concat(\n    [submission_v11, submission_v2], \n    ignore_index=True\n)\nprint(\"final_submission_yolo.shape =\", final_submission_yolo.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-20T07:11:05.142113Z","iopub.execute_input":"2025-01-20T07:11:05.142430Z","iopub.status.idle":"2025-01-20T07:11:05.164258Z","shell.execute_reply.started":"2025-01-20T07:11:05.142397Z","shell.execute_reply":"2025-01-20T07:11:05.163386Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# UNet","metadata":{}},{"cell_type":"code","source":"# Subsequent 3D Unet(Monai) inference + final fusion\n\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\n\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    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    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    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    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[x:x + patch_size, y:y + patch_size, z:z + patch_size]\n                    patches.append(patch)\n                    coordinates.append((x, y, z))\n    return patches, coordinates\n\ndef reconstruct_array(patches: List[np.ndarray], \n                     coordinates: List[Tuple[int, int, int]], \n                     original_shape: Tuple[int, int, int]) -> np.ndarray:\n    reconstructed = np.zeros(original_shape, dtype=np.int64)\n    patch_size = patches[0].shape[0]\n    \n    for patch, (x, y, z) in zip(patches, coordinates):\n        reconstructed[x:x + patch_size, y:y + patch_size, z:z + patch_size] = patch\n    return reconstructed\n\ndef calculate_patch_starts(dimension_size: int, patch_size: int) -> List[int]:\n    if dimension_size <= patch_size:\n        return [0]\n    n_patches = np.ceil(dimension_size / patch_size)\n    if n_patches == 1:\n        return [0]\n    \n    total_overlap = (n_patches * patch_size - dimension_size) / (n_patches - 1)\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:\n            positions.append(pos)\n    return positions\n\ndef dict_to_df(coord_dict, experiment_name):\n    all_coords = []\n    all_labels = []\n    for label, coords in coord_dict.items():\n        all_coords.append(coords)\n        all_labels.extend([label] * len(coords))\n    all_coords = np.vstack(all_coords)\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    return df\n\nTRAIN_DATA_DIR = \"/kaggle/input/create-numpy-dataset-exp-name\"\ncopick_config_path = TRAIN_DATA_DIR + \"/copick.config\"\nwith open(copick_config_path) as f:\n    copick_config = json.load(f)\ncopick_config['static_root'] = '/kaggle/input/czii-cryo-et-object-identification/test/static'\ncopick_test_config_path = 'copick_test.config'\nwith open(copick_test_config_path, 'w') as outfile:\n    json.dump(copick_config, outfile)\n\nroot = copick.from_file(copick_test_config_path)\ncopick_user_name = \"copickUtils\"\ncopick_segmentation_name = \"paintedPicks\"\nvoxel_size = 10\ntomo_type = \"denoised\"\n\nfrom monai.transforms import (\n    Compose, \n    EnsureChannelFirstd, \n    NormalizeIntensityd,\n    Orientationd\n)\ninference_transforms = Compose([\n    EnsureChannelFirstd(keys=[\"image\"], channel_dim=\"no_channel\"),\n    NormalizeIntensityd(keys=\"image\"),\n    Orientationd(keys=[\"image\"], axcodes=\"RAS\")\n])\n\n\nid_to_name = {\n    1: \"apo-ferritin\", \n    2: \"beta-amylase\",\n    3: \"beta-galactosidase\", \n    4: \"ribosome\", \n    5: \"thyroglobulin\", \n    6: \"virus-like-particle\"\n}\nBLOB_THRESHOLD = 200\nCERTAINTY_THRESHOLD = 0.05\nclasses = [1, 2, 3, 4, 5, 6]\n\ndef load_models(model_paths):\n    models_ = []\n    for mp in model_paths:\n        model_ = Model(channels=channels, strides=strides_pattern, num_res_units=num_res_units)\n        weights = torch.load(mp)['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]\nmodels_3d = load_models(model_paths)\n\ndef ensemble_prediction_tta(model_, input_tensor, threshold=0.5):\n    # TTA\n    data_copy0 = torch.flip(input_tensor.clone(), dims=[2])\n    data_copy1 = torch.flip(input_tensor.clone(), dims=[3])\n    data_copy2 = torch.flip(input_tensor.clone(), dims=[4])\n    data_copy3 = input_tensor.clone().rot90(1, dims=[3, 4])\n    \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        \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        \n        avg_probs = torch.mean(torch.stack([probs0, probs1, probs2, probs3]), dim=0)\n        thresh_probs = avg_probs > threshold\n        _, max_classes = thresh_probs.max(dim=0)\n    return max_classes\n\n# 3D Unet inference\nsub_3d_list = []\nfor model_3d in models_3d:\n    with torch.no_grad():\n        location_df_all = []\n        for run in root.runs:\n            tomo = run.get_voxel_spacing(10).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            \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(model_3d, input_tensor, threshold=CERTAINTY_THRESHOLD)\n                pred_masks.append(max_classes.cpu().numpy())\n            \n            reconstructed_mask = reconstruct_array(pred_masks, coordinates, tomo.shape)\n            \n            location_dict = {}\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_dict[id_to_name[c]] = xyz\n            \n            df_loc = dict_to_df(location_dict, run.name)\n            location_df_all.append(df_loc)\n        location_df_all = pd.concat(location_df_all)\n        location_df_all.insert(loc=0, column='id', value=np.arange(len(location_df_all)))\n    sub_3d_list.append(location_df_all)\n\nlocation_df_3d = sub_3d_list[0]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-20T07:11:05.165846Z","iopub.execute_input":"2025-01-20T07:11:05.166068Z","iopub.status.idle":"2025-01-20T07:12:27.804757Z","shell.execute_reply.started":"2025-01-20T07:11:05.166044Z","shell.execute_reply":"2025-01-20T07:12:27.803842Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Finally fusion","metadata":{}},{"cell_type":"code","source":"# Final fusion (splicing YOLO results with 3D Unet results and doing some de-duplication/clustering)\n\ndf_for_blend = pd.concat([final_submission_yolo, location_df_3d], ignore_index=True)\n#print(\"Blend shape:\", df_for_blend.shape)\n\nparticle_names_final = [\n    'apo-ferritin', \n    'beta-amylase', \n    'beta-galactosidase', \n    'ribosome', \n    'thyroglobulin', \n    'virus-like-particle'\n]\nparticle_radius_final = {\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_list = []\nfor pidx, p in enumerate(particle_names_final):\n    pdf = df_for_blend[df_for_blend['particle_type'] == p].reset_index(drop=True)\n    p_rad = particle_radius_final[p]\n    \n    grouped = pdf.groupby(['experiment'])\n    for exp, group in grouped:\n        group = group.reset_index(drop=True)\n        coords = group[['x', 'y', 'z']].values\n        if len(coords) <= 1:\n            final_list.append(group)\n            continue\n        \n        db = DBSCAN(eps=p_rad, min_samples=2, metric='euclidean').fit(coords)\n        labels = db.labels_\n        group['cluster'] = labels\n        \n        for cluster_id in np.unique(labels):\n            if cluster_id == -1:\n                continue\n            cluster_points = group[group['cluster'] == cluster_id]\n            avg_x = cluster_points['x'].mean()\n            avg_y = cluster_points['y'].mean()\n            avg_z = cluster_points['z'].mean()\n            group.loc[group['cluster'] == cluster_id, ['x', 'y', 'z']] = avg_x, avg_y, avg_z\n        \n        group = group.drop_duplicates(subset=['x', 'y', 'z'])\n        group = group.drop(columns=['cluster'])\n        final_list.append(group)\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-20T07:12:27.805826Z","iopub.execute_input":"2025-01-20T07:12:27.806530Z","iopub.status.idle":"2025-01-20T07:12:30.235652Z","shell.execute_reply.started":"2025-01-20T07:12:27.806504Z","shell.execute_reply":"2025-01-20T07:12:30.234732Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# submission","metadata":{}},{"cell_type":"code","source":"df_save = pd.concat(final_list, ignore_index=True).reset_index(drop=True)\ndf_save['id'] = np.arange(len(df_save))\ndf_save = df_save[['id', 'experiment', 'particle_type', 'x', 'y', 'z']]\ndf_save.to_csv('submission.csv', index=False)\n\nprint(\"[Done] Final submission saved as submission.csv\")\ndf_save.head()","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}