{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"anaconda-cloud":{},"toc":{"base_numbering":1,"nav_menu":{},"number_sections":true,"sideBar":true,"skip_h1_title":false,"title_cell":"Table of Contents","title_sidebar":"Contents","toc_cell":false,"toc_position":{},"toc_section_display":true,"toc_window_display":false},"colab":{"provenance":[]},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":91249,"databundleVersionId":11294684,"sourceType":"competition"},{"sourceId":11558496,"sourceType":"datasetVersion","datasetId":7244332},{"sourceId":226864880,"sourceType":"kernelVersion"},{"sourceId":235993707,"sourceType":"kernelVersion"}],"dockerImageVersionId":31011,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\n# for dirname, _, filenames in os.walk('/kaggle/input'):\n    # for filename in filenames:\n        # print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-25T05:56:28.816237Z","iopub.execute_input":"2025-04-25T05:56:28.816980Z","iopub.status.idle":"2025-04-25T05:56:29.115026Z","shell.execute_reply.started":"2025-04-25T05:56:28.816953Z","shell.execute_reply":"2025-04-25T05:56:29.114414Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!tar xfvz /kaggle/input/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-25T05:56:29.116251Z","iopub.execute_input":"2025-04-25T05:56:29.116638Z","iopub.status.idle":"2025-04-25T05:56:54.355371Z","shell.execute_reply.started":"2025-04-25T05:56:29.116610Z","shell.execute_reply":"2025-04-25T05:56:54.354526Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nfrom pathlib import Path\nimport time\nimport yaml\nimport glob\nimport json\nimport random\nimport shutil\n\nimport cv2\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image, ImageDraw\nimport matplotlib.pyplot as plt\nfrom matplotlib.patches import Rectangle\n\nimport torch\nimport sklearn.metrics\nfrom ultralytics import YOLO\nfrom tqdm.notebook import tqdm  # Use tqdm.notebook for Jupyter/Kaggle environments\n\nfrom contextlib import nullcontext\n\nimport threading\nfrom concurrent.futures import ThreadPoolExecutor","metadata":{"id":"mqKEYTQlZh3J","trusted":true,"execution":{"iopub.status.busy":"2025-04-25T05:56:54.356579Z","iopub.execute_input":"2025-04-25T05:56:54.356873Z","iopub.status.idle":"2025-04-25T05:56:59.372052Z","shell.execute_reply.started":"2025-04-25T05:56:54.356842Z","shell.execute_reply":"2025-04-25T05:56:59.371483Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Set random seed for reproducibility\nnp.random.seed(42)\ntorch.manual_seed(42)\nPATH = \"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025\"\n# Define paths\ntest_dir = os.path.join(PATH, \"test\")\nsubmission_path = \"/kaggle/working/submission.csv\"\n\n# Model path - adjust if your best model is saved in a different location\nmodel_path = \"/kaggle/input/mybest/best2.pt\"\n\n# Detection parameters\nCONFIDENCE_THRESHOLD = 0.45  # Lower threshold to catch more potential motors\nMAX_DETECTIONS_PER_TOMO = 5  # Keep track of top N detections per tomogram\nNMS_IOU_THRESHOLD = 0.2  # Non-maximum suppression threshold for 3D clustering\nCONCENTRATION = 1 # for fast submission","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-25T05:56:59.381391Z","iopub.execute_input":"2025-04-25T05:56:59.381660Z","iopub.status.idle":"2025-04-25T05:57:00.078109Z","shell.execute_reply.started":"2025-04-25T05:56:59.381637Z","shell.execute_reply":"2025-04-25T05:57:00.077487Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-25T05:57:00.079154Z","iopub.execute_input":"2025-04-25T05:57:00.079425Z","iopub.status.idle":"2025-04-25T05:57:00.172402Z","shell.execute_reply.started":"2025-04-25T05:57:00.079402Z","shell.execute_reply":"2025-04-25T05:57:00.171816Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def 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 = 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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-25T05:57:00.173193Z","iopub.execute_input":"2025-04-25T05:57:00.173486Z","iopub.status.idle":"2025-04-25T05:57:00.199150Z","shell.execute_reply.started":"2025-04-25T05:57:00.173430Z","shell.execute_reply":"2025-04-25T05:57:00.198352Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Run the submission pipeline\n\n# Time entire process\nstart_time = time.time()\n\n# Generate submission\nsubmission = generate_submission()\n\n# Print total execution time\nelapsed = time.time() - start_time\nprint(f\"\\nTotal execution time: {elapsed:.2f} seconds ({elapsed/60:.2f} minutes)\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-25T05:57:00.199885Z","iopub.execute_input":"2025-04-25T05:57:00.200130Z","execution_failed":"2025-04-25T05:58:15.909Z"}},"outputs":[],"execution_count":null}]}