{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"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,"sourceType":"competition"}],"dockerImageVersionId":30919,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport shutil\nfrom concurrent.futures import ProcessPoolExecutor\nfrom functools import partial\nfrom glob import glob\nfrom itertools import product\n\nimport cv2\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nimport torch.nn.functional as F\nimport torch","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-06T10:02:46.542713Z","iopub.execute_input":"2025-03-06T10:02:46.543025Z","iopub.status.idle":"2025-03-06T10:02:52.716588Z","shell.execute_reply.started":"2025-03-06T10:02:46.542995Z","shell.execute_reply":"2025-03-06T10:02:52.715296Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"COMP_DIR = \"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025\"\nLABEL = f\"{COMP_DIR}/train_labels.csv\"\nTRAIN_DIR = f\"{COMP_DIR}/train\"\nTAU = 1000\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def stack_slices(slice_paths: list[str]) -> np.ndarray:\n    tomo = []\n    for slice_ in slice_paths:\n        img = cv2.imread(slice_, -1)\n        img = img / 255\n        tomo.append(img)\n    return np.array(tomo)\n\n\ndef set_label(array: np.ndarray, radius: int, z: int, y: int, x: int) -> np.ndarray:\n    if z <= 0 or y <= 0 or x <= 0:\n        return array\n\n    shape = array.shape\n\n    z, y, x = int(z), int(y), int(x)\n\n    z_min = max(0, z - radius)\n    z_max = min(shape[0], z + radius + 1)\n    y_min = max(0, y - radius)\n    y_max = min(shape[1], y + radius + 1)\n    x_min = max(0, x - radius)\n    x_max = min(shape[2], x + radius + 1)\n\n    for zi, yi, xi in product(\n        range(z_min, z_max), range(y_min, y_max), range(x_min, x_max)\n    ):\n        if (zi - z) ** 2 + (yi - y) ** 2 + (xi - x) ** 2 <= radius**2:\n            array[zi, yi, xi] += 1\n    return array\n\n\ndef create_label_array(\n    tomo_shape: tuple[int, int, int], coords: np.ndarray, radius: int\n) -> np.ndarray:\n    label = np.zeros(tomo_shape)\n    for coord in coords:\n        z = coord[0]\n        y = coord[1]\n        x = coord[2]\n        label = set_label(label, radius, z, y, x)\n\n    label = np.clip(label, 0, 1)\n\n    return label\n\n\ndef scale_array(array: np.ndarray, scale: float) -> np.ndarray:\n    array = torch.Tensor(array).unsqueeze(0).unsqueeze(0)\n    array = F.interpolate(array, scale_factor=scale, mode=\"trilinear\")\n    array = array.squeeze().numpy()\n    return array\n\n\ndef process(\n    tomo_dir: str,\n    dataset_dir: str,\n    scale: float,\n    radius_multi: float,\n    df_label: pd.DataFrame,\n):\n    tomo_id = tomo_dir.split(\"/\")[-1]\n    df_label_ = df_label[df_label[\"tomo_id\"] == tomo_id]\n\n    tomo_save_path = f\"{dataset_dir}/{tomo_id}/tomo.npy\"\n    label_save_path = f\"{dataset_dir}/{tomo_id}/label.npy\"\n\n    os.makedirs(f\"{dataset_dir}/{tomo_id}\", exist_ok=True)\n\n    slice_paths = sorted(glob(f\"{tomo_dir}/*.jpg\"))\n    tomo = stack_slices(slice_paths)\n    tomo_shape = tomo.shape\n    tomo = scale_array(tomo, scale)\n    np.save(tomo_save_path, tomo)\n\n    coords = df_label_[[\"Motor axis 0\", \"Motor axis 1\", \"Motor axis 2\"]].values\n    spacing = df_label_[\"Voxel spacing\"].values[0]\n    radius = int(TAU / spacing * radius_multi)\n    label = create_label_array(tomo_shape, coords, radius)\n    label = scale_array(label, scale).astype(np.bool_)\n    np.save(label_save_path, label)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-06T10:02:52.717752Z","iopub.execute_input":"2025-03-06T10:02:52.718289Z","iopub.status.idle":"2025-03-06T10:02:52.734014Z","shell.execute_reply.started":"2025-03-06T10:02:52.718256Z","shell.execute_reply":"2025-03-06T10:02:52.731957Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"dataset_dir = \"/kaggle/working/dataset/train\"\ntrain_tomo_dirs = sorted(glob(f\"{TRAIN_DIR}/*\"))\ndf_label = pd.read_csv(LABEL)\nshutil.rmtree(dataset_dir, ignore_errors=True)\n\ntrain_tomo_dirs = train_tomo_dirs[:5]\nscale = 0.5\nradius_multi = 0.2\n\n# n_workers = 2\n# with ProcessPoolExecutor(max_workers=n_workers) as executor:\n#     process_func = partial(\n#         process_tomo,\n#         dataset_dir=dataset_dir,\n#         df_label=df_label,\n#         scale=scale\n#         radius_multi=radius_multi,\n#     )\n#     results = list(\n#         tqdm(executor.map(process_func, train_tomo_dirs), total=len(train_tomo_dirs))\n#     )\n\nfor tomo_dir in tqdm(train_tomo_dirs):\n    process(tomo_dir, dataset_dir, scale, radius_multi, df_label)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-06T10:02:54.145899Z","iopub.execute_input":"2025-03-06T10:02:54.146262Z","iopub.status.idle":"2025-03-06T10:04:42.409311Z","shell.execute_reply.started":"2025-03-06T10:02:54.146229Z","shell.execute_reply":"2025-03-06T10:04:42.408035Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import matplotlib.pyplot as plt\n\ntomo_dirs = sorted(glob(f\"{dataset_dir}/*\"))\nfor tomo_dir in tomo_dirs:\n    tomo_id = tomo_dir.split(\"/\")[-1]\n\n    tomo_path = f\"{tomo_dir}/tomo.npy\"\n    label_path = f\"{tomo_dir}/label.npy\"\n\n    tomo = np.load(tomo_path)\n    label = np.load(label_path)\n\n    if label.sum() == 0:\n        continue\n\n    tomo_shape = tomo.shape\n    for i_z in range(50, tomo_shape[0] - 50, 2):\n        fig, ax = plt.subplots(1, 2, figsize=(10, 4))\n        ax[0].imshow(tomo[i_z])\n        ax[1].imshow(label[i_z])\n        fig.suptitle(f\"{tomo_id} z={i_z}\")\n        plt.show()\n\n    break\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-06T10:04:57.858847Z","iopub.execute_input":"2025-03-06T10:04:57.859215Z","iopub.status.idle":"2025-03-06T10:05:07.872886Z","shell.execute_reply.started":"2025-03-06T10:04:57.859184Z","shell.execute_reply":"2025-03-06T10:05:07.871567Z"}},"outputs":[],"execution_count":null}]}