{"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":10494763,"sourceType":"datasetVersion","datasetId":6497826},{"sourceId":206640467,"sourceType":"kernelVersion"},{"sourceId":211097053,"sourceType":"kernelVersion"},{"sourceId":217866623,"sourceType":"kernelVersion"}],"dockerImageVersionId":30823,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"### In this notebook, I tried to integrate the original data and synthetic data using yolo training, and the LB score was **0.615**.\n\n### The LB score using only raw data is **0.618**.\n\n### The score for using only synthetic data was **0.682**.","metadata":{}},{"cell_type":"markdown","source":"# Import installation library","metadata":{}},{"cell_type":"code","source":"from IPython.display import clear_output\n!tar xfvz /kaggle/input/czii-ultralytics-for-offline-install-yolo/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-23T06:12:12.781817Z","iopub.execute_input":"2025-01-23T06:12:12.782172Z","iopub.status.idle":"2025-01-23T06:13:34.925714Z","shell.execute_reply.started":"2025-01-23T06:12:12.782147Z","shell.execute_reply":"2025-01-23T06:13:34.924777Z"}},"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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-23T06:13:34.927001Z","iopub.execute_input":"2025-01-23T06:13:34.927326Z","iopub.status.idle":"2025-01-23T06:13:37.275076Z","shell.execute_reply.started":"2025-01-23T06:13:34.927293Z","shell.execute_reply":"2025-01-23T06:13:37.274362Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Load model","metadata":{}},{"cell_type":"code","source":"# Define 2 YOLO models + related parameters to load\nmodel_paths_yolo = [\n    \"/kaggle/input/czii-yolo-l-trained-with-synthetic-data/best_synthetic.pt\", \n    \"/kaggle/input/czii-yolo11-baseline-weight/runs/detect/train/weights/best.pt\"\n]\nmodels_yolo = [YOLO(mp) for mp in model_paths_yolo]\n\nruns_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]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-23T06:13:37.276585Z","iopub.execute_input":"2025-01-23T06:13:37.276882Z","iopub.status.idle":"2025-01-23T06:13:38.941606Z","shell.execute_reply.started":"2025-01-23T06:13:37.276859Z","shell.execute_reply":"2025-01-23T06:13:38.940883Z"}},"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}\nindex_to_particle = {v:k for k,v 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}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-23T06:13:38.942832Z","iopub.execute_input":"2025-01-23T06:13:38.943082Z","iopub.status.idle":"2025-01-23T06:13:38.947706Z","shell.execute_reply.started":"2025-01-23T06:13:38.943060Z","shell.execute_reply":"2025-01-23T06:13:38.946843Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Function realization","metadata":{}},{"cell_type":"code","source":"# UnionFind + aggregator + TTA Ensemble\nclass UnionFind:\n    def __init__(self, size):\n        # Initialize the parent array with each element pointing to itself\n        self.parent = np.arange(size)\n        # Initialize the rank array with all elements set to 0\n        self.rank = np.zeros(size, dtype=int)\n\n    def find(self, u):\n        # If u is not the root of its set, recursively find the root and perform path compression\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        # Find the roots of u and v\n        u_root = self.find(u)\n        v_root = self.find(v)\n        # If u and v are already in the same set, do nothing\n        if u_root == v_root:\n            return\n        # Attach the smaller rank tree under the root of the higher rank tree\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 ranks are the same, increment the rank of the new 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        # Set the initial confidence threshold\n        self.first_conf = first_conf\n        # Set the confidence coefficient for aggregation\n        self.conf_coef = conf_coef\n        # Corresponding to 6 types of confidence thresholds\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        # Calculate the lower and upper percentiles of the volume\n        lower, upper = np.percentile(volume, (0.5, 99.5))\n        # Clip the volume to the calculated range\n        clipped = np.clip(volume, lower, upper)\n        # Scale the clipped volume to the range [0, 255] and convert to uint8\n        scaled = ((clipped - lower) / (upper - lower + 1e-12) * 255).astype(np.uint8)\n        return scaled\n\n    def predict_ensemble_tta(self, models_yolo, image_640, device_no):\n        \"\"\"\n        Perform multi-model, multi-TTA reasoning on the input image (640x640).\n        Original image + horizontal flip + vertical flip + rotate 90 degrees\n        Finally, do a simple NMS rehash of the result (or return it directly to the aggregator).\n        \"\"\"\n        # Lists to store all bounding boxes, confidences, and classes\n        all_boxes = []\n        all_confs = []\n        all_clss = []\n\n        def do_inference(img, invert_func):\n            # Perform inference on the input image using all YOLO models\n            for my_model in models_yolo:\n                res = my_model.predict(img, save=False, imgsz=640,\n                                       conf=self.first_conf, device=device_no, \n                                       batch=1, verbose=False)\n                for r in res:\n                    boxes = r.boxes\n                    # If no bounding boxes are detected, skip to the next iteration\n                    if boxes is None or len(boxes)==0:\n                        continue\n                    # Convert the bounding boxes, confidences, and classes to numpy arrays\n                    xyxy = boxes.xyxy.cpu().numpy()\n                    conf = boxes.conf.cpu().numpy()\n                    clss = boxes.cls.cpu().numpy().astype(int)\n                    # Apply the inverse transform to the bounding boxes\n                    xyxy_orig = invert_func(xyxy)\n                    all_boxes.append(xyxy_orig)\n                    all_confs.append(conf)\n                    all_clss.append(clss)\n\n        # Original image\n        do_inference(image_640, invert_func=lambda x:x)\n\n        # Horizontal flip\n        img_hflip = cv2.flip(image_640, 1)\n        def invert_hflip(xyxy):\n            # Invert the horizontal flip of the bounding boxes\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_inference(img_hflip, invert_func=invert_hflip)\n\n        # Vertical flip\n        img_vflip = cv2.flip(image_640, 0)\n        def invert_vflip(xyxy):\n            # Invert the vertical flip of the bounding boxes\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_inference(img_vflip, invert_func=invert_vflip)\n\n        # Rotate 90 degrees (clockwise)\n        img_rot90 = cv2.rotate(image_640, cv2.ROTATE_90_CLOCKWISE)\n        def invert_rot90(xyxy):\n            # Invert the 90-degree clockwise rotation of the bounding boxes\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\n        do_inference(img_rot90, invert_func=invert_rot90)\n\n        # If no bounding boxes are detected, return None\n        if len(all_boxes)==0:\n            return None\n\n        # Concatenate all bounding boxes, confidences, and classes\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\n        # Use non-maximum suppression to remove duplicate bounding boxes\n        from ultralytics.utils.ops import non_max_suppression\n\n        cat_data = np.concatenate([\n            boxes_cat,       # [N,4]\n            confs_cat[:,None],  # [N,1]\n            clss_cat[:,None]    # [N,1]\n        ], axis=1)  # shape [N,6]\n\n        cat_tensor = torch.from_numpy(cat_data).float().to(device_no)\n        cat_tensor = cat_tensor.unsqueeze(0)  # Add batch dimension => shape [1, N, 6]\n\n        nms_out = non_max_suppression(cat_tensor, iou_thres=0.5, max_det=300)\n        # If no bounding boxes are left after NMS, return None\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\n\n    def make_predictions(self, run_id, models_yolo, device_no=\"0\"):\n        \"\"\"\n        For the volume data of a certain run_id, load each slice -> Ensemble + TTA -> aggregator aggregation -> return df\n        \"\"\"\n        # Path to the volume data\n        volume_path = f'/kaggle/input/czii-cryo-et-object-identification/test/static/ExperimentRuns/{run_id}/VoxelSpacing10.000/denoised.zarr'\n        # Open the volume data\n        vol = zarr.open(volume_path, mode='r')[0]\n        # Convert the volume data to 8-bit\n        vol_8bit = self.convert_to_8bit(vol)\n        # Number of slices in the volume data\n        n_slices = vol_8bit.shape[0]\n\n        # Dictionary to store all detections\n        detections = {\n            'particle_type': [],\n            'confidence': [],\n            'x': [],\n            'y': [],\n            'z': []\n        }\n\n        for i in range(n_slices):\n            # Expand the single-channel slice to a three-channel image\n            img_3ch = np.stack([vol_8bit[i]]*3, axis=-1)\n            # Resize the image to 640x640\n            img_640 = cv2.resize(img_3ch, (640, 640))\n\n            # Perform ensemble and TTA prediction\n            final_nms = self.predict_ensemble_tta(models_yolo, img_640, device_no)\n            # If no detections are made, skip to the next slice\n            if final_nms is None:\n                continue\n\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                # Look up the particle name\n                ptype = index_to_particle[int(cls_)]  \n\n                # Calculate the center coordinates of the bounding box\n                xc = ((x1 + x2)/2.) * 10 * (63/64)\n                yc = ((y1 + y2)/2.) * 10 * (63/64)\n                zc = i * 10 + 5\n\n                # Store the detection information\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        # If no detections are made, return an empty DataFrame\n        if not detections['particle_type']:\n            return pd.DataFrame()\n\n        # Aggregate the detections using DFS/UnionFind\n        df = pd.DataFrame(detections)\n        particle_types = np.array(df['particle_type'])\n        confidences = np.array(df['confidence'])\n        xs = np.array(df['x'])\n        ys = np.array(df['y'])\n        zs = np.array(df['z'])\n\n        aggregated_data = []\n        for idx, particle in enumerate(particle_names):\n            if particle == 'beta-amylase':\n                # Skip beta-amylase particles\n                continue\n\n            # Mask for the current particle type\n            mask = (particle_types == particle)\n            # If no particles of the current type are detected, skip to the next particle\n            if not np.any(mask):\n                continue\n\n            # Extract the confidences, x, y, and z coordinates of the current particle type\n            part_conf = confidences[mask]\n            part_x = xs[mask]\n            part_y = ys[mask]\n            part_z = zs[mask]\n\n            # Stack the coordinates into a single array\n            coords = np.stack([part_x, part_y, part_z], axis=1)\n            # Maximum distance in the z direction\n            z_distance = 30\n            # Maximum distance in the xy plane\n            xy_distance = 20\n            # Maximum Euclidean distance\n            max_distance = math.sqrt(z_distance**2 + xy_distance**2)\n            # Build a KDTree for efficient nearest neighbor search\n            tree = cKDTree(coords)\n            # Find all pairs of points within the maximum distance\n            pairs = tree.query_pairs(r=max_distance, p=2)\n\n            # Initialize the UnionFind data structure\n            uf = UnionFind(len(coords))\n            for (u,v) in pairs:\n                # Calculate the difference in the z direction\n                z_diff = abs(coords[u,2] - coords[v,2])\n                if z_diff>z_distance:\n                    continue\n                # Calculate the difference in the xy plane\n                xy_diff = np.linalg.norm(coords[u,:2] - coords[v,:2])\n                if xy_diff>xy_distance:\n                    continue\n                # Union the two points if they are close enough\n                uf.union(u,v)\n\n            # Find the root of each point\n            roots = np.array([uf.find(k) for k in range(len(coords))])\n            # Find the unique roots, inverse indices, and counts\n            unique_roots, inverse_indices, counts = np.unique(roots, return_inverse=True, return_counts=True)\n            # Calculate the sum of confidences for each cluster\n            conf_sums = np.bincount(inverse_indices, weights=part_conf)\n            # Calculate the aggregated confidences\n            aggregated_confidences = conf_sums / (counts**self.conf_coef)\n            # Minimum number of points per cluster for each particle type\n            cluster_per_particle = [4,1,2,9,4,8]  # Original code\n            # Mask for valid clusters\n            valid_clusters = (counts >= cluster_per_particle[idx]) & (aggregated_confidences > self.particle_confs[idx])\n            # If no valid clusters are found, skip to the next particle\n            if not np.any(valid_clusters):\n                continue\n\n            # IDs of valid clusters\n            cluster_ids = unique_roots[valid_clusters]\n            # Sum of x coordinates for each cluster\n            sum_x = np.bincount(inverse_indices, weights=part_x)\n            # Sum of y coordinates for each cluster\n            sum_y = np.bincount(inverse_indices, weights=part_y)\n            # Sum of z coordinates for each cluster\n            sum_z = np.bincount(inverse_indices, weights=part_z)\n\n            # Calculate the center coordinates of each cluster\n            center_x = sum_x / counts\n            center_y = sum_y / counts\n            center_z = sum_z / counts\n\n            # Keep only the valid clusters\n            center_x = center_x[valid_clusters]\n            center_y = center_y[valid_clusters]\n            center_z = center_z[valid_clusters]\n\n            # Create a DataFrame for the aggregated data\n            aggregated_df = pd.DataFrame({\n                'experiment': [run_id]*len(center_x),\n                'particle_type': [particle]*len(center_x),\n                'x': center_x,\n                'y': center_y,\n                'z': center_z\n            })\n            aggregated_data.append(aggregated_df)\n\n        # If there is aggregated data, concatenate all DataFrames\n        if aggregated_data:\n            return pd.concat(aggregated_data, ignore_index=True)\n        else:\n            return pd.DataFrame()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-23T06:13:38.948662Z","iopub.execute_input":"2025-01-23T06:13:38.948902Z","iopub.status.idle":"2025-01-23T06:13:38.976408Z","shell.execute_reply.started":"2025-01-23T06:13:38.948881Z","shell.execute_reply":"2025-01-23T06:13:38.975378Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"aggregator = PredictionAggregator(first_conf=0.15, conf_coef=0.5)\nsubs = []\nstart_time = time.time()\n\nprint(\"Running on cuda:0 ...\")\nfor r in tqdm(runs, total=len(runs)):\n    df_pred = aggregator.make_predictions(r, models_yolo, device_no=\"cuda:0\")\n    subs.append(df_pred)\n\nend_time = time.time()\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-23T06:23:52.385886Z","iopub.execute_input":"2025-01-23T06:23:52.386370Z","iopub.status.idle":"2025-01-23T06:26:39.171534Z","shell.execute_reply.started":"2025-01-23T06:23:52.386338Z","shell.execute_reply":"2025-01-23T06:26:39.170724Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# submission","metadata":{}},{"cell_type":"code","source":"submission = pd.concat(subs, ignore_index=True)\nsubmission.insert(0, 'id', range(len(submission)))\nsubmission.to_csv(\"submission.csv\", index=False)\nprint(\"Saved submission.csv\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-23T06:13:46.230265Z","iopub.status.idle":"2025-01-23T06:13:46.230558Z","shell.execute_reply":"2025-01-23T06:13:46.230438Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-01-23T06:13:46.231380Z","iopub.status.idle":"2025-01-23T06:13:46.231687Z","shell.execute_reply":"2025-01-23T06:13:46.231558Z"}},"outputs":[],"execution_count":null}]}