{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.14","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":91249,"databundleVersionId":11294684,"sourceType":"competition"},{"sourceId":12074632,"sourceType":"datasetVersion","datasetId":7600588}],"dockerImageVersionId":30805,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!python -m pip install --no-index --find-links=/kaggle/input/solution-ds-byu-1st-place-metadata/packages -r /kaggle/input/solution-ds-byu-1st-place-metadata/packages/requirements.txt\nprint(\"done\")\n\nQUANTILE_THRESHOLD= 0.545 # 3xTTA","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-06T15:55:04.932475Z","iopub.execute_input":"2025-06-06T15:55:04.93273Z","iopub.status.idle":"2025-06-06T15:55:17.784972Z","shell.execute_reply.started":"2025-06-06T15:55:04.93269Z","shell.execute_reply":"2025-06-06T15:55:17.783897Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport sys\nsys.path.append('/kaggle/input/solution-ds-byu-1st-place-metadata/src')\n\n# Dir to store imgs\nfor path in [\"/tmp/working/\"]:\n    if not os.path.exists(path):\n        os.mkdir(path)\nprint(\"done\")","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-06-06T15:55:17.787775Z","iopub.execute_input":"2025-06-06T15:55:17.788483Z","iopub.status.idle":"2025-06-06T15:55:17.793848Z","shell.execute_reply.started":"2025-06-06T15:55:17.788437Z","shell.execute_reply":"2025-06-06T15:55:17.793014Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport shutil\n\ndef clean_working(directory_path: str = \"/kaggle/working/\"):\n    \"\"\"\n    Clean kaggle output directory.\n    \"\"\"\n    if os.path.exists(directory_path):\n        for item in os.listdir(directory_path):\n            if item == \"submission.csv\":\n                continue\n            item_path = os.path.join(directory_path, item)\n            os.remove(item_path) if os.path.isfile(item_path) else shutil.rmtree(item_path)\n        print(f\"All items in '{directory_path}' have been removed.\")\n    else:\n        print(f\"'{directory_path}' does not exist.\")\n        \nclean_working()\nclean_working(\"/tmp/working/\")\nprint(\"done\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-06T15:55:17.797893Z","iopub.execute_input":"2025-06-06T15:55:17.798162Z","iopub.status.idle":"2025-06-06T15:55:17.808378Z","shell.execute_reply.started":"2025-06-06T15:55:17.798137Z","shell.execute_reply":"2025-06-06T15:55:17.807562Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Infer","metadata":{}},{"cell_type":"code","source":"%%writefile ddp.py\n\nfrom types import SimpleNamespace\nimport os\nimport sys\nimport json\n\nif 'KAGGLE_URL_BASE' not in os.environ:\n    os.environ['CUDA_VISIBLE_DEVICES']= \"2,3\"\nelse:\n    sys.path.append('/kaggle/input/solution-ds-byu-1st-place-metadata/src')\n\nimport glob\nimport pickle\nfrom copy import deepcopy\nfrom tqdm import tqdm\n\nimport pandas as pd \nimport numpy as np\n\nimport torch\nimport torch.nn as nn \nimport torch.nn.functional as F\nfrom torch.amp import autocast, GradScaler\n\nfrom monai.inferers import sliding_window_inference\n\nfrom src.modules.utils import batch_to_device\nfrom src.models.utils import get_model\nfrom src.data._3d import CustomDataset\nfrom src.utils.torch import center_of_mass_3d\n\n\nimport torch.distributed as dist\nimport torch.multiprocessing as mp\nfrom torch.nn.parallel import DistributedDataParallel\nfrom torch.utils.data import DataLoader, DistributedSampler\n\n\ndef run_inference(rank, world_size):\n    # ========== Config ==========\n    if is_kaggle():\n        c= SimpleNamespace()\n        c.working_dir= \"/tmp/working/\"\n        c.model_dir= \"/kaggle/input/solution-ds-byu-1st-place-metadata/byu_models/*epoch400.pt\"\n        c.img_size= (128, 672, 672)\n        c.in_dir= \"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/test/\"\n        c.script_idx= int(os.environ.get(\"SCRIPT_IDX\"))\n        c.local_rank= rank\n    else:\n        raise ValueError(\"Are we in the Matrix?\")\n\n    # ========== Dataframe ===========\n    df= []\n    for tomo_id in os.listdir(c.in_dir):\n        df.append({\"tomo_id\": tomo_id, \"fold\": -1})\n    df= pd.DataFrame(df)\n\n    # ========== Models ===========\n    models= []\n    mpaths= sorted(glob.glob(c.model_dir))\n    total_models= len(mpaths)\n    mid= len(mpaths)//2\n    if c.script_idx == 0:\n        mpaths= mpaths[:mid]\n    else:\n        mpaths= mpaths[mid:]\n        \n    # Note: script runs twice (all models dont fit into vRAM)\n    if c.local_rank == 0: \n        print(\"=\"*25)\n        print(\"Loading {}/{} models..\".format(len(mpaths), total_models))\n        for m in mpaths:\n            print(m)\n            \n    for fpath in mpaths:\n        fpath_cfg= fpath.replace(\".pt\", \".pkl\")\n        \n        # Load cfg\n        with open(fpath_cfg, \"rb\") as f:\n            model_cfg= pickle.load(f)\n\n        model_cfg.local_rank= rank\n        model_cfg.weights_path= fpath\n        model_cfg.data_dir= c.in_dir\n        model_cfg.infer_cfg.sw_batch_size= 2\n        model_cfg.tta= True\n        model_cfg.img_size= c.img_size\n        model_cfg.infer_cfg.overlap= (0.875, 0.25, 0.25)\n\n        # Load model\n        m, _= get_model(model_cfg, inference_mode=True)\n        m= m.to(rank)\n        m= DistributedDataParallel(\n            m, \n            device_ids= [rank], \n            output_device= rank,\n            )\n        m= m.eval()\n        m= m.to(rank)\n\n        models.append({\n            \"model\": m,\n            \"cfg\": model_cfg,\n        })\n        \n    if c.local_rank == 0: \n        print(\"=\"*25)\n    cfg= deepcopy(model_cfg)\n\n    # ========== Datasets / Dataloader ===========\n    test_ds= CustomDataset(\n        cfg=model_cfg, \n        df=df, \n        mode=\"test\",\n        )\n\n    sampler= DistributedSampler(\n        test_ds,\n        num_replicas= world_size,\n        rank= rank,\n    )\n\n    test_dl= torch.utils.data.DataLoader(\n        dataset=test_ds,\n        batch_size=1, \n        num_workers=2,\n        sampler=sampler,\n        pin_memory=False,\n        shuffle=False, \n        drop_last=False,\n    )\n\n    # ========== Inference Loop ===========\n    preds_final= []\n    with torch.no_grad():\n        \n        # ROI weight map (downweights edge predictions)\n        pct = 0.30  # 30% edge\n        z, h, w = c.img_size\n        z_margin = int(z * pct)\n        h_margin = int(h * pct)\n        w_margin = int(w * pct)\n        roi_weight_map = torch.ones((z, h, w), device=rank)  # Initialize everything as 1.0\n        roi_weight_map[:z_margin] = 1e-3  # Top edge\n        roi_weight_map[-z_margin:] = 1e-3  # Bottom edge\n        roi_weight_map[:, :h_margin] = 1e-3  # Left edge\n        roi_weight_map[:, -h_margin:] = 1e-3  # Right edge\n        roi_weight_map[:, :, :w_margin] = 1e-3  # Front edge\n        roi_weight_map[:, :, -w_margin:] = 1e-3  # Back edge\n\n        for batch in tqdm(test_dl):\n            with autocast(cfg.device.type):\n\n                try:\n                    tomo_id= batch.pop(\"tomo_id\")[0]\n                    batch = batch_to_device(batch, device=rank) \n                    batch[\"input\"]= batch[\"input\"].float()\n\n                    # Sliding window\n                    preds= None\n                    for midx, row in enumerate(models):\n                        _preds = sliding_window_inference(\n                            inputs= batch[\"input\"],\n                            roi_size= row[\"cfg\"].roi_size,\n                            predictor= row[\"model\"],\n                            roi_weight_map= roi_weight_map,\n                            **vars(row[\"cfg\"].infer_cfg)\n                        )[0, 0, ...]\n                        _preds= torch.sigmoid(_preds)\n\n                        if preds is None:\n                            preds = _preds\n                        else:\n                            preds += _preds\n\n                    # Save Tensor\n                    preds= preds.half().cpu()\n                    fpath= \"{}{}_{}.pt\".format(c.working_dir, tomo_id, c.script_idx)\n                    torch.save(preds, fpath)\n                    outpath= fpath.replace(\".pt\", \".json\")\n\n                    with open(outpath, \"w+\") as f:\n                        json.dump({\n                            \"tomo_id\": tomo_id,\n                            \"script_id\": c.script_idx,\n                            \"z_shape\": batch[\"z_shape\"][0].item(),\n                            \"y_shape\": batch[\"y_shape\"][0].item(),\n                            \"x_shape\": batch[\"x_shape\"][0].item(),\n                        }, f)\n                except:\n                    import traceback\n                    print(traceback.format_exc())\n                    \n                    print(\"failed tomo_id:\", tomo_id)\n                    pass\n\n        return c, preds_final\n\n\ndef is_kaggle():\n    return os.getenv(\"KAGGLE_URL_BASE\") is not None\n\ndef run_DDP(run_fn, world_size):\n    mp.spawn(run_fn, args=(world_size,), nprocs=world_size, join=True)\n    \ndef setup(rank, world_size):\n    os.environ['MASTER_ADDR'] = 'localhost'\n    os.environ['MASTER_PORT'] = '9931'\n    torch.cuda.set_device(rank)\n    dist.init_process_group(\"nccl\", rank=rank, world_size=world_size)\n\ndef cleanup(rank):\n    dist.destroy_process_group()\n    if rank == 0:\n        print(\"DDP complete.\")\n\ndef run_all(rank, world_size):\n    print(f\"Running DDP code on rank {rank}.\")\n    setup(rank, world_size)\n\n    c, preds= run_inference(rank, world_size)\n        \n    cleanup(rank)\n    return\n\nif __name__ == \"__main__\":\n    # Setting seed for reproducabilty\n    np.random.seed(0)\n    torch.manual_seed(0)\n    torch.cuda.manual_seed(0)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n    \n    n_gpus = torch.cuda.device_count()\n    print(f\"total GPUs: {n_gpus}\")\n    world_size = n_gpus\n    run_DDP(run_all, world_size)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-06T15:55:17.809456Z","iopub.execute_input":"2025-06-06T15:55:17.809695Z","iopub.status.idle":"2025-06-06T15:55:17.827741Z","shell.execute_reply.started":"2025-06-06T15:55:17.809672Z","shell.execute_reply":"2025-06-06T15:55:17.826956Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import time\nstart_time= time.time()\n\nkicr = \"true\" if os.getenv('KAGGLE_IS_COMPETITION_RERUN') else \"\"\nkub = \"true\" if 'KAGGLE_URL_BASE' in os.environ else \"\"\n\n!KAGGLE_IS_COMPETITION_RERUN={kicr} KAGGLE_URL_BASE={kub} SCRIPT_IDX=0 python ddp.py","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-06T15:55:17.828667Z","iopub.execute_input":"2025-06-06T15:55:17.828904Z","iopub.status.idle":"2025-06-06T16:00:14.95255Z","shell.execute_reply.started":"2025-06-06T15:55:17.828882Z","shell.execute_reply":"2025-06-06T16:00:14.951404Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!KAGGLE_IS_COMPETITION_RERUN={kicr} KAGGLE_URL_BASE={kub} SCRIPT_IDX=1 python ddp.py","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-06T16:00:14.955928Z","iopub.execute_input":"2025-06-06T16:00:14.956696Z","iopub.status.idle":"2025-06-06T16:04:23.8895Z","shell.execute_reply.started":"2025-06-06T16:00:14.95666Z","shell.execute_reply":"2025-06-06T16:04:23.888393Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import json\nimport glob\nimport pandas as pd \nimport torch\n\n# =========== Load infer metadata ==========\nd= []\nfpaths= sorted(glob.glob(\"/tmp/working/*.json\"))\nfor fpath in fpaths:\n\n    # Load json        \n    with open(fpath, \"r\") as f:\n        metadata= json.load(f)\n        \n    metadata= metadata | {\"fpath\": fpath.replace(\".json\", \".pt\")}\n    d.append(metadata)\n    \ndf = pd.DataFrame(d)\n\n# Sanity check\n# DDP might predict 2x for same tomo\ndf = df.groupby('tomo_id', as_index=False).agg({\n    'tomo_id': 'first',\n    'script_id': 'first',\n    'z_shape': 'first',\n    'y_shape': 'first',\n    'x_shape': 'first',\n    'fpath': lambda x: list(set(x)),\n})\ndisplay(df)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-06T16:04:23.890905Z","iopub.execute_input":"2025-06-06T16:04:23.89124Z","iopub.status.idle":"2025-06-06T16:04:25.810004Z","shell.execute_reply.started":"2025-06-06T16:04:23.891213Z","shell.execute_reply":"2025-06-06T16:04:25.809136Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from tqdm import tqdm\nimport numpy as np\n\n\n# ========== Ensemble volumes ===========\nsub_rows= []\nfor i, row in tqdm(df.iterrows(), total=len(df)):\n    row= row.to_dict()\n    \n    # Ensemble\n    arr= []\n    for f in row[\"fpath\"]:\n        tmp= torch.load(f, weights_only=False)\n        arr.append(tmp)\n    arr= torch.stack(arr, axis=0)\n    \n    arr= torch.sum(arr, axis=0) # Mean ensemble\n    # arr= arr.median(dim=0).values # Median ensemble\n    # arr= (arr.clamp(min=0.01).log().mean(dim=0)).exp() # Geometric mean ensemble\n\n    # Argmax\n    coords = torch.argmax(arr)\n    coords = torch.unravel_index(coords, arr.shape)\n    coords = (\n        (coords[0].item() + 0.5) / arr.shape[0],\n        (coords[1].item() + 0.5) / arr.shape[1],\n        (coords[2].item() + 0.5) / arr.shape[2], \n        )\n\n    # Add\n    sub_rows.append({\n        \"tomo_id\": row[\"tomo_id\"],\n        \"z\": coords[0] * row[\"z_shape\"],\n        \"y\": coords[1] * row[\"y_shape\"], \n        \"x\": coords[2] * row[\"x_shape\"], \n        \"max\": torch.max(arr).item(),\n    })\n    \nsub= pd.DataFrame(sub_rows)\ndisplay(sub)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-06T16:04:25.810969Z","iopub.execute_input":"2025-06-06T16:04:25.811287Z","iopub.status.idle":"2025-06-06T16:04:26.126338Z","shell.execute_reply.started":"2025-06-06T16:04:25.811262Z","shell.execute_reply":"2025-06-06T16:04:26.125417Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Format Sub","metadata":{}},{"cell_type":"code","source":"# Apply Threshold\ncutoff= sub['max'].quantile(QUANTILE_THRESHOLD)\nsub.loc[sub[\"max\"] <= cutoff, [\"z\", \"y\", \"x\"]]= -1.0\nprint(\"=\"*25)\nprint(\"threshold:\", QUANTILE_THRESHOLD)\nprint(\"cutoff:\", cutoff)\nprint(\"=\"*25)\n\n# Format sub\ncol_map= {\n    \"z\": \"Motor axis 0\",\n    \"y\": \"Motor axis 1\",\n    \"x\": \"Motor axis 2\",\n}\nsub= sub.rename(columns=col_map)\nsub= sub[[\"tomo_id\", \"Motor axis 0\", \"Motor axis 1\", \"Motor axis 2\"]]\nsub.to_csv(\"submission.csv\", index=False)\n\nprint(sub)\n\nclean_working(\"/tmp/images/\")\nclean_working(\"/kaggle/working/\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-06T16:04:26.127482Z","iopub.execute_input":"2025-06-06T16:04:26.127848Z","iopub.status.idle":"2025-06-06T16:04:26.15435Z","shell.execute_reply.started":"2025-06-06T16:04:26.127821Z","shell.execute_reply":"2025-06-06T16:04:26.153395Z"}},"outputs":[],"execution_count":null}]}