{"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":"none","dataSources":[{"sourceId":91249,"databundleVersionId":11294684,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":11257843,"sourceType":"datasetVersion","datasetId":7035785}],"dockerImageVersionId":30918,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os, cv2\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nimport matplotlib.pyplot as plt\nfrom torch.amp import autocast, GradScaler","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"try:\n    import monai\n    print(\"✅ MONAI is already installed.\")\nexcept ImportError:\n    print(\"📦 Installing MONAI from local .whl file...\")\n    !pip install /kaggle/input/monai-1-3/monai-1.3.0-202310121228-py3-none-any.whl\n\nDATA_DIR = \"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025\"\nTRAIN_IMG_DIR = os.path.join(DATA_DIR, \"train\")\ntest_DIR = os.path.join(DATA_DIR,\"test\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 코드 실행 상태가 초기화되었기 때문에, 다시 필요한 라이브러리와 모델 정의 코드를 재실행합니다.\n\nfrom monai.networks.nets import UNet\nimport torch\n\n# 3D U-Net 모델 정의 (MONAI)\ndef get_unet_model(spatial_dims=3, in_channels=1, out_channels=1, pretrained_path=None):\n    \"\"\"\n    MONAI 3D U-Net 정의 및 사전학습 weight 로딩\n\n    Args:\n        spatial_dims: 3D 볼륨이므로 기본값 3음.\n        in_channels: 입력 채널 수 (흑백이면 1)\n        out_channels: 출력 채널 수 (binary segmentation이면 1)\n        pretrained_path: 사전학습된 모델 경로 (.pth 또는 .pt)\n\n    Returns:\n        PyTorch 모델\n    \"\"\"\n    from monai.networks.nets import UNet\n    \n    model = UNet(\n        spatial_dims=spatial_dims, #3이면 3D-Unet\n        in_channels=in_channels, #흑백 이미지면 1, RGB면 3\n        out_channels=out_channels, #예측할 클래스 개수(편모모터 한 종류이므로 1(편모모터의 개수가 아님, 예시로 편모모터, 세포핵 이렇게 2개의 종류 예측이면 2))\n        channels=(16, 32, 64, 128, 256),\n        strides=((1, 2, 2), (1, 2, 2), (1, 2, 2), (1, 2, 2)),  # 3D tuple\n        num_res_units=2,\n        norm='instance' #batch->일반적, instance->배치값이 작을때, group->batch가 1\n    )\n\n\n    if pretrained_path:\n        state_dict = torch.load(pretrained_path, map_location='cpu')\n        model.load_state_dict(state_dict)\n        print(\"✅ 사전학습된 weight 로딩 완료:\", pretrained_path)\n\n    return model\n\n# 테스트용 모델 인스턴스 생성\nunet3d = get_unet_model()\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport torch\nimport torch.nn.functional as F\n\ndef sliding_window_predict(volume, model, patch_size=(64, 256, 256), stride=(32, 128, 128), device='cuda'):\n    \"\"\"\n    Patch 기반으로 3D 볼륨에 대해 예측하는 함수.\n    한 파일을 슬라이싱하는 것이 아닌, 폴더 전체 즉, 3D이미지를 직육면체 모양으로 잘라서 분석한다고 보면된다.\n    \n    Args:\n        volume (np.ndarray): 입력 3D 볼륨 (Z, H, W)\n        model (torch.nn.Module): 학습된 3D U-Net 모델\n        patch_size (tuple): patch 크기 (z, y, x)\n        stride (tuple): 슬라이딩 윈도우 stride (z, y, x), 경계면 부분 정보손실을 막기위해 path_size보다 작게 설정.\n        device (str): 'cuda' or 'cpu'\n        \n    Returns:\n        full_pred (np.ndarray): 예측된 마스크 (Z, H, W)\n    \"\"\"\n    model.eval()\n    model.to(device)\n\n    z_max, y_max, x_max = volume.shape\n    patch_z, patch_y, patch_x = patch_size\n    stride_z, stride_y, stride_x = stride\n\n    # prediction과 counting을 위한 공간\n    output_pred = np.zeros(volume.shape, dtype=np.float32)\n    output_count = np.zeros(volume.shape, dtype=np.float32)\n\n    with torch.no_grad():\n        for z in range(0, z_max - patch_z + 1, stride_z):\n            for y in range(0, y_max - patch_y + 1, stride_y):\n                for x in range(0, x_max - patch_x + 1, stride_x):\n                    patch = volume[z:z + patch_z, y:y + patch_y, x:x + patch_x]\n                    patch_tensor = torch.tensor(patch, dtype=torch.float32).unsqueeze(0).unsqueeze(0).to(device)  # (1, 1, Z, H, W)\n                    \n                    pred = model(patch_tensor)  # (1, 1, Z, H, W)\n                    pred = torch.sigmoid(pred).squeeze().cpu().numpy()  # (Z, H, W)\n\n                    output_pred[z:z + patch_z, y:y + patch_y, x:x + patch_x] += pred\n                    output_count[z:z + patch_z, y:y + patch_y, x:x + patch_x] += 1\n\n    # count로 나눠서 평균\n    full_pred = output_pred / np.maximum(output_count, 1e-6)\n\n    return (full_pred > 0.5).astype(np.uint8)  # binary mask 반환 (0 또는 1)\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nfrom collections import Counter\n\n# 각 tomo 폴더당 슬라이스 개수 세기\nslice_counts = []\nfor tomo_id in os.listdir(TRAIN_IMG_DIR):\n    tomo_path = os.path.join(TRAIN_IMG_DIR, tomo_id)\n    if os.path.isdir(tomo_path):\n        count = len(os.listdir(tomo_path))\n        slice_counts.append(count)\n\n# 유니크한 개수 종류만 보기\nunique_counts = sorted(set(slice_counts))\nprint(\"🔍 Unique Z slice counts:\")\nprint(unique_counts)\n\n# (추가) 분포 보기\nfrom collections import Counter\nprint(\"\\n📊 Count distribution (how many tomos per slice count):\")\nfor k, v in sorted(Counter(slice_counts).items()):\n    print(f\"{k} slices: {v} tomos\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport numpy as np\nfrom PIL import Image\n\ndef load_volume_original(tomo_path):\n    \"\"\"\n    원본 해상도로 3D 볼륨을 로딩하는 함수.\n    (Z, H, W) 형태의 np.ndarray로 반환.\n\n    Args:\n        tomo_path (str): 슬라이스 이미지들이 들어있는 폴더 경로\n\n    Returns:\n        volume (np.ndarray): (Z, H, W) 크기의 3D 배열\n    \"\"\"\n    slice_files = sorted(os.listdir(tomo_path))  # jpg 파일들이 순서대로 있을 것\n    volume = []\n\n    for fname in slice_files:\n        if fname.lower().endswith(\".jpg\"):\n            img_path = os.path.join(tomo_path, fname)\n            img = Image.open(img_path).convert(\"L\")  # 흑백\n            arr = np.array(img)\n            volume.append(arr)\n\n    # 슬라이스들을 z축으로 쌓는다\n    volume = np.stack(volume, axis=0)  # (Z, H, W)\n\n    return volume\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# import os\n# # 전체 tomo 폴더 ID 목록 가져오기\n# tomo_ids = sorted(os.listdir(TRAIN_IMG_DIR))\n\n# for tomo_id in tomo_ids:\n#     tomo_path = os.path.join(TRAIN_IMG_DIR, tomo_id)\n\n#     # tomo_path가 폴더인지 확인 (jpg 파일만 있는 폴더라고 가정)\n#     if os.path.isdir(tomo_path):\n#         volume = load_volume_original(tomo_path)  # (Z, H, W)\n#         print(f\"{tomo_id}: volume shape = {volume.shape}\")\n#         # 여기서 volume으로 학습/추론 로직 연결 가능\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def extract_motor_centers(pred_mask, tomo_id):\n    \"\"\"\n    pred_mask: (Z, H, W) 이진 마스크\n    tomo_id:   현재 tomo ID\n    return:    list of dict [{tomo_id, Motor axis 0, Motor axis 1, Motor axis 2}, ...]\n    \"\"\"\n    results = []\n    labeled, num_components = label(pred_mask)\n\n    for comp_id in range(1, num_components + 1):\n        comp_mask = (labeled == comp_id)\n        zc, yc, xc = center_of_mass(comp_mask)\n        # 실수 좌표 => 반올림\n        zc, yc, xc = round(zc), round(yc), round(xc)\n        results.append({\n            \"tomo_id\": tomo_id,\n            \"Motor axis 0\": zc,\n            \"Motor axis 1\": yc,\n            \"Motor axis 2\": xc,\n        })\n    return results\n\n\n##############################\n# 4) 전체 루프\n##############################\ndef predict_and_save_csv(test_dir, model, output_csv=\"submission.csv\"):\n    \"\"\"\n    test_dir: 각 tomo 폴더가 들어있는 경로 (test set)\n    model:    학습된 3D U-Net\n    \"\"\"\n    tomo_ids = sorted(os.listdir(test_dir))\n    all_results = []\n\n    for tomo_id in tqdm(tomo_ids, desc=\"Predicting all tomos\"):\n        tomo_path = os.path.join(test_dir, tomo_id)\n        if not os.path.isdir(tomo_path):\n            continue\n\n        # 1) 볼륨 로딩\n        volume = load_volume_original(tomo_path)\n\n        # 2) 슬라이딩 윈도우 예측\n        pred_mask = sliding_window_predict(volume, model, patch_size=(64,256,256), stride=(32,128,128), device='cuda')\n\n        # 3) 마스크 → 중심 좌표 추출\n        centers = extract_motor_centers(pred_mask, tomo_id)\n        all_results.extend(centers)\n\n    # 4) 저장\n    df = pd.DataFrame(all_results)\n    df.to_csv(output_csv, index=False)\n    print(f\"✅ Saved {output_csv} with {len(all_results)} lines.\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"TEST_IMG_DIR = \"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/test\"\n\npredict_and_save_csv(\n    test_dir=TEST_IMG_DIR,\n    model=unet3d,\n    output_csv=\"submission.csv\"\n)\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}