{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"none","dataSources":[{"sourceId":91249,"databundleVersionId":11294684,"sourceType":"competition"},{"sourceId":11972004,"sourceType":"datasetVersion","datasetId":7447819,"isSourceIdPinned":true}],"dockerImageVersionId":31012,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# 1. References\n- {Notebook} [Parse Data](https://www.kaggle.com/code/andrewjdarley/parse-data)\n- {Discussion} [More Motor Annotations](https://www.kaggle.com/competitions/byu-locating-bacterial-flagellar-motors-2025/discussion/569921)\n- {Discussion} [CryoET Dataset with Pixel Anomalies Corrected](https://www.kaggle.com/competitions/byu-locating-bacterial-flagellar-motors-2025/discussion/575028)\n- {Notebook} [BYU - Download CroyoET Dataset (1/4)](https://www.kaggle.com/code/sunghoshim/byu-download-croyoet-dataset-1-4)","metadata":{}},{"cell_type":"code","source":"!pip install zarr cryoet_data_portal -q","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport shutil\n\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nimport yaml\nimport json\nfrom pathlib import Path\nfrom tqdm.notebook import tqdm\nfrom sklearn.model_selection import train_test_split\n\nimport zarr\nfrom cryoet_data_portal import Client, Dataset, Run","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-05-30T04:59:20.553129Z","iopub.execute_input":"2025-05-30T04:59:20.553448Z","iopub.status.idle":"2025-05-30T04:59:24.403740Z","shell.execute_reply.started":"2025-05-30T04:59:20.553427Z","shell.execute_reply":"2025-05-30T04:59:24.402833Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Dataset Info","metadata":{}},{"cell_type":"code","source":"with open('/kaggle/input/byu-cryoet-yolo-dataset/dataset_info.json', 'r', encoding='utf-8') as f:\n    dataset_info = json.load(f)\nprint(json.dumps(dataset_info, indent=4))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-30T05:00:50.912268Z","iopub.execute_input":"2025-05-30T05:00:50.912647Z","iopub.status.idle":"2025-05-30T05:00:50.924587Z","shell.execute_reply.started":"2025-05-30T05:00:50.912609Z","shell.execute_reply":"2025-05-30T05:00:50.923228Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_dataset = pd.read_csv(\"/kaggle/input/byu-cryoet-yolo-dataset/df_dataset.csv\")\ndf_dataset = df_dataset.drop(['z_diff_prev', 'z_diff_next'], axis=1)\ndf_dataset = df_dataset[df_dataset['split'] == 'test'].copy()\nprint(df_dataset.shape)\ndf_dataset.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-30T05:01:43.300489Z","iopub.execute_input":"2025-05-30T05:01:43.301753Z","iopub.status.idle":"2025-05-30T05:01:43.331055Z","shell.execute_reply.started":"2025-05-30T05:01:43.301708Z","shell.execute_reply":"2025-05-30T05:01:43.330126Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Test set - BYU","metadata":{}},{"cell_type":"code","source":"df_test_byu = df_dataset[df_dataset['dataset'] == 'byu']\ndf_test_byu.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-30T05:01:50.751402Z","iopub.execute_input":"2025-05-30T05:01:50.752669Z","iopub.status.idle":"2025-05-30T05:01:50.761808Z","shell.execute_reply.started":"2025-05-30T05:01:50.752621Z","shell.execute_reply":"2025-05-30T05:01:50.760627Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 1 motor\ntomos_byu = df_test_byu['tomo_id'].to_list()\nlen(tomos_byu)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-30T05:01:53.099257Z","iopub.execute_input":"2025-05-30T05:01:53.099638Z","iopub.status.idle":"2025-05-30T05:01:53.106988Z","shell.execute_reply.started":"2025-05-30T05:01:53.099609Z","shell.execute_reply":"2025-05-30T05:01:53.106167Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 0 motor\ndf_byu = pd.read_csv(\"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/train_labels.csv\")\ndf_byu_bg = df_byu[df_byu['Number of motors']==0].sort_values('tomo_id')\ndf_byu_bg['voxel_spacing_int'] = df_byu_bg['Voxel spacing'].astype(int)\nprint(df_byu_bg.shape)\ndf_byu_bg.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-30T05:01:57.729400Z","iopub.execute_input":"2025-05-30T05:01:57.729880Z","iopub.status.idle":"2025-05-30T05:01:57.759183Z","shell.execute_reply.started":"2025-05-30T05:01:57.729849Z","shell.execute_reply":"2025-05-30T05:01:57.758196Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_train, df_test = train_test_split(df_byu_bg, test_size=0.1, stratify=df_byu_bg['voxel_spacing_int'], random_state=42)\ndf_train.shape, df_test.shape","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-30T05:01:59.590041Z","iopub.execute_input":"2025-05-30T05:01:59.590379Z","iopub.status.idle":"2025-05-30T05:01:59.606998Z","shell.execute_reply.started":"2025-05-30T05:01:59.590353Z","shell.execute_reply":"2025-05-30T05:01:59.605509Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"tomos_bg = df_test['tomo_id'].to_list()\nlen(tomos_bg)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-30T05:02:02.153649Z","iopub.execute_input":"2025-05-30T05:02:02.154837Z","iopub.status.idle":"2025-05-30T05:02:02.164178Z","shell.execute_reply.started":"2025-05-30T05:02:02.154786Z","shell.execute_reply":"2025-05-30T05:02:02.162748Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"tomos_byu = tomos_byu + tomos_bg\nprint(len(tomos_byu))\ntomos_byu[:5]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-30T05:02:07.712998Z","iopub.execute_input":"2025-05-30T05:02:07.713319Z","iopub.status.idle":"2025-05-30T05:02:07.720161Z","shell.execute_reply.started":"2025-05-30T05:02:07.713296Z","shell.execute_reply":"2025-05-30T05:02:07.719133Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Save Images","metadata":{}},{"cell_type":"code","source":"path_test = Path().cwd() / \"test\"\npath_test.mkdir(parents=True, exist_ok=True)\npath_test","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-28T00:15:27.321493Z","iopub.execute_input":"2025-05-28T00:15:27.322241Z","iopub.status.idle":"2025-05-28T00:15:27.328353Z","shell.execute_reply.started":"2025-05-28T00:15:27.322218Z","shell.execute_reply":"2025-05-28T00:15:27.327464Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"path_train = Path(\"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/train\")\npath_train","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-28T00:17:18.594887Z","iopub.execute_input":"2025-05-28T00:17:18.596419Z","iopub.status.idle":"2025-05-28T00:17:18.604141Z","shell.execute_reply.started":"2025-05-28T00:17:18.596383Z","shell.execute_reply":"2025-05-28T00:17:18.603292Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def normalize_slice(slice_data):\n    \"\"\"\n    Normalize slice data using 2nd and 98th percentiles for better contrast\n    \"\"\"\n    p2 = np.percentile(slice_data, 2)\n    p98 = np.percentile(slice_data, 98)\n    clipped_data = np.clip(slice_data, p2, p98)\n    normalized = 255 * (clipped_data - p2) / (p98 - p2)\n    return np.uint8(normalized)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-28T00:17:18.936711Z","iopub.execute_input":"2025-05-28T00:17:18.937007Z","iopub.status.idle":"2025-05-28T00:17:18.942213Z","shell.execute_reply.started":"2025-05-28T00:17:18.936987Z","shell.execute_reply":"2025-05-28T00:17:18.941441Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def process_tomogram_set(tomo_id, cur_idx, total):\n    path_tomo = path_test / tomo_id\n    path_tomo.mkdir(parents=True, exist_ok=True)\n    print(f\"{tomo_id} ...\")\n\n    path_src = path_train / tomo_id\n    \n    for path in tqdm(list(path_src.iterdir()), desc=f'[{cur_idx:02d}/{total}] {tomo_id}'):\n        img = Image.open(path)\n        np_img = np.array(img)\n        np_normalized = normalize_slice(np_img)\n\n        path_dest = path_tomo / path.name\n        Image.fromarray(np_normalized).save(path_dest)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-28T00:28:02.342984Z","iopub.execute_input":"2025-05-28T00:28:02.343281Z","iopub.status.idle":"2025-05-28T00:28:02.348435Z","shell.execute_reply.started":"2025-05-28T00:28:02.343262Z","shell.execute_reply":"2025-05-28T00:28:02.347609Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%time\n\ntotal = len(tomos_byu)\nfor idx, tomo_id in enumerate(tomos_byu, start=1):\n    process_tomogram_set(tomo_id, idx, total)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-28T00:28:36.004284Z","iopub.execute_input":"2025-05-28T00:28:36.004910Z","iopub.status.idle":"2025-05-28T00:28:46.017278Z","shell.execute_reply.started":"2025-05-28T00:28:36.004889Z","shell.execute_reply":"2025-05-28T00:28:46.016531Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# CryoET","metadata":{}},{"cell_type":"code","source":"path_tmp = Path('/kaggle/working/tmp')\n\n# Tmp dir\nif path_tmp.is_dir():\n    shutil.rmtree(path_tmp)\npath_tmp.mkdir()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-30T05:02:30.097876Z","iopub.execute_input":"2025-05-30T05:02:30.098156Z","iopub.status.idle":"2025-05-30T05:02:30.104080Z","shell.execute_reply.started":"2025-05-30T05:02:30.098138Z","shell.execute_reply":"2025-05-30T05:02:30.102941Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_cryoet = df_dataset[df_dataset['dataset'] == 'cryoet'].copy()\nprint(df_cryoet.shape)\ndf_cryoet.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-30T05:02:31.165477Z","iopub.execute_input":"2025-05-30T05:02:31.165860Z","iopub.status.idle":"2025-05-30T05:02:31.188828Z","shell.execute_reply.started":"2025-05-30T05:02:31.165835Z","shell.execute_reply":"2025-05-30T05:02:31.187393Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"tomos_cryoet = df_cryoet['tomo_id'].to_list()\nlen(tomos_cryoet)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-30T05:02:33.074372Z","iopub.execute_input":"2025-05-30T05:02:33.074724Z","iopub.status.idle":"2025-05-30T05:02:33.081917Z","shell.execute_reply.started":"2025-05-30T05:02:33.074667Z","shell.execute_reply":"2025-05-30T05:02:33.080755Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"client = Client()\n\n# ========= Process Run ==========\nfor idx, tomo_id in enumerate(tomos_cryoet, start=1):\n    run = Run.find(client, query_filters=[Run.name == tomo_id])\n    if len(run) == 0:\n        print(\"MISSING: \", tomo_id)\n        continue\n    else:\n        run= run[0]\n\n    # Download tomo\n    try:\n        tomo= run.tomograms[0]\n        tomo.download_omezarr(dest_path=path_tmp)\n\n        # Load tomo\n        path_zarr = path_tmp / f\"{tomo_id}.zarr\"\n        arr = zarr.open(path_zarr, mode='r')\n        arr = arr['0']  # API changed?\n        shape = arr.shape\n\n        path_tomo = path_test / tomo_id\n        path_tomo.mkdir(parents=True, exist_ok=True)\n        print(f\"{tomo_id} ...\")\n        for i in tqdm(range(shape[0]), desc=f'[{idx:02d}/{len(tomos_cryoet)}] {tomo_id}'):\n            np_img = arr[i]\n            np_normalized = normalize_slice(np_img)\n            slice_filename = f\"slice_{i:04d}.jpg\"\n            path_dest = path_tomo / slice_filename\n            Image.fromarray(np_normalized).save(path_dest)\n    \n    except Exception as e:\n        print(e)\n        print(\"FAILED:\", tomo_id)\n\n    # Clear tmp\n    shutil.rmtree(path_tmp)\n    path_tmp.mkdir()\n\n    print(f\"[{idx:04d} / {len(tomos_cryoet)}] Done. {tomo_id}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-28T00:17:24.465832Z","iopub.execute_input":"2025-05-28T00:17:24.466738Z","iopub.status.idle":"2025-05-28T00:18:18.880078Z","shell.execute_reply.started":"2025-05-28T00:17:24.466713Z","shell.execute_reply":"2025-05-28T00:18:18.879153Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# solution_df","metadata":{}},{"cell_type":"code","source":"df_solution_byu = df_byu[df_byu['tomo_id'].isin(tomos_byu)].copy()\ndf_solution_byu['Has motor'] = (df_solution_byu['Number of motors'] > 0) + 0\nprint(df_solution_byu.shape)\ndf_solution_byu.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-30T05:02:44.666411Z","iopub.execute_input":"2025-05-30T05:02:44.666737Z","iopub.status.idle":"2025-05-30T05:02:44.690262Z","shell.execute_reply.started":"2025-05-30T05:02:44.666716Z","shell.execute_reply":"2025-05-30T05:02:44.688365Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_cryoet['Has motor'] = (df_cryoet['Number of motors'] > 0) + 0\ndf_cryoet = df_cryoet[df_solution_byu.columns]\nprint(df_cryoet.shape)\ndf_cryoet.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-30T05:02:48.702468Z","iopub.execute_input":"2025-05-30T05:02:48.703384Z","iopub.status.idle":"2025-05-30T05:02:48.721138Z","shell.execute_reply.started":"2025-05-30T05:02:48.703348Z","shell.execute_reply":"2025-05-30T05:02:48.720041Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_solution = pd.concat([df_solution_byu, df_cryoet])\nprint(df_solution.shape)\ndf_solution","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-30T05:02:52.467508Z","iopub.execute_input":"2025-05-30T05:02:52.467858Z","iopub.status.idle":"2025-05-30T05:02:52.488910Z","shell.execute_reply.started":"2025-05-30T05:02:52.467833Z","shell.execute_reply":"2025-05-30T05:02:52.487752Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_solution.to_csv('solution.csv', index=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-30T05:02:58.914013Z","iopub.execute_input":"2025-05-30T05:02:58.915071Z","iopub.status.idle":"2025-05-30T05:02:58.926803Z","shell.execute_reply.started":"2025-05-30T05:02:58.915040Z","shell.execute_reply":"2025-05-30T05:02:58.925782Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}