{"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"},{"sourceId":226368929,"sourceType":"kernelVersion"},{"sourceId":236157596,"sourceType":"kernelVersion"}],"dockerImageVersionId":30919,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# BYU Locating Flagellar Motors\n\n## YOLO Model Training Notebook\n\nThis is the third notebook in a series for the BYU Locating Bacterial Flagellar Motors 2025 Kaggle challenge. This notebook handles the training of YOLOv8 object detection models on our prepared dataset.\n\n### Notebook Series:\n1. **[Parse Data](https://www.kaggle.com/code/andrewjdarley/parse-data)**: 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 (Current)**: 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 training notebook implements a full YOLOv8 training pipeline for detecting bacterial flagellar motors in tomographic slices. The notebook:\n\n1. **Dataset Configuration**: Sets up and validates the YOLO-format dataset YAML configuration\n2. **Model Initialization**: Loads pre-trained YOLOv8 weights for transfer learning\n3. **Training Process**: Fine tunes the model with early stopping and periodic checkpoints\n4. **Loss Visualization**: Plots training and validation dfl loss curves to monitor progress\n5. **Performance Evaluation**: Tests the trained model on random validation samples\n6. **Model Export**: Saves the trained weights for use in the submission notebook","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19"}},{"cell_type":"code","source":"!pip install ultralytics","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-30T10:49:48.421128Z","iopub.execute_input":"2025-04-30T10:49:48.421489Z","iopub.status.idle":"2025-04-30T10:49:54.502064Z","shell.execute_reply.started":"2025-04-30T10:49:48.421461Z","shell.execute_reply":"2025-04-30T10:49:54.500534Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport torch\nimport numpy as np\nimport random\nfrom PIL import Image\nimport matplotlib.pyplot as plt\nfrom matplotlib.patches import Rectangle\nfrom ultralytics import YOLO\nimport yaml\nimport pandas as pd\nimport json\nimport torch.nn as nn\nfrom ultralytics import YOLO\nfrom ultralytics.nn.modules import Conv, C2f\nfrom typing import List, Optional, Dict, Any, Union\n\n\n# Set random seeds for reproducibility\nnp.random.seed(42)\nrandom.seed(42)\ntorch.manual_seed(42)\n\n# Define paths for Kaggle environment\nyolo_dataset_dir = \"/kaggle/input/parse-data/yolo_dataset\"\nyolo_weights_dir = \"/kaggle/working/yolo_weights\"\nyolo_pretrained_weights = \"yolov8n.pt\"  # Path to pre-downloaded weights\n\n# Create weights directory if it doesn't exist\nos.makedirs(yolo_weights_dir, exist_ok=True)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-30T10:49:54.503761Z","iopub.execute_input":"2025-04-30T10:49:54.504129Z","iopub.status.idle":"2025-04-30T10:49:59.319982Z","shell.execute_reply.started":"2025-04-30T10:49:54.504099Z","shell.execute_reply":"2025-04-30T10:49:59.318725Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from ultralytics.utils import LOGGER, checks\n# from ultralytics.utils.torch_utils import RANK\n# from ultralytics.utils.files import yaml_load\n# from ultralytics.utils.downloads import attempt_load_one_weight\nfrom ultralytics.cfg import DEFAULT_CFG_DICT, TASK2DATA\nimport types\n\ndef custom_train(self, trainer=None, **kwargs):\n    self._check_is_pytorch_model()\n    if hasattr(self.session, \"model\") and self.session.model.id:\n        if any(kwargs):\n            LOGGER.warning(\"using HUB training arguments, ignoring local training arguments.\")\n        kwargs = self.session.train_args\n\n    checks.check_pip_update_available()\n\n    overrides = yaml_load(checks.check_yaml(kwargs[\"cfg\"])) if kwargs.get(\"cfg\") else self.overrides\n    custom = {\n        \"data\": overrides.get(\"data\") or DEFAULT_CFG_DICT[\"data\"] or TASK2DATA[self.task],\n        \"model\": self.overrides[\"model\"],\n        \"task\": self.task,\n    }\n    args = {**overrides, **custom, **kwargs, \"mode\": \"train\"}\n    if args.get(\"resume\"):\n        args[\"resume\"] = self.ckpt_path\n\n    self.trainer = (trainer or self._smart_load(\"trainer\"))(overrides=args, _callbacks=self.callbacks)\n    if not args.get(\"resume\"):\n        if self.model:\n            self.trainer.model = self.model\n        else:\n            self.trainer.model = self.trainer.get_model(weights=self.model if self.ckpt else None, cfg=self.model.yaml)\n            self.model = self.trainer.model\n\n    self.trainer.hub_session = self.session\n    self.trainer.train()\n\n    if RANK in {-1, 0}:\n        ckpt = self.trainer.best if self.trainer.best.exists() else self.trainer.last\n        self.model, self.ckpt = attempt_load_one_weight(ckpt)\n        self.overrides = self.model.args\n        self.metrics = getattr(self.trainer.validator, \"metrics\", None)\n    return self.metrics\n\n# Patch it\n\n\nclass SelfAttention(nn.Module):\n    \"\"\"Self-attention module for YOLOv8 backbone.\"\"\"\n    \n    def __init__(self, channels, reduction_ratio=8):\n        super(SelfAttention, self).__init__()\n        self.query = nn.Conv2d(channels, channels // reduction_ratio, kernel_size=1)\n        self.key = nn.Conv2d(channels, channels // reduction_ratio, kernel_size=1)\n        self.value = nn.Conv2d(channels, channels, kernel_size=1)\n        self.gamma = nn.Parameter(torch.zeros(1))  # Learnable scaling parameter\n        self.softmax = nn.Softmax(dim=-1)\n        \n    def forward(self, x):\n        batch_size, channels, height, width = x.size()\n        \n        # Reshape for attention computation\n        proj_query = self.query(x).view(batch_size, -1, height * width).permute(0, 2, 1)  # B x HW x C'\n        proj_key = self.key(x).view(batch_size, -1, height * width)  # B x C' x HW\n        \n        # Compute attention map\n        energy = torch.bmm(proj_query, proj_key)  # B x HW x HW\n        attention = self.softmax(energy)  # B x HW x HW\n        \n        # Apply attention\n        proj_value = self.value(x).view(batch_size, -1, height * width)  # B x C x HW\n        out = torch.bmm(proj_value, attention.permute(0, 2, 1))\n        out = out.view(batch_size, channels, height, width)\n        \n        # Residual connection with learnable scale\n        out = self.gamma * out + x\n        \n        return out\n\n\nclass C2fAttention(C2f):\n    \"\"\"C2f block with added self-attention mechanism.\"\"\"\n    \n    def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):\n        # Initialize the parent C2f class\n        super().__init__(c1, c2, n, shortcut, g, e)\n        # Add the self-attention layer\n        self.attention = SelfAttention(c2)\n    \n    def forward(self, x):\n        # Use the original C2f forward pass - access through __call__ or super().forward\n        try:\n            # First try using direct forward call\n            y = super().forward(x)\n        except:\n            # If that fails, fall back to standard implementation similar to C2f\n            y = self.cv2(torch.cat((self.m(self.cv1(x)), x if self.shortcut else self.cv3(x)), 1))\n        \n        # Add self-attention mechanism\n        return self.attention(y)\n\n\n# def apply_attention_to_yolo(model):\n#     \"\"\"Apply strong self-attention to early + middle C2f layers of YOLO model for tiny object detection.\"\"\"\n#     detection_model = model.model\n#     backbone_indices = []\n\n#     for i, layer in enumerate(detection_model.model):\n#         if isinstance(layer, C2f):\n#             backbone_indices.append(i)\n    \n#     print(f\"Found {len(backbone_indices)} C2f layers at indices: {backbone_indices}\")\n\n#     # Modify early, middle, and (optionally) slightly later layers\n#     indices_to_modify = []\n#     if len(backbone_indices) >= 4:\n#         for i in range(len(backbone_indices)//2):\n#             indices_to_modify.append(backbone_indices[i])\n#             # = [\n#             #     backbone_indices[0],  # very early\n#             #     backbone_indices[1],  # second C2f (still early)\n#             #     backbone_indices[len(backbone_indices)//2],  # middle\n#             # ]\n#     else:\n#         indices_to_modify = backbone_indices\n\n#     print(f\"Modifying C2f layers at indices: {indices_to_modify}\")\n\n#     for idx in indices_to_modify:\n#         layer = detection_model.model[idx]\n#         c1 = layer.cv1.conv.in_channels\n#         c2 = layer.cv1.conv.out_channels * 2\n#         n = len(layer.m)\n#         shortcut = getattr(layer, 'shortcut', True)\n#         g = layer.cv1.conv.groups\n\n#         try:\n#             e = float(c1 / (layer.cv1.conv.out_channels * 2))\n#         except:\n#             e = 0.5\n        \n#         print(f\"Creating C2fAttention with params: c1={c1}, c2={c2}, n={n}, shortcut={shortcut}, g={g}, e={e}\")\n#         detection_model.model[idx] = C2fAttention(c1, c2, n, shortcut, g, e)\n\n#     return model\n\ndef apply_attention_to_yolo(model):\n    \"\"\"Apply strong self-attention to early + middle C2f layers of YOLO model for tiny object detection.\"\"\"\n    detection_model = model.model\n    backbone_indices = []\n\n    for i, layer in enumerate(detection_model.model):\n        if isinstance(layer, C2f):\n            backbone_indices.append(i)\n\n    print(f\"Found {len(backbone_indices)} C2f layers at indices: {backbone_indices}\")\n\n    indices_to_modify = []\n    if len(backbone_indices) >= 4:\n        for i in range(len(backbone_indices)//2):\n            indices_to_modify.append(backbone_indices[i])\n    else:\n        indices_to_modify = backbone_indices\n\n    print(f\"Modifying C2f layers at indices: {indices_to_modify}\")\n\n    for idx in indices_to_modify:\n        original_layer = detection_model.model[idx]\n\n        c1 = original_layer.cv1.conv.in_channels\n        c2 = original_layer.cv1.conv.out_channels * 2\n        n = len(original_layer.m)\n        shortcut = getattr(original_layer, 'shortcut', True)\n        g = original_layer.cv1.conv.groups\n\n        try:\n            e = float(c1 / (original_layer.cv1.conv.out_channels * 2))\n        except:\n            e = 0.5\n\n        print(f\"Creating C2fAttention with params: c1={c1}, c2={c2}, n={n}, shortcut={shortcut}, g={g}, e={e}\")\n        attn_layer = C2fAttention(c1, c2, n, shortcut, g, e)\n        attn_layer.f = getattr(original_layer, 'f', -1)  # ✅ preserve connection info\n\n        detection_model.model[idx] = attn_layer\n\n    return model\n\ndef 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    \n    # Read the original YAML\n    with open(yaml_path, 'r') as f:\n        yaml_data = yaml.safe_load(f)\n    \n    # Update paths to use actual dataset location\n    if 'path' in yaml_data:\n        yaml_data['path'] = yolo_dataset_dir\n    \n    # Create a new fixed YAML in the working directory\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\n\ndef prepare_dataset():\n    \"\"\"\n    Check if dataset exists and create a proper YAML if needed\n    \n    Returns:\n        str: Path to the YAML file to use for training\n    \"\"\"\n    # Check if images exist\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 directory existence status\n    print(f\"Directory status:\")\n    print(f\"- Train images dir exists: {os.path.exists(train_images_dir)}\")\n    print(f\"- Val images dir exists: {os.path.exists(val_images_dir)}\")\n    print(f\"- Train labels dir exists: {os.path.exists(train_labels_dir)}\")\n    print(f\"- Val labels dir exists: {os.path.exists(val_labels_dir)}\")\n    \n    # Check for original YAML file\n    original_yaml_path = os.path.join(yolo_dataset_dir, 'dataset.yaml')\n    \n    if os.path.exists(original_yaml_path):\n        print(f\"Found original dataset.yaml at {original_yaml_path}\")\n        # Fix the paths in the YAML\n        return fix_yaml_paths(original_yaml_path)\n    else:\n        print(f\"Original dataset.yaml not found, creating a new one\")\n        \n        # Create a new YAML file\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        \n        new_yaml_path = \"/kaggle/working/dataset.yaml\"\n        with open(new_yaml_path, 'w') as f:\n            yaml.dump(yaml_data, f)\n            \n        print(f\"Created new YAML at {new_yaml_path}\")\n        return new_yaml_path\n\n\n\nyaml_path = prepare_dataset()\n\n# Now use as normal\nmodel = YOLO(\"yolov8n.pt\")\nmodel.train = types.MethodType(custom_train, model)\nmodel = apply_attention_to_yolo(model)\n# model.model\nresults = model.train(data=yaml_path, epochs=1, freeze = 0)\nprint(model.model)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-30T10:51:01.786823Z","iopub.execute_input":"2025-04-30T10:51:01.787516Z","execution_failed":"2025-04-30T10:51:53.008Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class SelfAttention(nn.Module):\n    \"\"\"Self-attention module for YOLOv8 backbone.\"\"\"\n    \n    def __init__(self, channels, reduction_ratio=8):\n        super(SelfAttention, self).__init__()\n        self.query = nn.Conv2d(channels, channels // reduction_ratio, kernel_size=1)\n        self.key = nn.Conv2d(channels, channels // reduction_ratio, kernel_size=1)\n        self.value = nn.Conv2d(channels, channels, kernel_size=1)\n        self.gamma = nn.Parameter(torch.zeros(1))  # Learnable scaling parameter\n        self.softmax = nn.Softmax(dim=-1)\n        \n    def forward(self, x):\n        batch_size, channels, height, width = x.size()\n        \n        # Reshape for attention computation\n        proj_query = self.query(x).view(batch_size, -1, height * width).permute(0, 2, 1)  # B x HW x C'\n        proj_key = self.key(x).view(batch_size, -1, height * width)  # B x C' x HW\n        \n        # Compute attention map\n        energy = torch.bmm(proj_query, proj_key)  # B x HW x HW\n        attention = self.softmax(energy)  # B x HW x HW\n        \n        # Apply attention\n        proj_value = self.value(x).view(batch_size, -1, height * width)  # B x C x HW\n        out = torch.bmm(proj_value, attention.permute(0, 2, 1))\n        out = out.view(batch_size, channels, height, width)\n        \n        # Residual connection with learnable scale\n        out = self.gamma * out + x\n        \n        return out\n\n\nclass C2fAttention(C2f):\n    \"\"\"C2f block with added self-attention mechanism.\"\"\"\n    \n    def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):\n        # Initialize the parent C2f class\n        super().__init__(c1, c2, n, shortcut, g, e)\n        # Add the self-attention layer\n        self.attention = SelfAttention(c2)\n    \n    def forward(self, x):\n        # Use the original C2f forward pass - access through __call__ or super().forward\n        try:\n            # First try using direct forward call\n            y = super().forward(x)\n        except:\n            # If that fails, fall back to standard implementation similar to C2f\n            y = self.cv2(torch.cat((self.m(self.cv1(x)), x if self.shortcut else self.cv3(x)), 1))\n        \n        # Add self-attention mechanism\n        return self.attention(y)\n\n\ndef apply_attention_to_yolo(model):\n    \"\"\"Apply strong self-attention to early + middle C2f layers of YOLO model for tiny object detection.\"\"\"\n    detection_model = model.model\n    backbone_indices = []\n\n    for i, layer in enumerate(detection_model.model):\n        if isinstance(layer, C2f):\n            backbone_indices.append(i)\n    \n    print(f\"Found {len(backbone_indices)} C2f layers at indices: {backbone_indices}\")\n\n    # Modify early, middle, and (optionally) slightly later layers\n    indices_to_modify = []\n    if len(backbone_indices) >= 4:\n        indices_to_modify = [\n            backbone_indices[0],  # very early\n            backbone_indices[1],  # second C2f (still early)\n            backbone_indices[len(backbone_indices)//2],  # middle\n        ]\n    else:\n        indices_to_modify = backbone_indices\n\n    print(f\"Modifying C2f layers at indices: {indices_to_modify}\")\n\n    for idx in indices_to_modify:\n        layer = detection_model.model[idx]\n        c1 = layer.cv1.conv.in_channels\n        c2 = layer.cv1.conv.out_channels * 2\n        n = len(layer.m)\n        shortcut = getattr(layer, 'shortcut', True)\n        g = layer.cv1.conv.groups\n\n        try:\n            e = float(c1 / (layer.cv1.conv.out_channels * 2))\n        except:\n            e = 0.5\n        \n        print(f\"Creating C2fAttention with params: c1={c1}, c2={c2}, n={n}, shortcut={shortcut}, g={g}, e={e}\")\n        detection_model.model[idx] = C2fAttention(c1, c2, n, shortcut, g, e)\n\n    return model","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-30T10:40:55.196303Z","iopub.status.idle":"2025-04-30T10:40:55.196653Z","shell.execute_reply":"2025-04-30T10:40:55.196497Z"}},"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    \n    # Read the original YAML\n    with open(yaml_path, 'r') as f:\n        yaml_data = yaml.safe_load(f)\n    \n    # Update paths to use actual dataset location\n    if 'path' in yaml_data:\n        yaml_data['path'] = yolo_dataset_dir\n    \n    # Create a new fixed YAML in the working directory\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\n\ndef plot_dfl_loss_curve(run_dir):\n    \"\"\"\n    Plot the DFL loss curves for train and validation, marking the best model\n    \n    Args:\n        run_dir (str): Directory where the training results are stored\n    \"\"\"\n    # Path to the results CSV file\n    results_csv = os.path.join(run_dir, 'results.csv')\n    \n    if not os.path.exists(results_csv):\n        print(f\"Results file not found at {results_csv}\")\n        return\n    \n    # Read results CSV\n    results_df = pd.read_csv(results_csv)\n    \n    # Check if DFL loss columns exist\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    # Find the epoch with the best validation loss\n    best_epoch = results_df[val_dfl_col].idxmin()\n    best_val_loss = results_df.loc[best_epoch, val_dfl_col]\n    \n    # Create the plot\n    plt.figure(figsize=(10, 6))\n    \n    # Plot training and validation losses\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    \n    # Mark the best model with a vertical line\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    \n    # Add labels and legend\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    # Save the plot in the same directory as weights\n    plot_path = os.path.join(run_dir, 'dfl_loss_curve.png')\n    plt.savefig(plot_path)\n    \n    # Also save it to the working directory for easier access\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 the best epoch info\n    return best_epoch, best_val_loss\n\ndef train_yolo_model(yaml_path, pretrained_weights_path, epochs=30, batch_size=16, img_size=640, \n                    use_attention=True, yolo_weights_dir='runs/detect', devices=0):\n    \"\"\"\n    Train a YOLO model on the prepared dataset with option to use self-attention architecture\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        batch_size (int): Batch size for training\n        img_size (int): Image size for training\n        use_attention (bool): Whether to use the self-attention architecture\n        yolo_weights_dir (str): Directory to save training results\n    \"\"\"\n    print(f\"Loading pre-trained weights from: {pretrained_weights_path}\")\n    \n    # Load the standard YOLOv8 model\n    model = YOLO(pretrained_weights_path)\n    \n    # Apply attention mechanism if requested\n    if use_attention:\n        print(\"Applying self-attention to YOLOv8 architecture\")\n        try:\n            model = apply_attention_to_yolo(model)\n            print(\"Successfully applied self-attention to model\")\n        except Exception as e:\n            print(f\"Error applying self-attention: {e}\")\n            print(\"Falling back to standard YOLOv8 architecture\")\n            # Reload the original model to ensure it's in a clean state\n            model = YOLO(pretrained_weights_path)\n    else:\n        print(\"Using standard YOLOv8 architecture\")\n    \n    # Train the model with early stopping\n    results = model.train(\n        data=yaml_path,\n        epochs=epochs,\n        batch=batch_size,\n        imgsz=img_size,\n        project=yolo_weights_dir,\n        device=devices,\n        name='motor_detector',\n        exist_ok=True,\n        # patience=10,              # Early stopping if no improvement for 5 epochs\n        save_period=2,           # Save checkpoints every 5 epochs\n        val=True,                # Ensure validation is performed\n        verbose=True,             # Show detailed output during training\n        single_cls=True,\n        degrees=180,\n        translate=0.2,\n        shear=2.0,\n        perspective=0.001,\n        flipud=0.5,\n        fliplr=0.5,\n        mixup=0.25\n    )\n    \n    # Get the path to the run directory\n    run_dir = os.path.join(yolo_weights_dir, 'motor_detector')\n    \n    # Plot and save the loss curve\n    best_epoch_info = plot_dfl_loss_curve(run_dir)\n    \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\n\n\ndef predict_on_samples(model, num_samples=4, conf_threshold=0.25):\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    # Get validation images\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        # Try train directory instead if val doesn't exist\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    \n    if len(val_images) == 0:\n        print(\"No images found for prediction\")\n        return\n    \n    # Select random samples\n    num_samples = min(num_samples, len(val_images))\n    samples = random.sample(val_images, num_samples)\n    \n    # Create figure\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        \n        # Run prediction\n        results = model.predict(img_path, conf=conf_threshold, batch=4)[0]\n        \n        # Load and display the image\n        img = Image.open(img_path)\n        axes[i].imshow(np.array(img), cmap='gray')\n        \n        # Draw ground truth box if available (from filename)\n        try:\n            # This assumes your filenames contain coordinates in a specific format\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            \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                \n                box_size = 24\n                rect_gt = Rectangle((x_gt - box_size//2, y_gt - box_size//2), \n                              box_size, box_size, \n                              linewidth=1, edgecolor='g', facecolor='none')\n                axes[i].add_patch(rect_gt)\n        except:\n            pass  # Skip ground truth if parsing fails\n        \n        # Draw predicted boxes (red)\n        if len(results.boxes) > 0:\n            boxes = results.boxes.xyxy.cpu().numpy()\n            confs = results.boxes.conf.cpu().numpy()\n            \n            for box, conf in zip(boxes, confs):\n                x1, y1, x2, y2 = box\n                rect_pred = Rectangle((x1, y1), x2-x1, y2-y1, \n                                     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}\\nGround Truth (green) vs Prediction (red)\")\n    \n    plt.tight_layout()\n    \n    # Save the predictions plot\n    plt.savefig(os.path.join('/kaggle/working', 'predictions.png'))\n    plt.show()\n\n# Check and create a dataset YAML if needed\ndef prepare_dataset():\n    \"\"\"\n    Check if dataset exists and create a proper YAML if needed\n    \n    Returns:\n        str: Path to the YAML file to use for training\n    \"\"\"\n    # Check if images exist\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 directory existence status\n    print(f\"Directory status:\")\n    print(f\"- Train images dir exists: {os.path.exists(train_images_dir)}\")\n    print(f\"- Val images dir exists: {os.path.exists(val_images_dir)}\")\n    print(f\"- Train labels dir exists: {os.path.exists(train_labels_dir)}\")\n    print(f\"- Val labels dir exists: {os.path.exists(val_labels_dir)}\")\n    \n    # Check for original YAML file\n    original_yaml_path = os.path.join(yolo_dataset_dir, 'dataset.yaml')\n    \n    if os.path.exists(original_yaml_path):\n        print(f\"Found original dataset.yaml at {original_yaml_path}\")\n        # Fix the paths in the YAML\n        return fix_yaml_paths(original_yaml_path)\n    else:\n        print(f\"Original dataset.yaml not found, creating a new one\")\n        \n        # Create a new YAML file\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        \n        new_yaml_path = \"/kaggle/working/dataset.yaml\"\n        with open(new_yaml_path, 'w') as f:\n            yaml.dump(yaml_data, f)\n            \n        print(f\"Created new YAML at {new_yaml_path}\")\n        return new_yaml_path\n\n# Main execution\ndef main():\n    print(\"Starting YOLO training process...\")\n    \n    # Prepare dataset and get YAML path\n    yaml_path = prepare_dataset()\n    print(f\"Using YAML file: {yaml_path}\")\n    \n    # Print YAML file contents\n    with open(yaml_path, 'r') as f:\n        yaml_content = f.read()\n    print(f\"YAML file contents:\\n{yaml_content}\")\n    \n    # Train model\n    print(\"\\nStarting YOLO training...\")\n    model, results = train_yolo_model(\n        yaml_path,\n        pretrained_weights_path=yolo_pretrained_weights,\n        epochs=100,  # Using 30 epochs instead of 100 for faster training\n        devices=0,\n        use_attention=False\n    )\n    \n    print(\"\\nTraining complete!\")\n    \n    # Run predictions\n    print(\"\\nRunning predictions on sample images...\")\n    predict_on_samples(model, num_samples=4, conf_threshold=0.1)\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-27T16:49:32.032885Z","iopub.execute_input":"2025-04-27T16:49:32.033188Z"}},"outputs":[],"execution_count":null}]}