{"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":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nimport shutil\nimport time\nimport yaml\nimport cv2\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 = 2  # 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\nIMG_SIZE = 1024\n# 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)\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                )\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 in tqdm(motor_counts, desc=f\"Processing {set_name} motors\"):\n            # Calculate range of slices to include\n            z_min = max(0, z_center - trust)\n            z_max = min(z_max - 1, z_center + trust)\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                orig_width, orig_height = img.size\n                img_array = np.array(img)\n                scale = min(IMG_SIZE/orig_width, IMG_SIZE/orig_height)\n                new_width = int(orig_width * scale)\n                new_height = int(orig_height * scale)\n                resized_img = cv2.resize(img_array, (new_width, new_height), interpolation=cv2.INTER_AREA)\n                canvas = np.zeros((IMG_SIZE, IMG_SIZE), dtype=np.uint8)\n                # Normalize the image\n                x_offset = (IMG_SIZE - new_width) // 2\n                y_offset = (IMG_SIZE - new_height) // 2\n                canvas[y_offset:y_offset+new_height, x_offset:x_offset+new_width] = resized_img\n                normalized_img = normalize_slice(canvas)\n                \n                # Create destination filename (with unique identifier)\n                dest_filename = f\"{tomo_id}_z{z:04d}_y{y_center:04d}_x{x_center:04d}.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                scaled_x = x_center * scale\n                scaled_y = y_center * scale\n        \n                # Then add the padding offset\n                new_x_center = scaled_x + x_offset\n                new_y_center = scaled_y + y_offset\n                # Scale the box size\n                scaled_box_size = BOX_SIZE * scale\n                # Normalize to [0,1] relative to IMG_SIZE\n                x_center_norm = new_x_center / IMG_SIZE\n                y_center_norm = new_y_center / IMG_SIZE\n                box_width_norm = scaled_box_size / IMG_SIZE\n                box_height_norm = scaled_box_size / IMG_SIZE\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\"0 {x_center_norm} {y_center_norm} {box_width_norm} {box_height_norm}\\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'}\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-14T02:40:25.603293Z","iopub.execute_input":"2025-04-14T02:40:25.603719Z","iopub.status.idle":"2025-04-14T02:40:25.867009Z","shell.execute_reply.started":"2025-04-14T02:40:25.603664Z","shell.execute_reply":"2025-04-14T02:40:25.865237Z"}},"outputs":[],"execution_count":null}]}