{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":91249,"databundleVersionId":11294684,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":226368929,"sourceType":"kernelVersion"}],"dockerImageVersionId":30919,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"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-05-26T17:36:39.904948Z","iopub.execute_input":"2025-05-26T17:36:39.905235Z","iopub.status.idle":"2025-05-26T17:36:45.540457Z","shell.execute_reply.started":"2025-05-26T17:36:39.905202Z","shell.execute_reply":"2025-05-26T17:36:45.539643Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nimport os\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\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 = \"yolo11n.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\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 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=680):\n    \"\"\"\n    Train a YOLO model on the prepared dataset\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    \"\"\"\n    print(f\"Loading pre-trained weights from: {pretrained_weights_path}\")\n    \n    # Load a pre-trained YOLOv11 model\n    model = YOLO(pretrained_weights_path)\n     # Train the model with early stopping\n    results = model.train(\n        data=yaml_path,\n        epochs=epochs,\n        batch=batch_size,\n        imgsz=680,\n        project=yolo_weights_dir,\n        name='motor_detector',\n        exist_ok=True,\n        patience=10,              # Early stopping if no improvement for 5 epochs\n        save_period=5,           # Save checkpoints every 5 epochs\n        val=True,                # Ensure validation is performed\n        verbose=True,             # Show detailed output during training\n        dfl=0.3,\n       # cls=2,\n        box=2,\n        mosaic=1,  # Enable mosaic\n        scale=0.5,   # Random scaling for small objects\n       # hsv_h=0.015,\n       # hsv_s=0.7,\n       # hsv_v=0.4\n        #cos_lr=True,\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\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    # 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=0.25)[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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T17:49:57.06904Z","iopub.execute_input":"2025-05-26T17:49:57.069383Z","iopub.status.idle":"2025-05-26T17:49:57.091115Z","shell.execute_reply.started":"2025-05-26T17:49:57.06934Z","shell.execute_reply":"2025-05-26T17:49:57.09044Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"predict_on_samples(model, num_samples=4)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-25T21:50:24.833764Z","iopub.execute_input":"2025-05-25T21:50:24.834082Z","iopub.status.idle":"2025-05-25T21:50:27.555923Z","shell.execute_reply.started":"2025-05-25T21:50:24.834058Z","shell.execute_reply":"2025-05-25T21:50:27.554549Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"### print(\"Starting YOLO training process...\")\n    \n    # Prepare dataset and get YAML path\nyaml_path = prepare_dataset()\nprint(f\"Using YAML file: {yaml_path}\")\n    \n    # Print YAML file contents\nwith open(yaml_path, 'r') as f:\n    yaml_content = f.read()\nprint(f\"YAML file contents:\\n{yaml_content}\")\n\n\n    # Train model\nprint(\"\\nStarting YOLO training...\")\nmodel, results = train_yolo_model(\nyaml_path,\npretrained_weights_path=yolo_pretrained_weights,\nepochs=40  # Using 30 epochs instead of 100 for faster training\n    )\n    \npredict_on_samples(model, num_samples=4)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T17:50:08.75298Z","iopub.execute_input":"2025-05-26T17:50:08.753335Z","iopub.status.idle":"2025-05-26T18:27:36.948212Z","shell.execute_reply.started":"2025-05-26T17:50:08.753294Z","shell.execute_reply":"2025-05-26T18:27:36.947182Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nimport torch\nimport cv2\nfrom tqdm.notebook import tqdm\nfrom ultralytics import YOLO\nimport threading\nimport time\nfrom contextlib import nullcontext\nfrom concurrent.futures import ThreadPoolExecutor\n\n# Set random seed for reproducibility\nnp.random.seed(42)\ntorch.manual_seed(42)\n\n# Define paths\ndata_path = \"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/\"\ntest_dir = os.path.join(data_path, \"test\")\nsubmission_path = \"/kaggle/working/submission.csv\"\n\n# Model path - adjust if your best model is saved in a different location\n#model_path = \"/kaggle/input/train-yolo/yolo_weights/motor_detector/weights/best.pt\"\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T18:35:51.262671Z","iopub.execute_input":"2025-05-26T18:35:51.262997Z","iopub.status.idle":"2025-05-26T18:35:51.269012Z","shell.execute_reply.started":"2025-05-26T18:35:51.262971Z","shell.execute_reply":"2025-05-26T18:35:51.268263Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Detection parameters\nCONFIDENCE_THRESHOLD = 0.5  # Lower threshold to catch more potential motors\nMAX_DETECTIONS_PER_TOMO = 3  # Keep track of top N detections per tomogram\nNMS_IOU_THRESHOLD = 0.2  # Non-maximum suppression threshold for 3D clustering\nCONCENTRATION = 1 # ONLY PROCESS 1/20 slices for fast submission\n\n# GPU profiling context manager\nclass GPUProfiler:\n    def __init__(self, name):\n        self.name = name\n        self.start_time = None\n        \n    def __enter__(self):\n        if torch.cuda.is_available():\n            torch.cuda.synchronize()\n        self.start_time = time.time()\n        return self\n        \n    def __exit__(self, *args):\n        if torch.cuda.is_available():\n            torch.cuda.synchronize()\n        elapsed = time.time() - self.start_time\n        print(f\"[PROFILE] {self.name}: {elapsed:.3f}s\")\n\n# Check GPU availability and set up optimizations\ndevice = 'cuda:0' if torch.cuda.is_available() else 'cpu'\nBATCH_SIZE = 8  # Default batch size, will be adjusted dynamically if GPU available\n\nif device.startswith('cuda'):\n    # Set CUDA optimization flags\n    torch.backends.cudnn.benchmark = True\n    torch.backends.cudnn.deterministic = False\n    torch.backends.cuda.matmul.allow_tf32 = True  # Allow TF32 on Ampere GPUs\n    torch.backends.cudnn.allow_tf32 = True\n    \n    # Print GPU info\n    gpu_name = torch.cuda.get_device_name(0)\n    gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9  # Convert to GB\n    print(f\"Using GPU: {gpu_name} with {gpu_mem:.2f} GB memory\")\n    \n    # Get available GPU memory and set batch size accordingly\n    free_mem = gpu_mem - torch.cuda.memory_allocated(0) / 1e9\n    BATCH_SIZE = max(8, min(32, int(free_mem * 4)))  # 4 images per GB as rough estimate\n    print(f\"Dynamic batch size set to {BATCH_SIZE} based on {free_mem:.2f}GB free memory\")\nelse:\n    print(\"GPU not available, using CPU\")\n    BATCH_SIZE = 4  # Reduce batch size for CPU\n\ndef normalize_slice(slice_data):\n    \"\"\"\n    Normalize slice data using 2nd and 98th percentiles for better contrast\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\ndef preload_image_batch(file_paths):\n    \"\"\"Preload a batch of images to CPU memory\"\"\"\n    images = []\n    for path in file_paths:\n        img = cv2.imread(path)\n        if img is None:\n            # Try with PIL as fallback\n            img = np.array(Image.open(path))\n        images.append(img)\n    return images\n\ndef process_tomogram(tomo_id, model, index=0, total=1):\n    \"\"\"\n    Process a single tomogram and return the most confident motor detection\n    \"\"\"\n    print(f\"Processing tomogram {tomo_id} ({index}/{total})\")\n    \n    # Get all slice files for this tomogram\n    tomo_dir = os.path.join(test_dir, tomo_id)\n    slice_files = sorted([f for f in os.listdir(tomo_dir) if f.endswith('.jpg')])\n    \n    # Apply CONCENTRATION to reduce the number of slices processed\n    # This will process approximately CONCENTRATION fraction of all slices\n    selected_indices = np.linspace(0, len(slice_files)-1, int(len(slice_files) * CONCENTRATION))\n    selected_indices = np.round(selected_indices).astype(int)\n    slice_files = [slice_files[i] for i in selected_indices]\n    \n    print(f\"Processing {len(slice_files)} out of {len(os.listdir(tomo_dir))} slices based on CONCENTRATION={CONCENTRATION}\")\n    \n    # Create a list to store all detections\n    all_detections = []\n    \n    # Create CUDA streams for parallel processing if using GPU\n    if device.startswith('cuda'):\n        streams = [torch.cuda.Stream() for _ in range(min(4, BATCH_SIZE))]\n    else:\n        streams = [None]\n    \n    # Variables for preloading\n    next_batch_thread = None\n    next_batch_images = None\n    \n    # Process slices in batches\n    for batch_start in range(0, len(slice_files), BATCH_SIZE):\n        # Wait for previous preload thread if it exists\n        if next_batch_thread is not None:\n            next_batch_thread.join()\n            next_batch_images = None\n            \n        batch_end = min(batch_start + BATCH_SIZE, len(slice_files))\n        batch_files = slice_files[batch_start:batch_end]\n        \n        # Start preloading next batch\n        next_batch_start = batch_end\n        next_batch_end = min(next_batch_start + BATCH_SIZE, len(slice_files))\n        next_batch_files = slice_files[next_batch_start:next_batch_end] if next_batch_start < len(slice_files) else []\n        \n        if next_batch_files:\n            next_batch_paths = [os.path.join(tomo_dir, f) for f in next_batch_files]\n            next_batch_thread = threading.Thread(target=preload_image_batch, args=(next_batch_paths,))\n            next_batch_thread.start()\n        else:\n            next_batch_thread = None\n        \n        # Split batch across streams for parallel processing\n        sub_batches = np.array_split(batch_files, len(streams))\n        sub_batch_results = []\n        \n        for i, sub_batch in enumerate(sub_batches):\n            if len(sub_batch) == 0:\n                continue\n                \n            stream = streams[i % len(streams)]\n            with torch.cuda.stream(stream) if stream and device.startswith('cuda') else nullcontext():\n                # Process sub-batch\n                sub_batch_paths = [os.path.join(tomo_dir, slice_file) for slice_file in sub_batch]\n                sub_batch_slice_nums = [int(slice_file.split('_')[1].split('.')[0]) for slice_file in sub_batch]\n                \n                # Run inference with profiling\n                with GPUProfiler(f\"Inference batch {i+1}/{len(sub_batches)}\"):\n                    sub_results = model(sub_batch_paths, verbose=False)\n                \n                # Process each result in this sub-batch\n                for j, result in enumerate(sub_results):\n                    if len(result.boxes) > 0:\n                        boxes = result.boxes\n                        for box_idx, confidence in enumerate(boxes.conf):\n                            if confidence >= CONFIDENCE_THRESHOLD:\n                                # Get bounding box coordinates\n                                x1, y1, x2, y2 = boxes.xyxy[box_idx].cpu().numpy()\n                                \n                                # Calculate center coordinates\n                                x_center = (x1 + x2) / 2\n                                y_center = (y1 + y2) / 2\n                                \n                                # Store detection with 3D coordinates\n                                all_detections.append({\n                                    'z': round(sub_batch_slice_nums[j]),\n                                    'y': round(y_center),\n                                    'x': round(x_center),\n                                    'confidence': float(confidence)\n                                })\n        \n        # Synchronize streams\n        if device.startswith('cuda'):\n            torch.cuda.synchronize()\n    \n    # Clean up thread if still running\n    if next_batch_thread is not None:\n        next_batch_thread.join()\n    \n    # 3D Non-Maximum Suppression to merge nearby detections across slices\n    final_detections = perform_3d_nms(all_detections, NMS_IOU_THRESHOLD)\n    \n    # Sort detections by confidence (highest first)\n    final_detections.sort(key=lambda x: x['confidence'], reverse=True)\n    \n    # If there are no detections, return NA values\n    if not final_detections:\n        return {\n            'tomo_id': tomo_id,\n            'Motor axis 0': -1,\n            'Motor axis 1': -1,\n            'Motor axis 2': -1\n        }\n    \n    # Take the detection with highest confidence\n    best_detection = final_detections[0]\n    \n    # Return result with integer coordinates\n    return {\n        'tomo_id': tomo_id,\n        'Motor axis 0': round(best_detection['z']),\n        'Motor axis 1': round(best_detection['y']),\n        'Motor axis 2': round(best_detection['x'])\n    }\n\ndef perform_3d_nms(detections, iou_threshold):\n    \"\"\"\n    Perform 3D Non-Maximum Suppression on detections to merge nearby motors\n    \"\"\"\n    if not detections:\n        return []\n    \n    # Sort by confidence (highest first)\n    detections = sorted(detections, key=lambda x: x['confidence'], reverse=True)\n    \n    # List to store final detections after NMS\n    final_detections = []\n    \n    # Define 3D distance function\n    def distance_3d(d1, d2):\n        return np.sqrt((d1['z'] - d2['z'])**2 + \n                       (d1['y'] - d2['y'])**2 + \n                       (d1['x'] - d2['x'])**2)\n    \n    # Maximum distance threshold (based on box size and slice gap)\n    box_size = 24  # Same as annotation box size\n    distance_threshold = box_size * iou_threshold\n    \n    # Process each detection\n    while detections:\n        # Take the detection with highest confidence\n        best_detection = detections.pop(0)\n        final_detections.append(best_detection)\n        \n        # Filter out detections that are too close to the best detection\n        detections = [d for d in detections if distance_3d(d, best_detection) > distance_threshold]\n    \n    return final_detections\n\ndef debug_image_loading(tomo_id):\n    \"\"\"\n    Debug function to check image loading\n    \"\"\"\n    tomo_dir = os.path.join(test_dir, tomo_id)\n    slice_files = sorted([f for f in os.listdir(tomo_dir) if f.endswith('.jpg')])\n    \n    if not slice_files:\n        print(f\"No image files found in {tomo_dir}\")\n        return\n        \n    print(f\"Found {len(slice_files)} image files in {tomo_dir}\")\n    sample_file = slice_files[len(slice_files)//2]  # Middle slice\n    img_path = os.path.join(tomo_dir, sample_file)\n    \n    # Try different loading methods\n    try:\n        # Method 1: PIL\n        img_pil = Image.open(img_path)\n        img_array_pil = np.array(img_pil)\n        print(f\"PIL Image shape: {img_array_pil.shape}, dtype: {img_array_pil.dtype}\")\n        \n        # Method 2: OpenCV\n        img_cv2 = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)\n        print(f\"OpenCV Image shape: {img_cv2.shape}, dtype: {img_cv2.dtype}\")\n        \n        # Method 3: Convert to RGB\n        img_rgb = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB)\n        print(f\"OpenCV RGB Image shape: {img_rgb.shape}, dtype: {img_rgb.dtype}\")\n        \n        print(\"Image loading successful!\")\n    except Exception as e:\n        print(f\"Error loading image {img_path}: {e}\")\n        \n    # Also test with YOLO's built-in loader\n    try:\n        test_model = model#YOLO(model_path)\n        test_results = test_model([img_path], verbose=False)\n        print(\"YOLO model successfully processed the test image\")\n    except Exception as e:\n        print(f\"Error with YOLO processing: {e}\")\n\ndef generate_submission():\n    \"\"\"\n    Main function to generate the submission file\n    \"\"\"\n    # Get list of test tomograms\n    test_tomos = sorted([d for d in os.listdir(test_dir) if os.path.isdir(os.path.join(test_dir, d))])\n    total_tomos = len(test_tomos)\n    \n    print(f\"Found {total_tomos} tomograms in test directory\")\n    \n    # Debug image loading for the first tomogram\n    if test_tomos:\n        debug_image_loading(test_tomos[0])\n    \n    # Clear GPU cache before starting\n    if torch.cuda.is_available():\n        torch.cuda.empty_cache()\n    \n    # Initialize model once outside the processing loop\n   # print(f\"Loading YOLO model from {model_path}\")\n   # model = YOLO(model_path)\n    model.to(device)\n    \n    # Additional optimizations for inference\n    if device.startswith('cuda'):\n        # Fuse conv and bn layers for faster inference\n        model.fuse()\n        \n        # Enable model half precision (FP16) if on compatible GPU\n        if torch.cuda.get_device_capability(0)[0] >= 7:  # Volta or newer\n            model.model.half()\n            print(\"Using half precision (FP16) for inference\")\n    \n    # Process tomograms with parallelization\n    results = []\n    motors_found = 0\n    \n    # Using ThreadPoolExecutor with max_workers=1 since each worker uses the GPU already\n    # and we're parallelizing within each tomogram processing\n    with ThreadPoolExecutor(max_workers=1) as executor:\n        future_to_tomo = {}\n        \n        # Submit all tomograms for processing\n        for i, tomo_id in enumerate(test_tomos, 1):\n            future = executor.submit(process_tomogram, tomo_id, model, i, total_tomos)\n            future_to_tomo[future] = tomo_id\n        \n        # Process completed futures as they complete\n        for future in future_to_tomo:\n            tomo_id = future_to_tomo[future]\n            try:\n                # Clear CUDA cache between tomograms\n                if torch.cuda.is_available():\n                    torch.cuda.empty_cache()\n                    \n                result = future.result()\n                results.append(result)\n                \n                # Update motors found count\n                has_motor = not pd.isna(result['Motor axis 0'])\n                if has_motor:\n                    motors_found += 1\n                    print(f\"Motor found in {tomo_id} at position: \"\n                          f\"z={result['Motor axis 0']}, y={result['Motor axis 1']}, x={result['Motor axis 2']}\")\n                else:\n                    print(f\"No motor detected in {tomo_id}\")\n                    \n                print(f\"Current detection rate: {motors_found}/{len(results)} ({motors_found/len(results)*100:.1f}%)\")\n            \n            except Exception as e:\n                print(f\"Error processing {tomo_id}: {e}\")\n                # Create a default entry for failed tomograms\n                results.append({\n                    'tomo_id': tomo_id,\n                    'Motor axis 0': -1,\n                    'Motor axis 1': -1,\n                    'Motor axis 2': -1\n                })\n    \n    # Create submission dataframe\n    submission_df = pd.DataFrame(results)\n    \n    # Ensure proper column order\n    submission_df = submission_df[['tomo_id', 'Motor axis 0', 'Motor axis 1', 'Motor axis 2']]\n    \n    # Save the submission file\n    submission_df.to_csv(submission_path, index=False)\n    \n    print(f\"\\nSubmission complete!\")\n    print(f\"Motors detected: {motors_found}/{total_tomos} ({motors_found/total_tomos*100:.1f}%)\")\n    print(f\"Submission saved to: {submission_path}\")\n    \n    # Display first few rows of submission\n    print(\"\\nSubmission preview:\")\n    print(submission_df.head())\n    \n    return submission_df\n\n# Run the submission pipeline\nif __name__ == \"__main__\":\n    # Time entire process\n    start_time = time.time()\n    \n    # Generate submission\n    submission = generate_submission()\n    \n    # Print total execution time\n    elapsed = time.time() - start_time\n    print(f\"\\nTotal execution time: {elapsed:.2f} seconds ({elapsed/60:.2f} minutes)\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-26T18:37:11.151897Z","iopub.execute_input":"2025-05-26T18:37:11.152248Z","iopub.status.idle":"2025-05-26T18:38:02.077333Z","shell.execute_reply.started":"2025-05-26T18:37:11.152218Z","shell.execute_reply":"2025-05-26T18:38:02.076564Z"}},"outputs":[],"execution_count":null}]}