{"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":"gpu","dataSources":[{"sourceType":"competition","sourceId":91249,"databundleVersionId":11294684},{"sourceType":"kernelVersion","sourceId":211097053},{"sourceType":"kernelVersion","sourceId":226864880},{"sourceType":"kernelVersion","sourceId":227202990},{"sourceType":"kernelVersion","sourceId":231023075}],"dockerImageVersionId":30918,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Post processing and training Data for BYU\n\nMay thanks to:\n\nBase line notebook from [andrewjdarley](https://www.kaggle.com/code/andrewjdarley/parse-data)\n\nImprovement notebook from [sharifi76](https://www.kaggle.com/code/sharifi76/eda-visualization-yolov8)","metadata":{}},{"cell_type":"code","source":"!tar xfvz /kaggle/input/k/mathieuduverne/ultralytics-for-offline-install/archive.tar.gz\n!pip install --no-index --find-links=./packages ultralytics\n!rm -rf ./packages","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T14:38:48.379217Z","iopub.execute_input":"2025-04-01T14:38:48.379601Z","iopub.status.idle":"2025-04-01T14:39:39.497232Z","shell.execute_reply.started":"2025-04-01T14:38:48.379572Z","shell.execute_reply":"2025-04-01T14:39:39.496305Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import plotly.express as px\nfrom PIL import Image, ImageDraw\nimport random\nimport seaborn as sns\nfrom matplotlib.patches import Rectangle\nfrom ultralytics import YOLO\nimport yaml\nimport json\nimport os\nimport glob\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom tqdm.notebook import tqdm\nfrom sklearn.model_selection import train_test_split\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\nimport cv2\nimport threading\nimport time\nfrom contextlib import nullcontext\nfrom concurrent.futures import ThreadPoolExecutor\nimport math","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-04-01T14:39:39.498494Z","iopub.execute_input":"2025-04-01T14:39:39.498891Z","iopub.status.idle":"2025-04-01T14:39:48.713790Z","shell.execute_reply.started":"2025-04-01T14:39:39.498854Z","shell.execute_reply":"2025-04-01T14:39:48.713112Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Define global constants for dataset directories\nDATA_DIR = '/kaggle/input/byu-locating-bacterial-flagellar-motors-2025'\nTRAIN_CSV = os.path.join(DATA_DIR, 'train_labels.csv')\nTRAIN_DIR = os.path.join(DATA_DIR, 'train')\nTEST_DIR = os.path.join(DATA_DIR, 'test')\nOUTPUT_DIR = './'\nMODEL_DIR = './models'\n\n# Create output directories if they don't exist\nos.makedirs(OUTPUT_DIR, exist_ok=True)\nos.makedirs(MODEL_DIR, exist_ok=True)\n\n# Set device: Use GPU if available; otherwise, fall back to CPU\nDEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Using device: {DEVICE}\")\n\n# Set random seeds for reproducibility\nRANDOM_SEED = 42\nrandom.seed(RANDOM_SEED)\nnp.random.seed(RANDOM_SEED)\ntorch.manual_seed(RANDOM_SEED)\nif torch.cuda.is_available():\n    torch.cuda.manual_seed(RANDOM_SEED)\n    torch.backends.cudnn.deterministic = True\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T14:39:48.715141Z","iopub.execute_input":"2025-04-01T14:39:48.715679Z","iopub.status.idle":"2025-04-01T14:39:48.775828Z","shell.execute_reply.started":"2025-04-01T14:39:48.715654Z","shell.execute_reply":"2025-04-01T14:39:48.775014Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 3) **YOLO**\n\n## 3.1) Data Preprocessing for YOLO Dataset Preparation\n\nIn this section, we:\n\n- **Load Motor Annotations:** Read the tomogram and motor location annotations.\n- **Slice Extraction:** For each motor, extract 2D slices (with a configurable range above and below the motor slice).\n- **Image Normalization:** Normalize each slice using percentile-based contrast enhancement.\n- **Dataset Organization:** Save the normalized images and generate YOLO-format bounding box annotations.\n- **Train/Validation Split:** Split the data by tomogram to avoid any overlap between training and validation sets.\n- **Configuration File Generation:** Create a `dataset.yaml` file for YOLO training.","metadata":{}},{"cell_type":"code","source":"# Define YOLO dataset structure and parameters\ndata_path = \"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/\"\ntrain_dir = os.path.join(data_path, \"train\")\n\n# Output directories for YOLO dataset (adjust as needed)\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 necessary 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 for processing\nTRUST = 4       # Number of slices above and below center slice (total slices = 2*TRUST + 1)\nBOX_SIZE = 24   # Bounding box size (in pixels)\nTRAIN_SPLIT = 0.8  # 80% training, 20% validation\n\n# Define a helper function for image normalization using percentile-based contrast enhancement.\ndef normalize_slice(slice_data):\n    \"\"\"\n    Normalize slice data using the 2nd and 98th percentiles.\n    \n    Args:\n        slice_data (numpy.array): Input image slice.\n    \n    Returns:\n        np.uint8: Normalized image in the range [0, 255].\n    \"\"\"\n    p2 = np.percentile(slice_data, 2)\n    p98 = np.percentile(slice_data, 98)\n    clipped_data = np.clip(slice_data, p2, p98)\n    normalized = 255 * (clipped_data - p2) / (p98 - p2)\n    return np.uint8(normalized)\n\n# Define the preprocessing function to extract slices, normalize, and generate YOLO annotations.\ndef prepare_yolo_dataset(trust=TRUST, train_split=TRAIN_SPLIT):\n    \"\"\"\n    Extract slices containing motors and save images with corresponding YOLO annotations.\n    \n    Steps:\n    - Load the motor labels.\n    - Perform a train/validation split by tomogram.\n    - For each motor, extract slices in a range (± trust parameter).\n    - Normalize each slice and save it.\n    - Generate YOLO format bounding box annotations with a fixed box size.\n    - Create a YAML configuration file for YOLO training.\n    \n    Returns:\n        dict: A summary containing dataset statistics and file paths.\n    \"\"\"\n    # Load the labels CSV\n    labels_df = pd.read_csv(os.path.join(data_path, \"train_labels.csv\"))\n    \n    total_motors = labels_df['Number of motors'].sum()\n    print(f\"Total number of motors in the dataset: {total_motors}\")\n    \n    # Consider only tomograms with at least one motor\n    tomo_df = labels_df[labels_df['Number of motors'] > 0].copy()\n    unique_tomos = tomo_df['tomo_id'].unique()\n    print(f\"Found {len(unique_tomos)} unique tomograms with motors\")\n    \n    # Shuffle and split tomograms into train and validation sets\n    np.random.shuffle(unique_tomos)\n    split_idx = int(len(unique_tomos) * train_split)\n    train_tomos = unique_tomos[:split_idx]\n    val_tomos = unique_tomos[split_idx:]\n    print(f\"Split: {len(train_tomos)} tomograms for training, {len(val_tomos)} tomograms for validation\")\n    \n    # Helper function to process a list 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 motor annotations for the current 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        processed_slices = 0\n        \n        # Loop over each motor annotation\n        for tomo_id, z_center, y_center, x_center, z_max in tqdm(motor_counts, desc=f\"Processing {set_name} motors\"):\n            z_min = max(0, z_center - trust)\n            z_max_bound = min(z_max - 1, z_center + trust)\n            for z in range(z_min, z_max_bound + 1):\n                # Create the slice filename and source path\n                slice_filename = f\"slice_{z:04d}.jpg\"\n                src_path = os.path.join(train_dir, tomo_id, slice_filename)\n                if not os.path.exists(src_path):\n                    print(f\"Warning: {src_path} does not exist, skipping.\")\n                    continue\n                \n                # Load, normalize, and save the image slice\n                img = Image.open(src_path)\n                img_array = np.array(img)\n                normalized_img = normalize_slice(img_array)\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                Image.fromarray(normalized_img).save(dest_path)\n                \n                # Prepare YOLO bounding box annotation (normalized values)\n                img_width, img_height = img.size\n                x_center_norm = x_center / img_width\n                y_center_norm = y_center / img_height\n                box_width_norm = BOX_SIZE / img_width\n                box_height_norm = BOX_SIZE / img_height\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                \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    # Process validation tomograms\n    val_slices, val_motors = process_tomogram_set(val_tomos, yolo_images_val, yolo_labels_val, \"validation\")\n    \n    # Generate YAML configuration for YOLO training\n    yaml_content = {\n        'path': yolo_dataset_dir,\n        'train': 'images/train',\n        'val': 'images/val',\n        'names': {0: 'motor'}\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 {\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(\"\\nReady for YOLO training!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T14:39:48.777159Z","iopub.execute_input":"2025-04-01T14:39:48.777410Z","iopub.status.idle":"2025-04-01T14:42:49.464475Z","shell.execute_reply.started":"2025-04-01T14:39:48.777388Z","shell.execute_reply":"2025-04-01T14:42:49.463693Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Set random seeds for reproducibility\nnp.random.seed(42)\nrandom.seed(42)\ntorch.manual_seed(42)\n\n# Define paths for the Kaggle environment\nyolo_dataset_dir = \"/kaggle/working/yolo_dataset\"\nyolo_weights_dir = \"/kaggle/working/yolo_weights\"\nyolo_pretrained_weights = \"/kaggle/input/k/mathieuduverne/ultralytics-for-offline-install/yolov8n.pt\"  # Pre-downloaded weights\n\n# Create the weights directory if it does not exist\nos.makedirs(yolo_weights_dir, exist_ok=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T14:42:49.465296Z","iopub.execute_input":"2025-04-01T14:42:49.465584Z","iopub.status.idle":"2025-04-01T14:42:49.470918Z","shell.execute_reply.started":"2025-04-01T14:42:49.465548Z","shell.execute_reply":"2025-04-01T14:42:49.470272Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def fix_yaml_paths(yaml_path):\n    \"\"\"\n    Fix the paths in the YAML file to match the actual Kaggle directories.\n    \n    Args:\n        yaml_path (str): Path to the original dataset YAML file.\n        \n    Returns:\n        str: Path to the fixed YAML file.\n    \"\"\"\n    print(f\"Fixing YAML paths in {yaml_path}\")\n    with open(yaml_path, 'r') as f:\n        yaml_data = yaml.safe_load(f)\n    \n    if 'path' in yaml_data:\n        yaml_data['path'] = yolo_dataset_dir\n    \n    fixed_yaml_path = \"/kaggle/working/fixed_dataset.yaml\"\n    with open(fixed_yaml_path, 'w') as f:\n        yaml.dump(yaml_data, f)\n    \n    print(f\"Created fixed YAML at {fixed_yaml_path} with path: {yaml_data.get('path')}\")\n    return fixed_yaml_path","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T14:42:49.471741Z","iopub.execute_input":"2025-04-01T14:42:49.472052Z","iopub.status.idle":"2025-04-01T14:42:49.486792Z","shell.execute_reply.started":"2025-04-01T14:42:49.472023Z","shell.execute_reply":"2025-04-01T14:42:49.485962Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def plot_dfl_loss_curve(run_dir):\n    \"\"\"\n    Plot the DFL loss curves for training and validation, marking the best model.\n    \n    Args:\n        run_dir (str): Directory where the training results are stored.\n    \"\"\"\n    results_csv = os.path.join(run_dir, 'results.csv')\n    if not os.path.exists(results_csv):\n        print(f\"Results file not found at {results_csv}\")\n        return\n    \n    results_df = pd.read_csv(results_csv)\n    train_dfl_col = [col for col in results_df.columns if 'train/dfl_loss' in col]\n    val_dfl_col = [col for col in results_df.columns if 'val/dfl_loss' in col]\n    \n    if not train_dfl_col or not val_dfl_col:\n        print(\"DFL loss columns not found in results CSV\")\n        print(f\"Available columns: {results_df.columns.tolist()}\")\n        return\n    \n    train_dfl_col = train_dfl_col[0]\n    val_dfl_col = val_dfl_col[0]\n    \n    best_epoch = results_df[val_dfl_col].idxmin()\n    best_val_loss = results_df.loc[best_epoch, val_dfl_col]\n    \n    plt.figure(figsize=(10, 6))\n    plt.plot(results_df['epoch'], results_df[train_dfl_col], label='Train DFL Loss')\n    plt.plot(results_df['epoch'], results_df[val_dfl_col], label='Validation DFL Loss')\n    plt.axvline(x=results_df.loc[best_epoch, 'epoch'], color='r', linestyle='--', \n                label=f'Best Model (Epoch {int(results_df.loc[best_epoch, \"epoch\"])}, Val Loss: {best_val_loss:.4f})')\n    plt.xlabel('Epoch')\n    plt.ylabel('DFL Loss')\n    plt.title('Training and Validation DFL Loss')\n    plt.legend()\n    plt.grid(True, linestyle='--', alpha=0.7)\n    \n    plot_path = os.path.join(run_dir, 'dfl_loss_curve.png')\n    plt.savefig(plot_path)\n    plt.savefig(os.path.join('/kaggle/working', 'dfl_loss_curve.png'))\n    \n    print(f\"Loss curve saved to {plot_path}\")\n    plt.close()\n    \n    return best_epoch, best_val_loss","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T14:42:49.487618Z","iopub.execute_input":"2025-04-01T14:42:49.487904Z","iopub.status.idle":"2025-04-01T14:42:49.504017Z","shell.execute_reply.started":"2025-04-01T14:42:49.487873Z","shell.execute_reply":"2025-04-01T14:42:49.503405Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 3.3) YOLOv8 Training Pipeline for Flagellar Motor Detection\n\nThis notebook implements a full YOLOv8 training pipeline for detecting bacterial flagellar motors in tomographic slices.\n\n**Overview:**\n- **Dataset Configuration:** Sets up and validates the YOLO-format dataset YAML file.\n- **Model Initialization:** Loads pre-trained YOLOv8 weights for transfer learning.\n- **Training Process:** Fine tunes the model with early stopping and periodic checkpoints.\n- **Loss Visualization:** Plots training and validation DFL loss curves to monitor progress.\n- **Performance Evaluation:** Tests the trained model on random validation samples.\n- **Model Export:** Saves the trained weights for later use.\n\nLet's begin by importing the necessary libraries and setting up reproducibility.\n","metadata":{}},{"cell_type":"code","source":"def train_yolo_model(yaml_path, pretrained_weights_path, epochs=50, img_size=640):\n    \"\"\"\n    Train a YOLO model with stabilized hyperparameters to reduce loss fluctuations.\n\n    Args:\n        yaml_path (str): Path to the dataset YAML file.\n        pretrained_weights_path (str): Path to pre-downloaded weights file.\n        epochs (int): Number of training epochs.\n        img_size (int): Image size for training.\n\n    Returns:\n        model (YOLO): Trained YOLO model.\n        results: Training results.\n    \"\"\"\n    print(f\"Loading pre-trained weights from: {pretrained_weights_path}\")\n    model = YOLO(pretrained_weights_path)\n\n    results = model.train(\n        data=yaml_path,\n        epochs=epochs,\n        imgsz=img_size,\n        project=yolo_weights_dir,\n        name='motor_detector_stable',\n        exist_ok=True,\n        \n        # Batch and hardware settings\n        batch=16,                # Keep original batch size or increase if memory allows\n        workers=4,               # Data loading workers\n        device=0,                # Use GPU 0\n        amp=True,                # Mixed precision for efficiency\n        \n        # Learning rate and optimizer settings\n        optimizer=\"AdamW\",       # Using AdamW optimizer\n        lr0=0.0005,              # Lower initial learning rate for stability\n        lrf=0.05,                # Final learning rate factor (higher than before)\n        momentum=0.9,            # Slightly lower momentum\n        weight_decay=0.001,      # Increased weight decay for regularization\n        warmup_epochs=3,         # Add warmup period\n        warmup_momentum=0.8,     # Warmup momentum\n        warmup_bias_lr=0.1,      # Warmup bias learning rate\n        \n        # Loss function weights\n        box=7.5,                 # Higher box loss weight\n        cls=0.5,                 # Lower classification loss weight \n        dfl=1.5,                 # Adjusted DFL loss weight\n        \n        # Augmentation parameters  \n        mosaic=0.7,              # Reduce mosaic probability\n        mixup=0.15,              # Slightly reduced mixup\n        copy_paste=0.1,          # Add copy-paste augmentation\n        degrees=0.0,             # No rotation (specific to this task)\n        translate=0.1,           # Add translation augmentation\n        scale=0.5,               # Scale augmentation\n        shear=0.0,               # No shear (could destabilize training)\n        fliplr=0.5,              # Horizontal flip probability\n        flipud=0.0,              # No vertical flip (task specific)\n        hsv_h=0.015,             # Hue augmentation\n        hsv_s=0.5,               # Saturation augmentation\n        hsv_v=0.3,               # Value/brightness augmentation\n        perspective=0.0,         # No perspective shift for this task\n        \n        # Stability and convergence parameters\n        patience=15,             # Increase patience from 10 to 15\n        save_period=5,           # Save model every 5 epochs\n        close_mosaic=15,         # Disable mosaic augmentation later in training\n        overlap_mask=True,       # Overlap mask for better boundary detection\n        \n        # Model specific settings\n        dropout=0.1,             # Add dropout for regularization\n        val=True,                # Keep validation on\n        rect=False,              # Don't use rectangular training\n        multi_scale=True,        # Add multi-scale training for robustness\n    )\n\n    run_dir = os.path.join(yolo_weights_dir, 'motor_detector_stable')\n    \n    # Plot loss curves\n    if 'plot_dfl_loss_curve' in globals():\n        best_epoch_info = plot_dfl_loss_curve(run_dir)\n        if best_epoch_info:\n            best_epoch, best_val_loss = best_epoch_info\n            print(f\"\\nBest model found at epoch {best_epoch} with validation DFL loss: {best_val_loss:.4f}\")\n\n    return model, results","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T14:44:22.085687Z","iopub.execute_input":"2025-04-01T14:44:22.086037Z","iopub.status.idle":"2025-04-01T14:44:22.093023Z","shell.execute_reply.started":"2025-04-01T14:44:22.086013Z","shell.execute_reply":"2025-04-01T14:44:22.092057Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def plot_learning_rate_schedule(epochs=50, lr0=0.0005, lrf=0.05, warmup_epochs=3):\n    \"\"\"\n    Visualize the learning rate schedule with warmup and cosine decay.\n    \n    Args:\n        epochs (int): Total number of epochs\n        lr0 (float): Initial learning rate\n        lrf (float): Final learning rate factor\n        warmup_epochs (int): Number of warmup epochs\n    \"\"\"\n    epochs_array = np.arange(epochs)\n    lr_array = np.zeros(epochs)\n    \n    # Warmup phase\n    for i in range(warmup_epochs):\n        lr_array[i] = lr0 * ((i + 1) / warmup_epochs)\n    \n    # Cosine decay phase\n    for i in range(warmup_epochs, epochs):\n        lr_array[i] = lr0 * (((1 - math.cos(math.pi * (i - warmup_epochs) / (epochs - warmup_epochs))) / 2) * (lrf - 1) + 1)\n    \n    plt.figure(figsize=(10, 6))\n    plt.plot(epochs_array, lr_array)\n    plt.xlabel('Epoch')\n    plt.ylabel('Learning Rate')\n    plt.title('Learning Rate Schedule with Warmup and Cosine Decay')\n    plt.grid(True, linestyle='--', alpha=0.7)\n    plt.savefig(os.path.join('/kaggle/working', 'lr_schedule.png'))\n    plt.show()\n    \n    return lr_array","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T14:42:49.524535Z","iopub.execute_input":"2025-04-01T14:42:49.524736Z","iopub.status.idle":"2025-04-01T14:42:49.539905Z","shell.execute_reply.started":"2025-04-01T14:42:49.524718Z","shell.execute_reply":"2025-04-01T14:42:49.539222Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def analyze_training_stability(results_csv):\n    \"\"\"\n    Analyze and visualize training stability from results CSV.\n    \n    Args:\n        results_csv (str): Path to the results CSV file.\n    \"\"\"\n    results_df = pd.read_csv(results_csv)\n    \n    # Set up figure\n    fig, axs = plt.subplots(2, 2, figsize=(16, 12))\n    \n    # Plot 1: All losses\n    axs[0, 0].plot(results_df['epoch'], results_df['train/box_loss'], label='Train Box Loss')\n    axs[0, 0].plot(results_df['epoch'], results_df['val/box_loss'], label='Val Box Loss')\n    axs[0, 0].plot(results_df['epoch'], results_df['train/cls_loss'], label='Train Cls Loss')\n    axs[0, 0].plot(results_df['epoch'], results_df['val/cls_loss'], label='Val Cls Loss')\n    axs[0, 0].plot(results_df['epoch'], results_df['train/dfl_loss'], label='Train DFL Loss')\n    axs[0, 0].plot(results_df['epoch'], results_df['val/dfl_loss'], label='Val DFL Loss')\n    axs[0, 0].set_title('All Training and Validation Losses')\n    axs[0, 0].set_xlabel('Epoch')\n    axs[0, 0].set_ylabel('Loss')\n    axs[0, 0].legend()\n    axs[0, 0].grid(True, linestyle='--', alpha=0.7)\n    \n    # Plot 2: Performance metrics\n    axs[0, 1].plot(results_df['epoch'], results_df['metrics/precision(B)'], label='Precision')\n    axs[0, 1].plot(results_df['epoch'], results_df['metrics/recall(B)'], label='Recall')\n    axs[0, 1].plot(results_df['epoch'], results_df['metrics/mAP50(B)'], label='mAP50')\n    axs[0, 1].plot(results_df['epoch'], results_df['metrics/mAP50-95(B)'], label='mAP50-95')\n    axs[0, 1].set_title('Model Performance Metrics')\n    axs[0, 1].set_xlabel('Epoch')\n    axs[0, 1].set_ylabel('Metric Value')\n    axs[0, 1].legend()\n    axs[0, 1].grid(True, linestyle='--', alpha=0.7)\n    \n    # Plot 3: Box and Classification Loss Stability\n    # Calculate rolling mean and standard deviation to assess stability\n    window = 3  # 3-epoch window for smoothing\n    axs[1, 0].plot(results_df['epoch'], results_df['train/box_loss'].rolling(window=window, min_periods=1).mean(), \n                  label='Train Box Loss (MA)')\n    axs[1, 0].fill_between(results_df['epoch'], \n                         results_df['train/box_loss'].rolling(window=window, min_periods=1).mean() - \n                         results_df['train/box_loss'].rolling(window=window, min_periods=1).std(),\n                         results_df['train/box_loss'].rolling(window=window, min_periods=1).mean() + \n                         results_df['train/box_loss'].rolling(window=window, min_periods=1).std(),\n                         alpha=0.3)\n    axs[1, 0].plot(results_df['epoch'], results_df['val/box_loss'].rolling(window=window, min_periods=1).mean(), \n                  label='Val Box Loss (MA)')\n    axs[1, 0].fill_between(results_df['epoch'], \n                         results_df['val/box_loss'].rolling(window=window, min_periods=1).mean() - \n                         results_df['val/box_loss'].rolling(window=window, min_periods=1).std(),\n                         results_df['val/box_loss'].rolling(window=window, min_periods=1).mean() + \n                         results_df['val/box_loss'].rolling(window=window, min_periods=1).std(),\n                         alpha=0.3)\n    axs[1, 0].set_title('Box Loss Stability (Moving Average ± StdDev)')\n    axs[1, 0].set_xlabel('Epoch')\n    axs[1, 0].set_ylabel('Loss')\n    axs[1, 0].legend()\n    axs[1, 0].grid(True, linestyle='--', alpha=0.7)\n    \n    # Plot 4: Train vs Val Loss Ratio (stability indicator)\n    train_val_ratio_box = results_df['train/box_loss'] / results_df['val/box_loss']\n    train_val_ratio_cls = results_df['train/cls_loss'] / results_df['val/cls_loss']\n    train_val_ratio_dfl = results_df['train/dfl_loss'] / results_df['val/dfl_loss']\n    \n    axs[1, 1].plot(results_df['epoch'], train_val_ratio_box, label='Box Loss Ratio')\n    axs[1, 1].plot(results_df['epoch'], train_val_ratio_cls, label='Cls Loss Ratio')\n    axs[1, 1].plot(results_df['epoch'], train_val_ratio_dfl, label='DFL Loss Ratio')\n    axs[1, 1].axhline(y=1.0, color='r', linestyle='--', label='Ideal Ratio = 1.0')\n    axs[1, 1].set_title('Train/Val Loss Ratio (Stability Indicator)')\n    axs[1, 1].set_xlabel('Epoch')\n    axs[1, 1].set_ylabel('Train/Val Ratio')\n    axs[1, 1].legend()\n    axs[1, 1].grid(True, linestyle='--', alpha=0.7)\n    \n    plt.tight_layout()\n    plt.savefig(os.path.join('/kaggle/working', 'training_stability_analysis.png'))\n    plt.show()\n    \n    # Calculate and print stability metrics\n    print(\"Training Stability Analysis:\")\n    print(f\"Box Loss Stability (StdDev): Train = {results_df['train/box_loss'].std():.4f}, Val = {results_df['val/box_loss'].std():.4f}\")\n    print(f\"Cls Loss Stability (StdDev): Train = {results_df['train/cls_loss'].std():.4f}, Val = {results_df['val/cls_loss'].std():.4f}\")\n    print(f\"DFL Loss Stability (StdDev): Train = {results_df['train/dfl_loss'].std():.4f}, Val = {results_df['val/dfl_loss'].std():.4f}\")\n    \n    # Calculate final vs. initial loss ratio (lower is better)\n    loss_reduction_train = results_df['train/box_loss'].iloc[-1] / results_df['train/box_loss'].iloc[0]\n    loss_reduction_val = results_df['val/box_loss'].iloc[-1] / results_df['val/box_loss'].iloc[0]\n    print(f\"Loss Reduction (Final/Initial): Train = {loss_reduction_train:.4f}, Val = {loss_reduction_val:.4f}\")\n    \n    return results_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T14:42:49.540654Z","iopub.execute_input":"2025-04-01T14:42:49.540865Z","iopub.status.idle":"2025-04-01T14:42:49.559718Z","shell.execute_reply.started":"2025-04-01T14:42:49.540846Z","shell.execute_reply":"2025-04-01T14:42:49.558998Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\ndef predict_on_samples(model, num_samples=4):\n    \"\"\"\n    Run predictions on random validation samples and display results.\n    \n    Args:\n        model: Trained YOLO model.\n        num_samples (int): Number of random samples to test.\n    \"\"\"\n    val_dir = os.path.join(yolo_dataset_dir, 'images', 'val')\n    if not os.path.exists(val_dir):\n        print(f\"Validation directory not found at {val_dir}\")\n        val_dir = os.path.join(yolo_dataset_dir, 'images', 'train')\n        print(f\"Using train directory for predictions instead: {val_dir}\")\n        \n    if not os.path.exists(val_dir):\n        print(\"No images directory found for predictions\")\n        return\n    \n    val_images = os.listdir(val_dir)\n    if len(val_images) == 0:\n        print(\"No images found for prediction\")\n        return\n    \n    num_samples = min(num_samples, len(val_images))\n    samples = random.sample(val_images, num_samples)\n    \n    fig, axes = plt.subplots(2, 2, figsize=(12, 12))\n    axes = axes.flatten()\n    \n    for i, img_file in enumerate(samples):\n        if i >= len(axes):\n            break\n            \n        img_path = os.path.join(val_dir, img_file)\n        results = model.predict(img_path, conf=0.25)[0]\n        img = Image.open(img_path)\n        axes[i].imshow(np.array(img), cmap='gray')\n        \n        # Draw ground truth box if available (extracted from filename)\n        try:\n            parts = img_file.split('_')\n            y_part = [p for p in parts if p.startswith('y')]\n            x_part = [p for p in parts if p.startswith('x')]\n            if y_part and x_part:\n                y_gt = int(y_part[0][1:])\n                x_gt = int(x_part[0][1:].split('.')[0])\n                box_size = 24\n                rect_gt = Rectangle((x_gt - box_size//2, y_gt - box_size//2), box_size, box_size,\n                                      linewidth=1, edgecolor='g', facecolor='none')\n                axes[i].add_patch(rect_gt)\n        except:\n            pass\n        \n        if len(results.boxes) > 0:\n            boxes = results.boxes.xyxy.cpu().numpy()\n            confs = results.boxes.conf.cpu().numpy()\n            for box, conf in zip(boxes, confs):\n                x1, y1, x2, y2 = box\n                rect_pred = Rectangle((x1, y1), x2-x1, y2-y1, linewidth=1, edgecolor='r', facecolor='none')\n                axes[i].add_patch(rect_pred)\n                axes[i].text(x1, y1-5, f'{conf:.2f}', color='red')\n        \n        axes[i].set_title(f\"Image: {img_file}\\nGT (green) vs Pred (red)\")\n    \n    plt.tight_layout()\n    plt.savefig(os.path.join('/kaggle/working', 'predictions.png'))\n    plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T14:42:49.560390Z","iopub.execute_input":"2025-04-01T14:42:49.560579Z","iopub.status.idle":"2025-04-01T14:42:49.576662Z","shell.execute_reply.started":"2025-04-01T14:42:49.560556Z","shell.execute_reply":"2025-04-01T14:42:49.575962Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def prepare_dataset():\n    \"\"\"\n    Check if the dataset exists and create/fix a proper YAML file for training.\n    \n    Returns:\n        str: Path to the YAML file to use for training.\n    \"\"\"\n    train_images_dir = os.path.join(yolo_dataset_dir, 'images', 'train')\n    val_images_dir = os.path.join(yolo_dataset_dir, 'images', 'val')\n    train_labels_dir = os.path.join(yolo_dataset_dir, 'labels', 'train')\n    val_labels_dir = os.path.join(yolo_dataset_dir, 'labels', 'val')\n    \n    print(f\"Directory status:\")\n    print(f\"- Train images exists: {os.path.exists(train_images_dir)}\")\n    print(f\"- Val images exists: {os.path.exists(val_images_dir)}\")\n    print(f\"- Train labels exists: {os.path.exists(train_labels_dir)}\")\n    print(f\"- Val labels exists: {os.path.exists(val_labels_dir)}\")\n    \n    original_yaml_path = os.path.join(yolo_dataset_dir, 'dataset.yaml')\n    if os.path.exists(original_yaml_path):\n        print(f\"Found original dataset.yaml at {original_yaml_path}\")\n        return fix_yaml_paths(original_yaml_path)\n    else:\n        print(\"Original dataset.yaml not found, creating a new one\")\n        yaml_data = {\n            'path': yolo_dataset_dir,\n            'train': 'images/train',\n            'val': 'images/train' if not os.path.exists(val_images_dir) else 'images/val',\n            'names': {0: 'motor'}\n        }\n        new_yaml_path = \"/kaggle/working/dataset.yaml\"\n        with open(new_yaml_path, 'w') as f:\n            yaml.dump(yaml_data, f)\n        print(f\"Created new YAML at {new_yaml_path}\")\n        return new_yaml_path\n\ndef main():\n    print(\"Starting improved YOLO training process...\")\n    yaml_path = prepare_dataset()\n    print(f\"Using YAML file: {yaml_path}\")\n    \n    # Visualize learning rate schedule before training\n    print(\"\\nVisualizing learning rate schedule...\")\n    lr_schedule = plot_learning_rate_schedule(epochs=50, lr0=0.0005, lrf=0.05, warmup_epochs=3)\n    \n    print(\"\\nStarting YOLO training with stabilized hyperparameters...\")\n    model, results = train_yolo_model(\n        yaml_path,\n        pretrained_weights_path=yolo_pretrained_weights,\n        epochs=50  # Increase epochs for better convergence\n    )\n    \n    print(\"\\nTraining complete!\")\n    \n    # Analyze training stability\n    results_csv = os.path.join(yolo_weights_dir, 'motor_detector_stable', 'results.csv')\n    if os.path.exists(results_csv):\n        print(\"\\nAnalyzing training stability...\")\n        stability_analysis = analyze_training_stability(results_csv)\n    \n    print(\"\\nRunning predictions on sample images...\")\n    predict_on_samples(model, num_samples=4)\n    \n    print(\"\\nTraining process complete. Check the output directory for results and visualizations.\")\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-01T14:44:29.636709Z","iopub.execute_input":"2025-04-01T14:44:29.637001Z","iopub.status.idle":"2025-04-01T14:51:17.266224Z","shell.execute_reply.started":"2025-04-01T14:44:29.636975Z","shell.execute_reply":"2025-04-01T14:51:17.264814Z"}},"outputs":[],"execution_count":null}]}