{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.11.11"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":91249,"databundleVersionId":11294684,"sourceType":"competition"},{"sourceId":12192538,"sourceType":"datasetVersion","datasetId":7679875},{"sourceId":12192573,"sourceType":"datasetVersion","datasetId":7679902}],"dockerImageVersionId":31041,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true},"papermill":{"default_parameters":{},"duration":468.883529,"end_time":"2025-06-04T11:49:10.301272","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2025-06-04T11:41:21.417743","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Setup\n\n## Configure Session\nOn the right, there is Session Options:\n- Accelerator: GPU T4 x2\n- Language: Python\n- Persistence: No persistence\n- Environment: Pin to original environment\n- Internet: Internet off\n\n\n## Add Inputs\nAdd the following inputs on the right\n- https://www.kaggle.com/datasets/st3v3d/2nd-place-byu-challenge-packages\n- https://www.kaggle.com/datasets/st3v3d/2nd-place-byu-challenge-checkpoints\n- byu-locating-bacterial-flagellar-motors-2025\n","metadata":{"papermill":{"duration":0.002617,"end_time":"2025-06-04T11:41:26.016439","exception":false,"start_time":"2025-06-04T11:41:26.013822","status":"completed"},"tags":[]}},{"cell_type":"code","source":"import os\nos.listdir('/kaggle/input')","metadata":{"papermill":{"duration":0.011132,"end_time":"2025-06-04T11:41:26.029852","exception":false,"start_time":"2025-06-04T11:41:26.018720","status":"completed"},"tags":[],"trusted":true,"execution":{"iopub.status.busy":"2025-06-17T08:07:04.708831Z","iopub.execute_input":"2025-06-17T08:07:04.709075Z","iopub.status.idle":"2025-06-17T08:07:04.717366Z","shell.execute_reply.started":"2025-06-17T08:07:04.709057Z","shell.execute_reply":"2025-06-17T08:07:04.716686Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Install packages","metadata":{"papermill":{"duration":0.001999,"end_time":"2025-06-04T11:41:26.034038","exception":false,"start_time":"2025-06-04T11:41:26.032039","status":"completed"},"tags":[]}},{"cell_type":"code","source":"!pip install --no-index --find-links /kaggle/input/2nd-place-byu-challenge-packages/ /kaggle/input/2nd-place-byu-challenge-packages/*.whl","metadata":{"papermill":{"duration":190.153214,"end_time":"2025-06-04T11:44:36.189256","exception":false,"start_time":"2025-06-04T11:41:26.036042","status":"completed"},"tags":[],"trusted":true,"execution":{"iopub.status.busy":"2025-06-17T08:07:04.718511Z","iopub.execute_input":"2025-06-17T08:07:04.719273Z","iopub.status.idle":"2025-06-17T08:11:11.500129Z","shell.execute_reply.started":"2025-06-17T08:07:04.719252Z","shell.execute_reply":"2025-06-17T08:11:11.499219Z"},"scrolled":true,"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Create Inference Script","metadata":{"papermill":{"duration":0.009366,"end_time":"2025-06-04T11:44:36.208898","exception":false,"start_time":"2025-06-04T11:44:36.199532","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# write the entire inference script to a file\nscript = r\"\"\"\nimport os\nimport argparse\nimport ast\nfrom concurrent.futures import ThreadPoolExecutor\n\nimport torch\nimport numpy as np\nimport pandas as pd\nfrom batchgenerators.utilities.file_and_folder_operations import subdirs, join\nfrom torch.nn.functional import interpolate\n\nfrom nnunetv2.dataset_conversion.kaggle_byu.official_data_to_nnunet import convert_coordinates, load_jpgs\nfrom nnunetv2.inference.predict_from_raw_data import nnUNetPredictor\nfrom nnunetv2.utilities.helpers import empty_cache\n\n\n@torch.inference_mode()\ndef resize_image(image: np.ndarray, edge_length: int, device: torch.device) -> torch.Tensor:\n    zoom = edge_length / max(image.shape)\n    new_shape = [round(s * zoom) for s in image.shape]\n    t = torch.from_numpy(image).to(device).float()\n    t = interpolate(t[None, None], new_shape, mode='area')[0, 0]\n    t = torch.clip(torch.round(t), 0, 255).byte()\n    empty_cache(device)\n    return t\n\n\ndef parse_args():\n    p = argparse.ArgumentParser()\n    p.add_argument('--input-dir',\n                   default=\"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/test\")\n    p.add_argument('--output-file',\n                   default=\"/kaggle/working/submission.csv\",\n                   help=\"output-file path\")\n    p.add_argument('--ckpt-dir')\n    p.add_argument(\n        '--fold',\n        type=ast.literal_eval,\n        help=\"tuple of fold identifiers, e.g. ('all',) or (0,1,2)\"\n    )\n    p.add_argument('--threshold', type=float)\n    p.add_argument('--min-dist', type=int, default=13)\n    p.add_argument('--edge', type=int, default=512)\n    p.add_argument('--gpu-id', type=int, default=0,\n                   help=\"This process's GPU index (0 to num_gpus-1)\")\n    p.add_argument('--num-gpus', type=int, default=1,\n                   help=\"Total number of GPUs being used\")\n    return p.parse_args()\n\ndef main():\n    args = parse_args()\n    DEVICE = torch.device(f\"cuda:{args.gpu_id}\")\n\n    # ensure unique filename per GPU\n    out_path = args.output_file\n\n    all_tomos = sorted(subdirs(args.input_dir, join=False))\n    tomos = all_tomos[args.gpu_id::args.num_gpus]\n\n    pred = nnUNetPredictor(\n        tile_step_size=0.5,\n        use_gaussian=True,\n        use_mirroring=True,\n        perform_everything_on_device=True,\n        device=DEVICE,\n        verbose=False,\n        verbose_preprocessing=False,\n        allow_tqdm=True\n    )\n    pred.initialize_from_trained_model_folder(args.ckpt_dir, args.fold)\n    pred.label_manager._all_labels = [0]\n\n    results = []\n    with ThreadPoolExecutor(max_workers=1) as executor:\n        future = executor.submit(load_jpgs, join(args.input_dir, tomos[0]))\n        for i, tomo in enumerate(tomos):\n            img_np = future.result()\n            if i + 1 < len(tomos):\n                future = executor.submit(load_jpgs, join(args.input_dir, tomos[i+1]))\n\n            orig_shape = img_np.shape\n            img = resize_image(img_np, args.edge, DEVICE).float()\n            img = (img - img.mean()) / img.std()\n\n            out = pred.predict_logits_from_preprocessed_data(img[None], out_device=DEVICE).float()[None]\n            out = torch.sigmoid(out)[0, 0]\n            coords = torch.argwhere((out == torch.max(out)) & (out > args.threshold))\n            ps = [out[tuple(c)].item() for c in coords]\n\n            if len(ps) == 0:\n                results.append({'tomo_id': tomo,\n                                'Motor axis 0': -1,\n                                'Motor axis 1': -1,\n                                'Motor axis 2': -1})\n            else:\n                # all motors equally likely, pick first\n                best = coords[0].tolist()\n                xyz = convert_coordinates([best], img.shape, orig_shape)[0]\n                results.append({'tomo_id': tomo,\n                                'Motor axis 0': xyz[0],\n                                'Motor axis 1': xyz[1],\n                                'Motor axis 2': xyz[2]})\n\n            # free up memory\n            del img, out, coords, ps, img_np\n            empty_cache(DEVICE)\n\n    # write out clean CSV\n    df = pd.DataFrame(results, columns=['tomo_id','Motor axis 0','Motor axis 1','Motor axis 2'])\n    df.to_csv(out_path, index=False)\n    print(f\"Saved predictions to {out_path}\")\n\nif __name__ == \"__main__\":\n    main()\n\"\"\"\nscript_path = '/kaggle/working/inference.py'\nwith open(script_path, 'w') as f:\n    f.write(script)\nprint(f\"Saved {script_path}\")\n","metadata":{"papermill":{"duration":0.019654,"end_time":"2025-06-04T11:44:36.237861","exception":false,"start_time":"2025-06-04T11:44:36.218207","status":"completed"},"tags":[],"trusted":true,"execution":{"iopub.status.busy":"2025-06-17T08:11:11.501382Z","iopub.execute_input":"2025-06-17T08:11:11.501709Z","iopub.status.idle":"2025-06-17T08:11:11.509488Z","shell.execute_reply.started":"2025-06-17T08:11:11.501652Z","shell.execute_reply":"2025-06-17T08:11:11.508713Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Execute","metadata":{"papermill":{"duration":0.009299,"end_time":"2025-06-04T11:44:36.256867","exception":false,"start_time":"2025-06-04T11:44:36.247568","status":"completed"},"tags":[]}},{"cell_type":"markdown","source":"## Spawn two processes with each one GPU attached","metadata":{"papermill":{"duration":0.009298,"end_time":"2025-06-04T11:44:36.275492","exception":false,"start_time":"2025-06-04T11:44:36.266194","status":"completed"},"tags":[]}},{"cell_type":"code","source":"import os, sys, subprocess\n\nos.environ['nnUNet_compile'] = 'True'\nos.environ['torch.backends.cudnn.benchmark'] = 'True'\n\nbase = [\n    sys.executable,\n    \"/kaggle/working/inference.py\",\n    \"--num-gpus\", \"2\",\n    \"--input-dir\", \"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/test\",\n    \"--ckpt-dir\", \"/kaggle/input/2nd-place-byu-challenge-checkpoints/MotorRegressionTrainer_BCEtopK20Loss_moreDA_3_5kep_EDT25__nnUNetResEncUNetMPlans__3d_fullres_bs16_ps128_256_256\",\n    \"--fold\", \"('all', )\", \n    \"--threshold\", \"0.15\",\n]\n\np0 = subprocess.Popen(base + [\"--gpu-id\", \"0\", \"--output-file\", \"/kaggle/working/submission_gpu0.csv\"])\np1 = subprocess.Popen(base + [\"--gpu-id\", \"1\", \"--output-file\", \"/kaggle/working/submission_gpu1.csv\"])\np0.wait(); p1.wait()","metadata":{"papermill":{"duration":273.44299,"end_time":"2025-06-04T11:49:09.727931","exception":false,"start_time":"2025-06-04T11:44:36.284941","status":"completed"},"tags":[],"trusted":true,"execution":{"iopub.status.busy":"2025-06-17T08:11:11.510898Z","iopub.execute_input":"2025-06-17T08:11:11.511093Z","iopub.status.idle":"2025-06-17T08:15:45.295845Z","shell.execute_reply.started":"2025-06-17T08:11:11.511079Z","shell.execute_reply":"2025-06-17T08:15:45.295081Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Merge resulting CSVs","metadata":{"papermill":{"duration":0.011774,"end_time":"2025-06-04T11:49:09.752406","exception":false,"start_time":"2025-06-04T11:49:09.740632","status":"completed"},"tags":[]}},{"cell_type":"code","source":"# write header from first file, then append data (skipping headers) from both\n!(head -n 1 submission_gpu0.csv && tail -n +2 -q submission_gpu0.csv submission_gpu1.csv) > submission.csv","metadata":{"papermill":{"duration":0.18388,"end_time":"2025-06-04T11:49:09.948159","exception":false,"start_time":"2025-06-04T11:49:09.764279","status":"completed"},"tags":[],"trusted":true,"execution":{"iopub.status.busy":"2025-06-17T08:15:45.296824Z","iopub.execute_input":"2025-06-17T08:15:45.297100Z","iopub.status.idle":"2025-06-17T08:15:45.443721Z","shell.execute_reply.started":"2025-06-17T08:15:45.297074Z","shell.execute_reply":"2025-06-17T08:15:45.442835Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# print the final prediction file\n!head submission.csv","metadata":{"papermill":{"duration":0.011906,"end_time":"2025-06-04T11:49:09.971984","exception":false,"start_time":"2025-06-04T11:49:09.960078","status":"completed"},"tags":[],"trusted":true,"execution":{"iopub.status.busy":"2025-06-17T08:15:45.444774Z","iopub.execute_input":"2025-06-17T08:15:45.445020Z","iopub.status.idle":"2025-06-17T08:15:45.563602Z","shell.execute_reply.started":"2025-06-17T08:15:45.444997Z","shell.execute_reply":"2025-06-17T08:15:45.562695Z"}},"outputs":[],"execution_count":null}]}