{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","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":31040,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Data Exploration and Visualization\n\nThis notebook serves as the foundational step for the Kaggle competition **“BYU - Locating Bacterial Flagellar Motors 2025”**. It focuses on in-depth exploration and visualization of the cryo-electron tomography (cryo-ET) dataset, provided as stacks of 2D JPEG slices forming 3D tomograms.\n\n### Objectives:\n- Understand the structure and diversity of tomograms\n- Analyze the spatial distribution of flagellar motors\n- Visualize data via slices, 3D scatter plots, and heatmaps\n- Assess voxel dimensions, pixel resolution, and data volume sizes\n- Identify potential preprocessing steps for downstream modeling\n\n### Key Insights:\n- 3D visualization of motor positions using scatter plots and heatmaps\n- Zoom-in views of motor areas to guide region proposals\n- Statistical summaries of tomogram dimensions and voxel spacing\n- Class balance between tomograms with and without motors\n- Interactive volume previews and axis-wise motor distribution","metadata":{}},{"cell_type":"code","source":"import os\nimport pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport imageio.v2 as imageio\nimport torch\nimport warnings\nfrom mpl_toolkits.mplot3d import Axes3D\nfrom IPython.display import display\nfrom scipy.ndimage import gaussian_filter\n\n\nINPUT_DIR = \"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025\"\nTRAIN_JPG = os.path.join(INPUT_DIR, \"train\")\nTRAIN_LABELS = os.path.join(INPUT_DIR, \"train_labels.csv\")\nTEST_JPG = os.path.join(INPUT_DIR, \"test\")\nSAMPLE_SUB = os.path.join(INPUT_DIR, \"sample_submission.csv\")\n\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# Load the CSV file containing motor labels\nlabels_df = pd.read_csv(TRAIN_LABELS)\ndisplay(labels_df.head())\n\n# Plot: Number of motors per tomogram\nmotor_counts = labels_df.groupby('tomo_id').size().reset_index(name='motor_count')\nplt.figure(figsize=(10, 4))\nsns.histplot(motor_counts['motor_count'], bins=20, kde=False)\nplt.title(\"Distribution of motors per tomogram\")\nplt.xlabel(\"Number of motors\")\nplt.ylabel(\"Number of tomograms\")\nplt.show()\n\n# Load a limited number of slices from a given tomogram as a 3D volume\ndef load_volume(tomo_id, max_slices=200):\n    folder = os.path.join(TRAIN_JPG, tomo_id)\n    slices = sorted(os.listdir(folder))[:max_slices]\n    volume = [imageio.imread(os.path.join(folder, s)) for s in slices]\n    return np.stack(volume)\n\n# Select a sample tomogram with at least one motor\ntomo_sample = labels_df[labels_df['Number of motors'] > 0]['tomo_id'].iloc[0]\nvolume = load_volume(tomo_sample)\nprint(f\"Loaded volume shape: {volume.shape}\")\n\n# Display the middle slice of the volume\nplt.imshow(volume[volume.shape[0] // 2], cmap='gray')\nplt.title(f\"Middle slice of {tomo_sample}\")\nplt.axis('off')\nplt.show()\n\n# Extract and plot motor coordinates in 3D\ncoords = labels_df[labels_df['tomo_id'] == tomo_sample][['Motor axis 0', 'Motor axis 1', 'Motor axis 2']].values\ncoords = np.array([c for c in coords if np.all(np.isfinite(c)) and np.all(np.array(c) >= 0)])\nfig = plt.figure(figsize=(8, 6))\nax = fig.add_subplot(111, projection='3d')\nax.scatter(coords[:,2], coords[:,1], coords[:,0], c='r', marker='o')\nax.set_title(f\"Motor positions in {tomo_sample}\")\nax.set_xlabel('X axis')\nax.set_ylabel('Y axis')\nax.set_zlabel('Z axis (slice index)')\nplt.show()\n\n# Generate 3D heatmap from motor positions\ndef create_heatmap(volume_shape, coordinates, sigma=3):\n    heatmap = np.zeros(volume_shape, dtype=np.float32)\n    for z, y, x in coordinates:\n        z, y, x = int(z), int(y), int(x)\n        if 0 <= z < volume_shape[0] and 0 <= y < volume_shape[1] and 0 <= x < volume_shape[2]:\n            heatmap[z, y, x] = 1.0\n    heatmap = gaussian_filter(heatmap, sigma=sigma)\n    return heatmap\n\nheatmap = create_heatmap(volume.shape, coords)\n\n# Visualize several slices of the heatmap around the strongest activation\nmax_z = np.argmax(np.max(np.max(heatmap, axis=1), axis=1))\nslice_range = range(max(0, max_z - 3), min(volume.shape[0], max_z + 3))\n\nfig, axs = plt.subplots(2, len(slice_range), figsize=(18, 6))\nfor i, z in enumerate(slice_range):\n    axs[0, i].imshow(volume[z], cmap='gray')\n    axs[0, i].set_title(f\"Volume Slice z={z}\")\n    axs[0, i].axis('off')\n\n    axs[1, i].imshow(heatmap[z], cmap='hot')\n    axs[1, i].set_title(f\"Heatmap Slice z={z}\")\n    axs[1, i].axis('off')\n\nplt.tight_layout()\nplt.show()\n\n# Summary statistics of the dataset\ndata_summary = labels_df.copy()\ndata_summary['has_motor'] = data_summary['Motor axis 0'] >= 0\nsummary = data_summary.groupby('tomo_id').agg(\n    number_of_motors=('has_motor', 'sum'),\n    z_slices=('Array shape (axis 0)', 'first'),\n    height=('Array shape (axis 1)', 'first'),\n    width=('Array shape (axis 2)', 'first'),\n    voxel_spacing=('Voxel spacing', 'first')\n).reset_index()\n\nprint(\"Summary of Tomograms:\")\ndisplay(summary.head())\n\n# Additional analysis: pixel size statistics\nsummary['pixel_area'] = summary['height'] * summary['width']\nmean_pixel_size = summary['pixel_area'].mean()\nmin_pixel_size = summary['pixel_area'].min()\nmax_pixel_size = summary['pixel_area'].max()\n\nmin_shape = summary.loc[summary['pixel_area'].idxmin(), ['height', 'width']].values\nmax_shape = summary.loc[summary['pixel_area'].idxmax(), ['height', 'width']].values\n\nprint(\"\\nPixel size analysis:\")\nprint(f\"Average pixel area: {mean_pixel_size:.2f} (in pixels)\")\nprint(f\"Smallest image: {min_shape[0]}x{min_shape[1]} (area={min_pixel_size})\")\nprint(f\"Largest image: {max_shape[0]}x{max_shape[1]} (area={max_pixel_size})\")\n\n# Zoom-in images around motor positions\nsample_tomos = labels_df[labels_df['Number of motors'] > 0]['tomo_id'].unique()[:6]\nfig, axs = plt.subplots(1, 6, figsize=(20, 4))\nfor i, tomo_id in enumerate(sample_tomos):\n    volume = load_volume(tomo_id)\n    motor_pos = labels_df[labels_df['tomo_id'] == tomo_id][['Motor axis 0', 'Motor axis 1', 'Motor axis 2']].values[0]\n    z, y, x = int(motor_pos[0]), int(motor_pos[1]), int(motor_pos[2])\n    z = np.clip(z, 0, volume.shape[0]-1)\n    y1, y2 = max(0, y-32), min(volume.shape[1], y+32)\n    x1, x2 = max(0, x-32), min(volume.shape[2], x+32)\n    zoom = volume[z, y1:y2, x1:x2]\n    axs[i].imshow(zoom, cmap='gray')\n    axs[i].set_title(f\"{tomo_id}\\nz={z}\")\n    axs[i].axis('off')\n\nplt.tight_layout()\nplt.show()\n\n\n# Balance analysis: tomograms with vs. without motors\nlabels_df['has_motor'] = labels_df['Motor axis 0'] >= 0\nbalance_df = labels_df.groupby('tomo_id')['has_motor'].any().value_counts().rename(index={True: 'Has Motor', False: 'No Motor'}).reset_index()\nbalance_df.columns = ['Class', 'Count']\ndisplay(balance_df)\n\nplt.figure(figsize=(6, 4))\nsns.barplot(x='Class', y='Count', data=balance_df)\nplt.title(\"Balance of Tomograms (with vs. without motors)\")\nplt.ylabel(\"Number of tomograms\")\nplt.xlabel(\"Class\")\nplt.show()\n\n# Voxel spacing distribution and implications\nplt.figure(figsize=(6, 4))\nsns.histplot(summary['voxel_spacing'], bins=30)\nplt.title(\"Voxel spacing distribution\")\nplt.xlabel(\"Angstrom per voxel\")\nplt.ylabel(\"Count\")\nplt.show()\n\nprint(\"Voxel spacing range:\")\nprint(f\"Min: {summary['voxel_spacing'].min()} Å, Max: {summary['voxel_spacing'].max()} Å\")\n\n# Volume size statistics (for memory planning)\nsummary['volume_size'] = summary['z_slices'] * summary['height'] * summary['width']\nplt.figure(figsize=(6, 4))\nsns.histplot(summary['volume_size'] / 1e6, bins=30)\nplt.title(\"Distribution of volume sizes\")\nplt.xlabel(\"Volume size (in millions of voxels)\")\nplt.ylabel(\"Number of tomograms\")\nplt.show()\n\nprint(\"Volume size stats (in voxels):\")\nprint(f\"Min: {summary['volume_size'].min()} | Max: {summary['volume_size'].max()} | Avg: {summary['volume_size'].mean():.2f}\")\nprint(\"Voxel spacing range:\")\nprint(f\"Min: {summary['voxel_spacing'].min()} Å, Max: {summary['voxel_spacing'].max()} Å\")\n\n# Volume size statistics (for memory planning)\nsummary['volume_size'] = summary['z_slices'] * summary['height'] * summary['width']\nplt.figure(figsize=(6, 4))\nsns.histplot(summary['volume_size'] / 1e6, bins=30)\nplt.title(\"Distribution of volume sizes\")\nplt.xlabel(\"Volume size (in millions of voxels)\")\nplt.ylabel(\"Number of tomograms\")\nplt.show()\n\nprint(\"Volume size stats (in voxels):\")\nprint(f\"Min: {summary['volume_size'].min()} | Max: {summary['volume_size'].max()} | Avg: {summary['volume_size'].mean():.2f}\")\n\n...\n\n# 3D Volume Preview using Plotly\nimport plotly.graph_objects as go\n\nfig = go.Figure(data=go.Volume(\n    x=np.repeat(np.arange(volume.shape[2]), volume.shape[0] * volume.shape[1]),\n    y=np.tile(np.repeat(np.arange(volume.shape[1]), volume.shape[2]), volume.shape[0]),\n    z=np.tile(np.arange(volume.shape[0]), volume.shape[1] * volume.shape[2]),\n    value=volume.flatten(),\n    opacity=0.1,\n    surface_count=15,\n    colorscale='Gray'\n))\nfig.update_layout(title='3D Volume Preview (Plotly)', scene=dict(zaxis_title='Z', yaxis_title='Y', xaxis_title='X'))\nfig.show()\n\n# Distribution of motors along Z axis\nmotor_z = labels_df[labels_df['Motor axis 0'] >= 0]['Motor axis 0']\nplt.figure(figsize=(8, 4))\nsns.histplot(motor_z, bins=50)\nplt.title(\"Distribution of motors along Z-axis\")\nplt.xlabel(\"Z (Slice index)\")\nplt.ylabel(\"Count\")\nplt.show()\n\n# Analysis: Number of motors vs tomogram size\nsummary['motors_per_million_voxels'] = summary['number_of_motors'] / (summary['volume_size'] / 1e6)\nplt.figure(figsize=(6, 4))\nsns.scatterplot(x='volume_size', y='number_of_motors', data=summary)\nplt.title(\"Number of motors vs. volume size\")\nplt.xlabel(\"Volume size (voxels)\")\nplt.ylabel(\"Number of motors\")\nplt.show()\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-05-23T10:52:19.042944Z","iopub.execute_input":"2025-05-23T10:52:19.043259Z","iopub.status.idle":"2025-05-23T10:53:57.791792Z","shell.execute_reply.started":"2025-05-23T10:52:19.043239Z","shell.execute_reply":"2025-05-23T10:53:57.790775Z"}},"outputs":[],"execution_count":null}]}