{"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":"# BYU Locating Flagellar Motors\n\n## Data Preprocessing Pipeline\n\nThis is the first notebook in a series representing my submission for the BYU Locating Bacterial Flagellar Motors 2025 Kaggle challenge. The notebooks in this series cover the full workflow from data preprocessing to model training and submission generation.\n\n### Notebook Series:\n1. **Parse Data (Current)**: Extracting and preparing 2D slices containing motors to make a YOLO dataset\n2. **[Visualize Data](https://www.kaggle.com/code/andrewjdarley/visualize-data)**: Exploratory data analysis and visualization of annotated motor locations\n3. **[Train YOLO](https://www.kaggle.com/code/andrewjdarley/train-yolo)**: Fine tuning an YOLOv8 object detection model on the prepared dataset\n4. **[Submission Notebook](https://www.kaggle.com/code/andrewjdarley/submission-notebook)**: Running inference and generating submission files\n \n## About this Notebook\n\nThis notebook preprocesses the tomographic data for training. It performs the following:\n\n1. **Data Loading**: Reads the tomograms and motor location annotations\n2. **Slice Extraction**: Extracts 2D slices containing motors and surrounding slices (± TRUST parameter, little experimentation has been done modifying this parameter)\n3. **Image Normalization**: Normalizes slice intensity using percentile-based contrast enhancement (standard across all my work)\n4. **Dataset Preparation**: Organizes data into YOLO-compatible format with:\n   - Train/validation split of motors with no overlap in source tomograms (80/20 split)\n   - Bounding box annotations for each motor (I chose an arbitrary box size. It could be optimized)\n   - Proper directory structure for YOLO training\n5. **Configuration**: Generates dataset.yaml","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\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\nyolo_dataset_dir = \"/kaggle/working/yolo_dataset\"\nyolo_images_train = os.path.join(yolo_dataset_dir, \"images\", \"train\")\nyolo_images_val = os.path.join(yolo_dataset_dir, \"images\", \"val\")\nyolo_labels_train = os.path.join(yolo_dataset_dir, \"labels\", \"train\")\nyolo_labels_val = os.path.join(yolo_dataset_dir, \"labels\", \"val\")\n\n# Create directories\nfor 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\nTRUST = 24  # Number of slices above and below center slice (total 2*TRUST + 1 slices)\nBOX_SIZE = 24  # Bounding box size for annotations (in pixels)\nTRAIN_SPLIT = 0.8  # 80% for training, 20% for validation\n\n# Image processing functions\ndef normalize_slice(slice_data):\n    \"\"\"\n    Normalize slice data using 2nd and 98th percentiles\n    \"\"\"\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(255*(slice_data-0.5)/0.5)\n\ndef prepare_yolo_dataset(trust=TRUST, train_split=TRAIN_SPLIT):\n    \"\"\"\n    Extract slices containing motors from tomograms and save to YOLO structure with annotations\n    \"\"\"\n    # Load the labels CSV\n    labels_df = pd.read_csv(os.path.join(data_path, \"train_labels.csv\"))\n    \n    # Count total number of motors\n    total_motors = labels_df['Number of motors'].sum()\n    print(f\"Total number of motors in the dataset: {total_motors}\")\n    \n    # Get unique tomograms that have motors\n    tomo_df = labels_df[labels_df['Number of motors'] > 0].copy()\n    unique_tomos = tomo_df['tomo_id'].unique()\n    \n    print(f\"Found {len(unique_tomos)} unique tomograms with motors\")\n    \n    # Perform the train-val split at the tomogram level (not motor level)\n    # This ensures all slices from a single tomogram go to either train or val\n    np.random.shuffle(unique_tomos)  # Shuffle the tomograms\n    split_idx = int(len(unique_tomos) * train_split)\n    train_tomos = unique_tomos[:split_idx]\n    val_tomos = unique_tomos[split_idx:]\n    \n    print(f\"Split: {len(train_tomos)} tomograms for training, {len(val_tomos)} tomograms for validation\")\n    \n    # Function to process a set of tomograms\n    def process_tomogram_set(tomogram_ids, images_dir, labels_dir, set_name):\n        motor_counts = []\n        for tomo_id in tomogram_ids:\n            # Get all motors for this tomogram\n            tomo_motors = labels_df[labels_df['tomo_id'] == tomo_id]\n            for _, motor in tomo_motors.iterrows():\n                if pd.isna(motor['Motor axis 0']):\n                    continue\n                motor_counts.append(\n                    (tomo_id, \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        # print(f\"Will process approximately {len(motor_counts) * (2 * trust + 1)} slices for {set_name}\")\n        \n        # Process each motor\n        processed_slices = 0\n        \n        for tomo_id, z_center, y_center, x_center, z_max, voxelsz in tqdm(motor_counts, desc=f\"Processing {set_name} motors\"):\n            \n            BOX_SIZE = 400/voxelsz\n            # trust = int((BOX_SIZE +1)//2) # BOX_SIZE +16 slices along z (rounded up)\n            \n            # Calculate range of slices to include => always 20\n            z_center = np.round(z_center)\n            z_min = max(0, z_center - 5)\n            Z = min(z_max - 1, z_center + 5)\n            if Z - z_min <10:\n                if z_min==0:\n                    Z = 10 \n                elif Z==z_max-1:\n                    z_min = Z - 10\n            z_max = Z\n            if z_max - z_min != 10:\n                print(\"Weird : \", z_min, z_max)\n            trust = int(0.5*BOX_SIZE)  #\n            \n            # Process each slice in the range\n            for z in range(z_min, z_max + 1):\n                # Create slice filename\n                slice_filename = f\"slice_{z:04d}.jpg\"\n                \n                # Source path for the slice\n                src_path = os.path.join(train_dir, tomo_id, slice_filename)\n                \n                if not os.path.exists(src_path):\n                    print(f\"Warning: {src_path} does not exist, skipping.\")\n                    continue\n                \n                # Load and normalize the slice\n                img = Image.open(src_path)\n                img_array = np.array(img)\n                \n                # Normalize the image\n                normalized_img = img_array # normalize_slice(img_array)\n                shape = 640 ; shape2 = shape//2\n                x1, x2 = max(0, x_center - shape2), min(normalized_img.shape[1], x_center +shape2)\n                y1, y2 = max(0, y_center - shape2), min(normalized_img.shape[0], y_center +shape2)\n                # cropped_image = 0.5*np.ones_like(normalized_img[:shape,:shape]) \n                # cropped_image[:x2-x1,:y2-y1] = normalized_img[x1:x2, y1:y2] \n                normalized_img = normalized_img[y1:y2, x1:x2]\n                \n                # Create destination filename (with unique identifier)\n                dest_filename = f\"{tomo_id}_zc{z_center:04d}_yc{y_center:04d}_xc{x_center:04d}_{z}.jpg\"\n                dest_path = os.path.join(images_dir, dest_filename)\n                \n                # Save the normalized image\n                Image.fromarray(normalized_img).save(dest_path)\n\n                if z<z_center + trust and z>z_center - trust:\n                    # Get image dimensions\n                    img_width, img_height = x2-x1, y2-y1\n                    \n                    # Create YOLO format label\n                    # YOLO format: <class> <x_center> <y_center> <width> <height>\n                    # Values are normalized to [0, 1]\n                    x_center_norm = (x_center - x1) / img_width\n                    y_center_norm = (y_center - y1) / img_height\n                    box_width_norm = BOX_SIZE / img_width\n                    box_height_norm = BOX_SIZE / img_height \n                    # Write label file\n                    label_path = os.path.join(labels_dir, dest_filename.replace('.jpg', '.txt'))\n                    with open(label_path, 'w') as f:\n                        # dist = abs(z - z_center)\n                        f.write(f\"0 {x_center_norm} {y_center_norm} {box_width_norm} {box_height_norm}\\n\")\n                else:\n                    # Write label file\n                    label_path = os.path.join(labels_dir, dest_filename.replace('.jpg', '.txt'))\n                    with open(label_path, 'w') as f:\n                        f.write(f\"\\n\")\n                \n                processed_slices += 1\n                \n        return processed_slices, len(motor_counts)\n    \n    # Process training tomograms\n    train_slices, train_motors = process_tomogram_set(train_tomos, yolo_images_train, yolo_labels_train, \"training\")\n    \n    # Process validation tomograms\n    val_slices, val_motors = process_tomogram_set(val_tomos, yolo_images_val, yolo_labels_val, \"validation\")\n    \n    # Create YAML configuration file for YOLO\n    yaml_content = {\n        'path': yolo_dataset_dir,\n        'train': 'images/train',\n        'val': 'images/val',\n        'names': {0: 'motor'} # {0: 'near', 1: 'far'}\n    }\n    \n    with open(os.path.join(yolo_dataset_dir, 'dataset.yaml'), 'w') as f:\n        yaml.dump(yaml_content, f, default_flow_style=False)\n    \n    print(f\"\\nProcessing Summary:\")\n    print(f\"- Train set: {len(train_tomos)} tomograms, {train_motors} motors, {train_slices} slices\")\n    print(f\"- Validation set: {len(val_tomos)} tomograms, {val_motors} motors, {val_slices} slices\")\n    print(f\"- Total: {len(train_tomos) + len(val_tomos)} tomograms, {train_motors + val_motors} motors, {train_slices + val_slices} slices\")\n    \n    # Return summary info\n    return {\n        \"dataset_dir\": yolo_dataset_dir,\n        \"yaml_path\": os.path.join(yolo_dataset_dir, 'dataset.yaml'),\n        \"train_tomograms\": len(train_tomos),\n        \"val_tomograms\": len(val_tomos),\n        \"train_motors\": train_motors,\n        \"val_motors\": val_motors,\n        \"train_slices\": train_slices,\n        \"val_slices\": val_slices\n    }\n\n# Run the preprocessing\nsummary = prepare_yolo_dataset(TRUST)\nprint(f\"\\nPreprocessing Complete:\")\nprint(f\"- Training data: {summary['train_tomograms']} tomograms, {summary['train_motors']} motors, {summary['train_slices']} slices\")\nprint(f\"- Validation data: {summary['val_tomograms']} tomograms, {summary['val_motors']} motors, {summary['val_slices']} slices\")\nprint(f\"- Dataset directory: {summary['dataset_dir']}\")\nprint(f\"- YAML configuration: {summary['yaml_path']}\")\nprint(f\"\\nReady for YOLO training!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-04T15:56:22.254386Z","iopub.execute_input":"2025-04-04T15:56:22.254771Z","iopub.status.idle":"2025-04-04T15:58:24.565909Z","shell.execute_reply.started":"2025-04-04T15:56:22.254740Z","shell.execute_reply":"2025-04-04T15:58:24.564830Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import random\nimport matplotlib.pyplot as plt\nfrom PIL import Image, ImageDraw\nimport os\nimport numpy as np\nimport glob\n\n# Define base_dir - this was missing in the original code\n# In Kaggle, we can use the working directory as base or remove it completely\n# since we're using absolute paths\nbase_dir = \"/kaggle/working\"  # or simply use \"\" if using absolute paths\n\n# Updated paths without concatenating with base_dir since they're already absolute\nimages_train_dir = yolo_dataset_dir + \"/images/train/\"\nlabels_train_dir = yolo_dataset_dir + \"/labels/train/\"\n\n# Box size for highlighting the motor\nBOX_SIZE = 24\n\ndef visualize_random_training_samples(num_samples=4):\n    \"\"\"\n    Visualize random training samples with YOLO annotations\n    \n    Args:\n        num_samples (int): Number of random images to display\n    \"\"\"\n    # Get all image files from the train directory\n    image_files = []\n    for ext in ['*.jpg', '*.jpeg', '*.png']:\n        image_files.extend(glob.glob(os.path.join(images_train_dir, \"**\", ext), recursive=True))\n    \n    # Make sure we have enough images\n    if len(image_files) == 0:\n        print(\"No image files found in the train directory!\")\n        return\n        \n    num_samples = min(num_samples, len(image_files))\n    \n    # Select random images\n    random_images = random.sample(image_files, num_samples)\n    \n    # Create a figure with subplots\n    rows = int(np.ceil(num_samples / 5))\n    cols = 5\n    fig, axes = plt.subplots(rows, cols, figsize=(14, 3 * rows))\n    \n    # Handle the case of a single subplot\n    if num_samples == 1:\n        axes = np.array([axes])\n    \n    # Flatten axes array for easy indexing\n    axes = axes.flatten()\n    \n    # Process each selected image\n    for i, img_path in enumerate(random_images):\n        try:\n            # Get corresponding label file\n            # YOLO labels have same name but .txt extension instead of image extension\n            relative_path = os.path.relpath(img_path, images_train_dir)\n            label_path = os.path.join(labels_train_dir, os.path.splitext(relative_path)[0] + '.txt')\n            \n            # Load the image\n            img = Image.open(img_path)\n            img_width, img_height = img.size\n            \n            # Normalize image using percentiles for better visualization\n            img_array = np.array(img)\n            p2 = np.percentile(img_array, 2)\n            p98 = np.percentile(img_array, 98)\n            normalized = np.clip(img_array, p2, p98)\n            normalized = 255 * (normalized - p2) / (p98 - p2)\n            img_normalized = Image.fromarray(np.uint8(normalized))\n            \n            # Convert image to RGB for colored box\n            img_rgb = img_normalized.convert('RGB')\n            \n            # Create a transparent overlay\n            overlay = Image.new('RGBA', img_rgb.size, (0, 0, 0, 0))\n            draw = ImageDraw.Draw(overlay)\n            \n            # Load YOLO format annotations if they exist\n            annotations = []\n            if os.path.exists(label_path):\n                with open(label_path, 'r') as f:\n                    for line in f:\n                        # YOLO format: class x_center y_center width height\n                        # All values are normalized from 0 to 1\n                        values = line.strip().split()\n                        class_id = int(values[0])\n                        x_center = float(values[1]) * img_width\n                        y_center = float(values[2]) * img_height\n                        width = float(values[3]) * img_width\n                        height = float(values[4]) * img_height\n                        \n                        annotations.append({\n                            'class_id': class_id,\n                            'x_center': x_center,\n                            'y_center': y_center,\n                            'width': width,\n                            'height': height\n                        })\n            \n            # Draw all annotations\n            for ann in annotations:\n                x_center = ann['x_center']\n                y_center = ann['y_center']\n                width = ann['width']\n                height = ann['height']\n                \n                # Calculate bounding box coordinates\n                x1 = max(0, int(x_center - width/2))\n                y1 = max(0, int(y_center - height/2))\n                x2 = min(img_width, int(x_center + width/2))\n                y2 = min(img_height, int(y_center + height/2))\n                \n                # Draw semi-transparent red rectangle\n                draw.rectangle([x1, y1, x2, y2], fill=(255, 0, 0, 64), outline=(255, 0, 0, 200))\n                \n                # Draw label\n                label_text = f\"Class {ann['class_id']}\"\n                draw.text((x1, y1-10), label_text, fill=(255, 0, 0, 255))\n            \n            # If no annotations found, indicate this\n            if not annotations:\n                draw.text((10, 10), \"No annotations found\", fill=(255, 0, 0, 255))\n            \n            # Composite the overlay onto the original image\n            img_rgb = Image.alpha_composite(img_rgb.convert('RGBA'), overlay).convert('RGB')\n            \n            # Display the image with annotations\n            axes[i].imshow(np.array(img_rgb))\n            img_name = os.path.basename(img_path)\n            axes[i].set_title(f\"Image: {img_name}\\nAnnotations: {len(annotations)}\")\n            axes[i].axis('on')\n            \n        except Exception as e:\n            print(f\"Error processing image {img_path}: {e}\")\n            axes[i].text(0.5, 0.5, f\"Error loading image: {os.path.basename(img_path)}\", \n                       horizontalalignment='center', verticalalignment='center')\n            axes[i].axis('off')\n    \n    # Handle extra subplots if any\n    for j in range(i + 1, len(axes)):\n        axes[j].axis('off')\n    \n    plt.tight_layout()\n    plt.show()\n    \n    # Print summary\n    print(f\"Displayed {num_samples} random images with YOLO annotations\")\n\n# Run the visualization\nvisualize_random_training_samples(15)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-04T15:56:12.309809Z","iopub.execute_input":"2025-04-04T15:56:12.310205Z","iopub.status.idle":"2025-04-04T15:56:16.078139Z","shell.execute_reply.started":"2025-04-04T15:56:12.310170Z","shell.execute_reply":"2025-04-04T15:56:16.076486Z"}},"outputs":[],"execution_count":null}]}