{"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":9867543,"sourceType":"datasetVersion","datasetId":6040935},{"sourceId":10127593,"sourceType":"datasetVersion","datasetId":6240616},{"sourceId":10445850,"sourceType":"datasetVersion","datasetId":6465904},{"sourceId":211097053,"sourceType":"kernelVersion"}],"dockerImageVersionId":30822,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# CZII YOLO11 Submission Baseline with KDTree\n\nWhen @ITK8191 shared his [great notebook](https://www.kaggle.com/code/itsuki9180/czii-yolo11-submission-baseline) I got inspired by it. My initial idea was to train a 2dUNET and use his approach to aggregate the results. However, I thought this wouldn't be feasible, as the YOLO notebook at that time required 10 hours to run. Because of this I tried to modify it by using KDTree. I managed to run a 2dUNET within the time limit but didn't got great results. I decided to share my results with the yolo model.\n\nLater on @min fuka shared his [multi processing notebook](https://www.kaggle.com/code/minfuka/czii-yolo11-submission-baseline-speed-up-ver) and I added his ideas.\n\n\nTime with KDTree was reduced to ~6500 seconds and with multiprocessing was reduced to ~4500 seconds.\n\n\nIf you think this notebook is good, please upvote the [@ITK8191 original notebook](https://www.kaggle.com/code/itsuki9180/czii-yolo11-submission-baseline), [@min fuka multi processing notebook](https://www.kaggle.com/code/minfuka/czii-yolo11-submission-baseline-speed-up-ver) (and this note)!","metadata":{}},{"cell_type":"markdown","source":"## Update\nThis is the furthest I got using YOLO. I've trained with the additional data, weights for the model can be find here [CZII YOLO L trained with synthetic data\n](https://www.kaggle.com/datasets/sersasj/czii-yolo-l-trained-with-synthetic-data).  \nThe notebook to create with additional data is here [YOLO dataset with syntethic data\n](https://www.kaggle.com/code/sersasj/czii-making-datasets-for-yolo-synthetic-data).  \nModel was trained with TS_5_4, TS_69_2 TS_6_4 TS_6_6 as validation.\n\nParameters (z_distance, zy_distance, first_conf, conf_coef) were found with OPTUNA optimization + various submissions for evaluating. \nPerfomance on validation was 0.73 \n\n```\nScore for TS_5_4: 0.658783812957661,{'apo-ferritin': {'total_tp': 42, 'total_fp': 20, 'total_fn': 2, 'fbeta': 0.9321148825065274}, 'beta-galactosidase': {'total_tp': 5, 'total_fp': 27, 'total_fn': 7, 'fbeta': 0.3794642857142858}, 'ribosome': {'total_tp': 20, 'total_fp': 29, 'total_fn': 10, 'fbeta': 0.6427221172022684}, 'thyroglobulin': {'total_tp': 23, 'total_fp': 104, 'total_fn': 7, 'fbeta': 0.6441515650741352}, 'virus-like-particle': {'total_tp': 11, 'total_fp': 2, 'total_fn': 0, 'fbeta': 0.9894179894179894}} \n\nScore for TS_69_2: 0.8191956150699464,{'apo-ferritin': {'total_tp': 35, 'total_fp': 25, 'total_fn': 0, 'fbeta': 0.9596774193548387}, 'beta-galactosidase': {'total_tp': 13, 'total_fp': 46, 'total_fn': 3, 'fbeta': 0.7015873015873016}, 'ribosome': {'total_tp': 35, 'total_fp': 15, 'total_fn': 2, 'fbeta': 0.926791277258567}, 'thyroglobulin': {'total_tp': 28, 'total_fp': 84, 'total_fn': 6, 'fbeta': 0.7256097560975611}, 'virus-like-particle': {'total_tp': 9, 'total_fp': 1, 'total_fn': 0, 'fbeta': 0.9935064935064936}}\n\nScore for TS_6_4: 0.685180923434018,{'apo-ferritin': {'total_tp': 45, 'total_fp': 34, 'total_fn': 12, 'fbeta': 0.7719475277497477}, 'beta-galactosidase': {'total_tp': 7, 'total_fp': 29, 'total_fn': 5, 'fbeta': 0.5219298245614036}, 'ribosome': {'total_tp': 54, 'total_fp': 59, 'total_fn': 12, 'fbeta': 0.7852865697177076}, 'thyroglobulin': {'total_tp': 24, 'total_fp': 77, 'total_fn': 6, 'fbeta': 0.70223752151463}, 'virus-like-particle': {'total_tp': 8, 'total_fp': 4, 'total_fn': 2, 'fbeta': 0.7906976744186046}}\n\nScore for TS_6_6: 0.7575532250952666,{'apo-ferritin': {'total_tp': 37, 'total_fp': 39, 'total_fn': 2, 'fbeta': 0.8985714285714286}, 'beta-galactosidase': {'total_tp': 8, 'total_fp': 43, 'total_fn': 3, 'fbeta': 0.5991189427312775}, 'ribosome': {'total_tp': 17, 'total_fp': 11, 'total_fn': 6, 'fbeta': 0.7297979797979798}, 'thyroglobulin': {'total_tp': 31, 'total_fp': 120, 'total_fn': 4, 'fbeta': 0.7412095639943742}, 'virus-like-particle': {'total_tp': 19, 'total_fp': 2, 'total_fn': 0, 'fbeta': 0.9938461538461538}}\n```\n\n\n# If you have others ideas or suggestions that might improve this approach feel free to share with the community!","metadata":{}},{"cell_type":"code","source":"!tar xfvz /kaggle/input/ultralytics-for-offline-install/archive.tar.gz\n!pip install --no-index --find-links=./packages ultralytics\n!rm -rf ./packages\n\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","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"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},"outputs":[],"execution_count":null},{"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},"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\n#add by @minfuka\nruns2 = runs[sp:]\nruns2[:5]\n\n#add by @minfuka\nassert torch.cuda.device_count() == 2","metadata":{"trusted":true},"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]\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}\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# add by @sesasj\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            \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, 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            \n            img = volume_8bit[slice_idx]\n            input_image = cv2.resize(np.stack([img]*3, axis=-1), (640, 640))\n\n            results = model.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) # 63/64 because of the resize\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\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            # -------------modified by @sersasj ------------------------\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] # Update\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\n","metadata":{"trusted":true},"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\n\n#add by @minfuka\nfrom concurrent.futures import ProcessPoolExecutor #add by @minfuka\n\n#add by @minfuka\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)\n        subs.append(df)\n    \n    return subs\n\n\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},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#change by @minfuka\nsubmission0 = pd.concat(results[0])\nsubmission1 = pd.concat(results[1])\nsubmission = pd.concat([submission0, submission1]).reset_index(drop=True)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission.insert(0, 'id', range(len(submission)))\nsubmission.to_csv(\"submission.csv\", index=False)\nsubmission.head()","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}