{"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":30918,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# BYU Animation visualizer","metadata":{"execution":{"iopub.status.busy":"2025-03-06T11:37:04.456320Z","iopub.execute_input":"2025-03-06T11:37:04.456659Z","iopub.status.idle":"2025-03-06T11:37:06.656315Z","shell.execute_reply.started":"2025-03-06T11:37:04.456628Z","shell.execute_reply":"2025-03-06T11:37:06.655126Z"}}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nimport cv2\nimport os\nfrom matplotlib import animation, rc; rc('animation', html='jshtml')\nimport matplotlib.pyplot as plt\nimport matplotlib\nimport re\nimport gc\nfrom glob import glob","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-06T19:01:26.155898Z","iopub.execute_input":"2025-03-06T19:01:26.156281Z","iopub.status.idle":"2025-03-06T19:01:26.161312Z","shell.execute_reply.started":"2025-03-06T19:01:26.156241Z","shell.execute_reply":"2025-03-06T19:01:26.160302Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%matplotlib inline\nmatplotlib.rcParams['animation.embed_limit'] = 524288","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-06T19:01:26.162798Z","iopub.execute_input":"2025-03-06T19:01:26.163048Z","iopub.status.idle":"2025-03-06T19:01:26.181662Z","shell.execute_reply.started":"2025-03-06T19:01:26.163022Z","shell.execute_reply":"2025-03-06T19:01:26.180890Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"ROOT_DIR = \"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-06T19:01:26.183431Z","iopub.execute_input":"2025-03-06T19:01:26.183753Z","iopub.status.idle":"2025-03-06T19:01:26.197892Z","shell.execute_reply.started":"2025-03-06T19:01:26.183698Z","shell.execute_reply":"2025-03-06T19:01:26.196299Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df = pd.read_csv(f\"{ROOT_DIR}/train_labels.csv\")\ndf.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-06T19:01:26.199195Z","iopub.execute_input":"2025-03-06T19:01:26.199552Z","iopub.status.idle":"2025-03-06T19:01:26.233943Z","shell.execute_reply.started":"2025-03-06T19:01:26.199515Z","shell.execute_reply":"2025-03-06T19:01:26.232826Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df.describe()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-06T19:01:26.234963Z","iopub.execute_input":"2025-03-06T19:01:26.235323Z","iopub.status.idle":"2025-03-06T19:01:26.268654Z","shell.execute_reply.started":"2025-03-06T19:01:26.235289Z","shell.execute_reply":"2025-03-06T19:01:26.267651Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def create_animation(df, tomo_id, fps=20, imsize=384, frame_skip=1, save_as_gif=False, save_dir=\"./\"):\n    # Get paths\n    img_paths = sorted(glob(f\"{ROOT_DIR}/train/{tomo_id}/*\"))\n    annot_paths = sorted(glob(f\"{ROOT_DIR}/train/{tomo_id}/*\"))\n\n    pdf = df[df['tomo_id']==tomo_id]\n    orig_y = pdf[\"Array shape (axis 1)\"].values[0]\n    orig_x = pdf[\"Array shape (axis 2)\"].values[0]\n    ratio_x = imsize/orig_x\n    ratio_y = imsize/orig_y\n    ratio_m = min(ratio_x, ratio_y)\n\n    vox_sp = pdf[\"Voxel spacing\"].values[0]\n    \n    \n    # Get images\n    images = np.stack([cv2.imread(path) for path in img_paths],axis=0)\n    images = [np.array(cv2.resize(img, (imsize,imsize))) for img in images]\n    images = (images-np.min(images))/(np.max(images)-np.min(images)+1e-6)\n    images = (images*255).astype(np.uint8)\n\n    annot = np.stack([cv2.imread(path) for path in annot_paths],axis=0)\n    annot = [np.array(cv2.resize(img, (imsize,imsize))) for img in annot]\n    annot = (annot-np.min(annot))/(np.max(annot)-np.min(annot)+1e-6)\n    annot = (annot*255).astype(np.uint8)   \n\n    mat_imgs = np.zeros_like(annot)\n    \n    for k, row in pdf.iterrows():\n        m_ax_z = row[\"Motor axis 0\"]\n        m_ax_y = row[\"Motor axis 1\"]\n        m_ax_x = row[\"Motor axis 2\"]\n        tmp_mask = annot\n        for j, (i, a) in enumerate((zip(images, annot))):\n            base_rad = 1000 / vox_sp * ratio_m\n            rad = (base_rad**2 - abs(m_ax_z-j)**2)\n            rad = max(0, rad) **0.5\n            rad = int(np.round(rad))\n            cv2.circle(tmp_mask[j], (int(m_ax_x*ratio_x), int(m_ax_y*ratio_y)), rad, (255*(1&k), 255*(2&k), 255*(4&k)), thickness=-1)\n\n            \n            if j==m_ax_z:\n                cv2.putText(tmp_mask[j], 'GT_slice', (imsize-145, imsize-20), cv2.FONT_HERSHEY_COMPLEX, 1.0, (255*(1&k), 255*(2&k), 255*(4&k)), thickness=2)\n        for j in range(len(images)):\n            mat_imgs[j] = cv2.addWeighted(tmp_mask[j], 0.5, images[j], 0.5, 0)\n            cv2.putText(mat_imgs[j], f'{j:03d}/{len(images)-1:03d}', (5, imsize-20), cv2.FONT_HERSHEY_COMPLEX, 1.0, (255, 0, 0), thickness=2)\n\n    ims_sgs = [np.concatenate([images[i], mat_imgs[i]], axis=1) for i in range(len(images))]\n    \n    # Stack images\n    if frame_skip==0:\n        frame_skip = len(images)//100\n    animation_arr = np.stack(ims_sgs, axis=0)[::frame_skip]\n\n    del images, ims_sgs\n    gc.collect()\n    \n    # Initialise plot\n    fig = plt.figure(figsize=(6,3), dpi=128)  # if size is too big then gif gets truncated\n\n    im = plt.imshow(animation_arr[0], cmap='bone')\n    plt.axis('off')\n    plt.title(f\"{tomo_id}\", fontweight=\"bold\")\n    \n    # Load next frame\n    def animate_func(i):\n        im.set_array(animation_arr[i])\n        return [im]\n    plt.close()\n    \n    # Animation function\n    anim = animation.FuncAnimation(fig, animate_func, frames = animation_arr.shape[0], interval = 200//fps)\n   \n    \n    # Save\n    if save_as_gif:\n        os.makedirs(save_dir, exist_ok=True)\n        anim.save(os.path.join(save_dir, f\"patient_{tomo_id}.gif\"), fps=fps, writer='imagemagick')\n        \n    return anim","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-06T19:01:26.269655Z","iopub.execute_input":"2025-03-06T19:01:26.269960Z","iopub.status.idle":"2025-03-06T19:01:26.285042Z","shell.execute_reply.started":"2025-03-06T19:01:26.269937Z","shell.execute_reply":"2025-03-06T19:01:26.283870Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"example = 'tomo_00e047'\ncreate_animation(df, example, fps=5, imsize=256, frame_skip=0, save_as_gif=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-06T19:01:26.285934Z","iopub.execute_input":"2025-03-06T19:01:26.286270Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"example = 'tomo_00e463'\ncreate_animation(df, example, fps=5, imsize=256, frame_skip=0, save_as_gif=False)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df10 = df[df[\"Number of motors\"]==10]\nid10 = df10[\"tomo_id\"].iloc[0]\ndf10","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"example = id10\ncreate_animation(df, example, fps=5, imsize=256, frame_skip=0, save_as_gif=False)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 🚧 WIP 🚧","metadata":{}}]}