{"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,"sourceType":"competition"}],"dockerImageVersionId":30918,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"This notebook builds upon  [Andrew Darley's 'Parse Data' notebook](https://www.kaggle.com/code/andrewjdarley/parse-data).\nThe purpose is visualizing YZ and ZX slices.","metadata":{}},{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nimport shutil\nimport time\nimport yaml\nfrom pathlib import Path\nfrom tqdm.notebook import tqdm  # Use tqdm.notebook for Jupyter/Kaggle environments\n\nimport matplotlib.pyplot as plt\n\n# Set random seed for reproducibility\nnp.random.seed(42)\n\n# Define Kaggle paths\ndata_path = \"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/\"\ntrain_dir = os.path.join(data_path, \"train\")\n\n# # Define YOLO dataset structure\n# yolo_dataset_dir = \"/kaggle/working/yolo_dataset\"\n# yolo_images_train = os.path.join(yolo_dataset_dir, \"images\", \"train\")\n# yolo_images_val = os.path.join(yolo_dataset_dir, \"images\", \"val\")\n# yolo_labels_train = os.path.join(yolo_dataset_dir, \"labels\", \"train\")\n# yolo_labels_val = os.path.join(yolo_dataset_dir, \"labels\", \"val\")\n\n# # Create directories\n# for dir_path in [yolo_images_train, yolo_images_val, yolo_labels_train, yolo_labels_val]:\n#     os.makedirs(dir_path, exist_ok=True)\n\n# # Define constants\n# TRUST = 4  # Number of slices above and below center slice (total 2*TRUST + 1 slices)\n# BOX_SIZE = 24  # Bounding box size for annotations (in pixels)\n# TRAIN_SPLIT = 0.8  # 80% for training, 20% for validation","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Image processing functions\ndef normalize_slice(slice_data):\n    \"\"\"\n    Normalize slice data using 2nd and 98th percentiles\n    \"\"\"\n    # Calculate percentiles\n    p2 = np.percentile(slice_data, 2)\n    p98 = np.percentile(slice_data, 98)\n    \n    # Clip the data to the percentile range\n    clipped_data = np.clip(slice_data, p2, p98)\n    \n    # Normalize to [0, 255] range\n    normalized = 255 * (clipped_data - p2) / (p98 - p2)\n    \n    return np.uint8(normalized)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Load the label CSV\nlabels_df = pd.read_csv(os.path.join(data_path, \"train_labels.csv\"))\n\n# Extract tomograms that have motors\ntomo_df = labels_df[labels_df['Number of motors'] > 0].copy()\nunique_tomos = tomo_df['tomo_id'].unique()\nnum_unique_tomos = len(unique_tomos)\n\n# Loop through and visualize slices for each tomogram (from 0 to 361)\nfor num in range(num_unique_tomos):\n    src_path = os.path.join(train_dir, unique_tomos[num])\n    file_list = [f for f in os.listdir(src_path) if os.path.isfile(os.path.join(src_path, f))]\n    file_list = sorted(file_list, key=lambda x: int(x.split('_')[1].split('.')[0]))\n\n    # Get the motor labels for this tomogram\n    tomo_motors = labels_df[labels_df['tomo_id'] == unique_tomos[num]]\n    motor_counts = []\n    for _, motor in tomo_motors.iterrows():\n        if pd.isna(motor['Motor axis 0']):\n            continue\n        motor_counts.append(\n            (unique_tomos[num], \n             int(motor['Motor axis 0']), \n             int(motor['Motor axis 1']), \n             int(motor['Motor axis 2']),\n             int(motor['Array shape (axis 0)']),\n             int(motor['Voxel spacing']))\n        )\n\n    z_dim = len(file_list)  # Z-axis = number of image slices\n\n    # Get the size from the first image\n    sample_img = Image.open(os.path.join(src_path, file_list[0]))\n    y_dim, x_dim = np.array(sample_img).shape\n\n    # Initialize 3D array\n    stacked_images = np.zeros((z_dim, y_dim, x_dim), dtype=np.float32)\n\n    # Load and stack images one by one (no normalization)\n    for z, file in enumerate(file_list):\n        file_path = os.path.join(src_path, file)\n        img = Image.open(file_path)\n        img_array = np.array(img, dtype=np.float32)\n        stacked_images[z] = img_array\n\n    # ---- Convert slice image to RGB and draw red lines ----\n    def add_red_lines(image, horizontal=None, vertical=None):\n        \"\"\" Convert grayscale image to RGB and add red lines at specified positions \"\"\"\n        rgb_img = np.stack([image] * 3, axis=-1)  # Convert to RGB\n        rgb_img = (rgb_img / rgb_img.max() * 255).astype(np.uint8)  # Scale to 0–255\n\n        if horizontal is not None:\n            rgb_img[horizontal, :, :] = [255, 0, 0]  # Horizontal red line (Y-axis)\n        if vertical is not None:\n            rgb_img[:, vertical, :] = [255, 0, 0]  # Vertical red line (X-axis)\n\n        return rgb_img\n\n    # Plot slices based on each motor location\n    plt.figure(figsize=(15, len(motor_counts) * 5))\n\n    for idx, motor in enumerate(motor_counts):\n        middle_x = motor[3]\n        middle_y = motor[2]\n        middle_z = motor[1]\n\n        # Extract XY, YZ, ZX slices\n        xy_slice = normalize_slice(stacked_images[middle_z, :, :])  # XY slice (Z-axis)\n        yz_slice = normalize_slice(stacked_images[:, :, middle_x])  # YZ slice (X-axis)\n        zx_slice = normalize_slice(stacked_images[:, middle_y, :])  # ZX slice (Y-axis)\n\n        # Draw slice positions on each plane\n        xy_rgb = add_red_lines(xy_slice, horizontal=middle_y, vertical=middle_x)\n        yz_rgb = add_red_lines(yz_slice, horizontal=middle_z, vertical=middle_y)\n        zx_rgb = add_red_lines(zx_slice, horizontal=middle_z, vertical=middle_x)\n\n        # ---- Display slice images for each motor ----\n        plt.subplot(len(motor_counts), 3, idx * 3 + 1)\n        plt.title(f\"No.{num}, File {unique_tomos[num]}, (Voxel {motor[5]}), Motor {idx} - XY Slice\")\n        plt.imshow(xy_rgb)\n        plt.axis('off')\n\n        plt.subplot(len(motor_counts), 3, idx * 3 + 2)\n        plt.title(f\"YZ Slice\")\n        plt.imshow(yz_rgb)\n        plt.axis('off')\n\n        plt.subplot(len(motor_counts), 3, idx * 3 + 3)\n        plt.title(f\"ZX Slice\")\n        plt.imshow(zx_rgb)\n        plt.axis('off')\n\n    plt.tight_layout()\n    plt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}