{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.10.12"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":91249,"databundleVersionId":11294684,"sourceType":"competition"},{"sourceId":12017793,"sourceType":"datasetVersion","datasetId":7560908},{"sourceId":12566777,"sourceType":"datasetVersion","datasetId":7935925},{"sourceId":327336,"sourceType":"modelInstanceVersion","modelInstanceId":274744,"modelId":295634}],"dockerImageVersionId":30919,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true},"papermill":{"default_parameters":{},"duration":245.894411,"end_time":"2025-05-31T15:22:38.977300","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2025-05-31T15:18:33.082889","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"0c931b7d","cell_type":"code","source":"!cp -r /kaggle/input/mhafyolo/pytorch/default/1/MHAF-YOLO-main /kaggle/working/\nprint('PIP INSTALL OK!!!')\n\n\nimport os\nimport sys\nimport tarfile\nfrom pathlib import Path\n\n# Clear existing ultralytics from memory\nfor module in list(sys.modules.keys()):\n    if module.startswith('ultralytics'):\n        del sys.modules[module]\n\n# Extract ultralytics from dataset\ndataset_path = \"/kaggle/input/ultralytics-for-offline-install2\"\nworking_dir = Path(\"/kaggle/working\")\n\n# Find and extract archive\nfor file in os.listdir(dataset_path):\n    if file.endswith(('.tar.gz', '.zip', '.tar')):\n        archive_path = os.path.join(dataset_path, file)\n        extract_dir = working_dir / \"ultralytics_custom\"\n        \n        print(f\"Extracting {file}...\")\n        with tarfile.open(archive_path, 'r:*') as tar:\n            tar.extractall(extract_dir)\n        \n        # Find ultralytics directory\n        for item in extract_dir.rglob(\"*\"):\n            if item.is_dir() and \"ultralytics\" in item.name.lower():\n                ultralytics_path = item.parent if item.name == \"ultralytics\" else item\n                break\n        else:\n            ultralytics_path = next(extract_dir.iterdir())  # Use first directory\n        \n        # Add to Python path\n        sys.path.insert(0, str(ultralytics_path))\n        print(f\"✓ Custom ultralytics ready from {ultralytics_path}\")\n        break\n\n# Setup MHAF-YOLO\n!cp -r /kaggle/input/mhafyolo/pytorch/default/1/MHAF-YOLO-main /kaggle/working/\nmhaf_dir = Path(\"/kaggle/working/MHAF-YOLO-main\")\nsys.path.insert(1, str(mhaf_dir))\nos.chdir(mhaf_dir)\n\n# Test imports\ntry:\n    from ultralytics import YOLOv10\n    print(\"✅ YOLOv10 imported successfully!\")\nexcept ImportError:\n    from ultralytics import YOLO as YOLOv10  # Fallback\n    print(\"✅ Using YOLO as YOLOv10 fallback\")\n\nprint(\"Ready to proceed!\")\n\n\n\nimport os\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nimport torch\nimport cv2\nfrom tqdm.notebook import tqdm\nfrom pathlib import Path\nimport threading\nimport time\nfrom contextlib import nullcontext\nfrom concurrent.futures import ThreadPoolExecutor\nfrom ultralytics.utils.ops import non_max_suppression\n# NEW IMPORTS - Add these\nimport albumentations as A\nfrom sklearn.cluster import DBSCAN\nimport scipy.ndimage as ndi\nfrom scipy.cluster.hierarchy import fcluster, linkage\nfrom scipy.spatial.distance import pdist\n\ncurrent_dir = Path.cwd()\nprint(\"this_dir:\", current_dir)\n\ntarget_dir = Path(\"/kaggle/working/MHAF-YOLO-main\") \nos.chdir(target_dir)\n\nfrom ultralytics import YOLOv10","metadata":{"execution":{"iopub.status.busy":"2025-06-02T18:37:25.101954Z","iopub.execute_input":"2025-06-02T18:37:25.102309Z","iopub.status.idle":"2025-06-02T18:37:32.695618Z","shell.execute_reply.started":"2025-06-02T18:37:25.102281Z","shell.execute_reply":"2025-06-02T18:37:32.694537Z"},"papermill":{"duration":51.022123,"end_time":"2025-05-31T15:19:26.783494","exception":false,"start_time":"2025-05-31T15:18:35.761371","status":"completed"},"tags":[],"trusted":true,"jupyter":{"outputs_hidden":true},"collapsed":true},"outputs":[],"execution_count":null},{"id":"7546e755","cell_type":"code","source":"# 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\nmodel_path = \"/kaggle/input/mahf-yolo-train/mayolov2f.pt\"\n\n# Detection parameters\nCONFIDENCE_THRESHOLD = 0.78  # Lower threshold to catch more potential motors\nMAX_DETECTIONS_PER_TOMO = 1  # 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\nSIZE = 1024\nEARLY_STOPPING_TH = 0.01\nADAPTIVE_CONFIDENCE = False\nUSE_NOISE_AUGMENTATION = False      # Enable Gaussian noise\nUSE_BRIGHTNESS_AUGMENTATION = True # Enable brightness/contrast variations\nNOISE_FACTOR = 0.01               # Gaussian noise strength \nBRIGHTNESS_ALPHA_RANGE = [0.9, 1.1] # Brightness multiplier range\nCONTRAST_ALPHA_RANGE = [0.85, 1.15] # Contrast multiplier range\n\nprint(f\"TTA Configuration:\")\nprint(f\"- Noise augmentation: {USE_NOISE_AUGMENTATION}\")\nprint(f\"- Brightness augmentation: {USE_BRIGHTNESS_AUGMENTATION}\")\nprint(f\"- Noise factor: {NOISE_FACTOR}\")","metadata":{"execution":{"iopub.status.busy":"2025-06-02T18:37:32.697182Z","iopub.execute_input":"2025-06-02T18:37:32.697535Z","iopub.status.idle":"2025-06-02T18:37:32.706022Z","shell.execute_reply.started":"2025-06-02T18:37:32.697501Z","shell.execute_reply":"2025-06-02T18:37:32.705216Z"},"papermill":{"duration":0.014114,"end_time":"2025-05-31T15:19:26.800811","exception":false,"start_time":"2025-05-31T15:19:26.786697","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"id":"c4f7af2a","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'\n# device = ['cuda:0', 'cuda:1']\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 = 8  # Reduce batch size for CPU","metadata":{"execution":{"iopub.status.busy":"2025-06-02T18:37:32.707733Z","iopub.execute_input":"2025-06-02T18:37:32.707962Z","iopub.status.idle":"2025-06-02T18:37:32.718986Z","shell.execute_reply.started":"2025-06-02T18:37:32.707943Z","shell.execute_reply":"2025-06-02T18:37:32.718215Z"},"papermill":{"duration":0.082233,"end_time":"2025-05-31T15:19:26.885813","exception":false,"start_time":"2025-05-31T15:19:26.803580","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"id":"e534b074","cell_type":"code","source":"def quick_detection_check(model, image_np, device, img_size, threshold=0.1):\n    \"\"\"\n    Quick check to see if slice has any potential detections\n    \"\"\"\n    try:\n        results = model(image_np, imgsz=img_size, device=device, verbose=False)\n        \n        max_confidence = 0.0\n        for result in results:\n            if result.boxes and len(result.boxes) > 0:\n                confs = result.boxes.conf.cpu().numpy()\n                max_confidence = max(max_confidence, np.max(confs))\n        \n        return max_confidence >= threshold\n        \n    except Exception as e:\n        print(f\"Quick check failed: {e}\")\n        return True\n\ndef add_gaussian_noise(image, noise_factor=0.02):\n    \"\"\"\n    Add subtle Gaussian noise to image\n    \"\"\"\n    try:\n        # Convert to float for noise addition\n        img_float = image.astype(np.float32) / 255.0\n        \n        # Add Gaussian noise\n        noise = np.random.normal(0, noise_factor, img_float.shape)\n        noisy_img = img_float + noise\n        \n        # Clip to valid range and convert back\n        noisy_img = np.clip(noisy_img, 0, 1)\n        return (noisy_img * 255).astype(np.uint8)\n    except:\n        return image  # Return original if noise addition fails\n\ndef adjust_brightness_contrast(image, alpha=1.0, beta=0):\n    \"\"\"\n    Adjust brightness and contrast\n    alpha: contrast (1.0 = no change, >1.0 = more contrast, <1.0 = less contrast)\n    beta: brightness (0 = no change, positive = brighter, negative = darker)\n    \"\"\"\n    try:\n        return cv2.convertScaleAbs(image, alpha=alpha, beta=beta)\n    except:\n        return image\n\ndef get_adaptive_confidence_threshold(slice_number, total_slices, base_threshold=CONFIDENCE_THRESHOLD):\n    \"\"\"\n    Calculate adaptive confidence threshold based on slice position.\n    Motors are more likely in middle slices, so we can be more permissive at edges.\n    \n    Args:\n        slice_number: Current slice number (0-based)\n        total_slices: Total number of slices in tomogram\n        base_threshold: Base confidence threshold (default 0.8)\n    \n    Returns:\n        Adjusted confidence threshold\n    \"\"\"\n    if total_slices <= 1:\n        return base_threshold\n    \n    # Normalize slice position to 0-1 range\n    normalized_position = slice_number / (total_slices - 1)\n    \n    # Conservative approach: Only make edges slightly more permissive\n    if normalized_position <= 0.15 or normalized_position >= 0.85:\n        # Edge slices (first/last 15%) - slightly lower threshold\n        return base_threshold * 0.9  # 0.8 -> 0.72\n    elif normalized_position <= 0.25 or normalized_position >= 0.75:\n        # Near-edge slices (15-25% from edges) - slightly lower threshold  \n        return base_threshold * 0.95  # 0.8 -> 0.76\n    else:\n        # Middle slices (25-75%) - keep original threshold\n        return base_threshold  # Keep 0.8\n\n\n\ndef weighted_box_fusion(boxes, iou_threshold=0.4):\n    \"\"\"Applies Weighted Box Fusion to combine overlapping bounding boxes.\"\"\"\n    fused_boxes = []\n    # print(type(boxes))\n    used = [False] * len(boxes)\n\n    for i in range(len(boxes)):\n        if used[i]:\n            continue\n        \n        similar_boxes = [boxes[i]]\n        used[i] = True\n\n        for j in range(i + 1, len(boxes)):\n            if used[j]:\n                continue\n\n            if iou(boxes[i], boxes[j]) > iou_threshold:\n                similar_boxes.append(boxes[j])\n                used[j] = True\n\n        # Compute weighted average for the final box\n        # similar_boxes = np.array(similar_boxes)\n        # similar_boxes = similar_boxes.cpu().numpy()\n        similar_boxes = np.array([t.cpu() for t in similar_boxes])\n        # print(\"similar_boxes: \",similar_boxes)\n        confidences = similar_boxes[:,:, 4]\n        weights = confidences / confidences.sum()\n\n        fused_x1 = np.sum(similar_boxes[:, :, 0] * weights)\n        fused_y1 = np.sum(similar_boxes[:, :, 1] * weights)\n        fused_x2 = np.sum(similar_boxes[:, :, 2] * weights)\n        fused_y2 = np.sum(similar_boxes[:, :, 3] * weights)\n        fused_confidence = np.mean(confidences) #np.max(confidences)  # Take max confidence\n\n        fused_boxes.append((fused_x1, fused_y1, fused_x2, fused_y2, fused_confidence))\n\n    return fused_boxes\n\ndef predict_ensemble_tta(single_model, image_np, device, img_size):\n    \"\"\"\n    Enhanced TTA with noise and brightness/contrast variations\n    \"\"\"\n    all_boxes = []\n    all_confs = []\n    all_clss = []\n\n    def do_infer(img_tta, invert_func):\n        res = single_model(img_tta, \n                        imgsz=img_size, \n                        device=device, \n                        verbose=False)\n        for r in res:\n            boxes = r.boxes\n            if boxes is None or len(boxes)==0:\n                continue\n            xyxy = boxes.xyxy.cpu().numpy()\n            confs = boxes.conf.cpu().numpy()\n            clss = boxes.cls.cpu().numpy().astype(int)\n            # Inverse transformation\n            xyxy_orig = invert_func(xyxy)\n            all_boxes.append(xyxy_orig)\n            all_confs.append(confs)\n            all_clss.append(clss)\n\n    # Original image\n    do_infer(image_np, invert_func=lambda x: x)\n    org_img_h, org_img_w = image_np.shape[:2]\n    \n    # Horizontal flip\n    img_hflip = cv2.flip(image_np, 1)\n    def invert_hflip(xyxy):\n        new_ = xyxy.copy()\n        x1 = org_img_w - xyxy[:,2]\n        x2 = org_img_w - xyxy[:,0]\n        new_[:,0] = x1\n        new_[:,2] = x2\n        return new_\n    do_infer(img_hflip, invert_func=invert_hflip)\n\n    # Vertical flip\n    # img_vflip = cv2.flip(image_np, 0)\n    # def invert_vflip(xyxy):\n    #     new_ = xyxy.copy()\n    #     y1 = org_img_h - xyxy[:,3]\n    #     y2 = org_img_h - xyxy[:,1]\n    #     new_[:,1] = y1\n    #     new_[:,3] = y2\n    #     return new_\n    # do_infer(img_vflip, invert_func=invert_vflip)\n    \n    # NEW: Gaussian noise augmentation\n    img_noise = add_gaussian_noise(image_np, noise_factor=0.02)\n    do_infer(img_noise, invert_func=lambda x: x)\n    \n    # NEW: Brightness variations\n    img_bright = adjust_brightness_contrast(image_np, alpha=1.1, beta=10)\n    do_infer(img_bright, invert_func=lambda x: x)\n    \n    img_dark = adjust_brightness_contrast(image_np, alpha=0.9, beta=-10)\n    do_infer(img_dark, invert_func=lambda x: x)\n    \n    # NEW: Contrast variations  \n    img_high_contrast = adjust_brightness_contrast(image_np, alpha=1.15, beta=0)\n    do_infer(img_high_contrast, invert_func=lambda x: x)\n    \n    img_low_contrast = adjust_brightness_contrast(image_np, alpha=0.85, beta=0)\n    do_infer(img_low_contrast, invert_func=lambda x: x)\n\n    if len(all_boxes)==0:\n        return None\n\n    boxes_cat = np.concatenate(all_boxes, axis=0)\n    confs_cat = np.concatenate(all_confs, axis=0)\n    clss_cat  = np.concatenate(all_clss, axis=0)\n    cat_data = np.column_stack([boxes_cat, confs_cat, clss_cat])\n\n    # Apply weighted box fusion (your existing NMS)\n    cat_tensor = torch.from_numpy(cat_data).float().unsqueeze(0).to(device)\n    nms_out = weighted_box_fusion(cat_tensor, iou_threshold=0.5)\n    \n    if len(nms_out)==0 or nms_out[0] is None or len(nms_out[0])==0:\n        return None\n    \n    final_nms = list(nms_out[0])\n    return final_nms","metadata":{"execution":{"iopub.status.busy":"2025-06-02T18:37:32.720261Z","iopub.execute_input":"2025-06-02T18:37:32.720502Z","iopub.status.idle":"2025-06-02T18:37:32.741373Z","shell.execute_reply.started":"2025-06-02T18:37:32.720480Z","shell.execute_reply":"2025-06-02T18:37:32.740496Z"},"papermill":{"duration":0.017538,"end_time":"2025-05-31T15:19:26.906314","exception":false,"start_time":"2025-05-31T15:19:26.888776","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"id":"c23ff605","cell_type":"code","source":"def make_predict(sub_path, model, device, img_size=1024):\n    res_list = []\n    for img in sub_path:\n        img_np = cv2.imread(img)\n\n        if img_np is None:\n            res_list.append((0, 0, 0, 0, 0.0))\n            continue\n            \n        should_process = quick_detection_check(model, img_np, device, SIZE, threshold=EARLY_STOPPING_TH)\n        if not should_process:\n            continue\n        res_nms = predict_ensemble_tta(model, img_np, device, img_size=img_size)\n\n        if res_nms is None:\n            res_list.append((0, 0, 0, 0, 0.0))\n        else:\n            res_list.append(res_nms)\n\n    return res_list","metadata":{"execution":{"iopub.status.busy":"2025-06-02T18:37:32.742392Z","iopub.execute_input":"2025-06-02T18:37:32.742693Z","iopub.status.idle":"2025-06-02T18:37:32.758213Z","shell.execute_reply.started":"2025-06-02T18:37:32.742646Z","shell.execute_reply":"2025-06-02T18:37:32.757463Z"},"papermill":{"duration":0.008279,"end_time":"2025-05-31T15:19:26.917667","exception":false,"start_time":"2025-05-31T15:19:26.909388","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"id":"ec3d2321","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,SIZE=SIZE):\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(8, 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                    sub_results = make_predict(sub_batch_paths, model, device, SIZE)\n                    # print(sub_results)\n                \n                # Process each result in this sub-batch\n                # for result in sub_results:\n                    # print('nms_res1:', result)\n                # Get total slices for adaptive thresholding\n                total_original_slices = len(os.listdir(tomo_dir))\n                \n                for j,res in enumerate(sub_results):\n                    x1,y1,x2,y2, confidence = res\n                    \n                    # Use adaptive or fixed threshold\n                    if ADAPTIVE_CONFIDENCE:\n                        threshold_to_use = get_adaptive_confidence_threshold(\n                            sub_batch_slice_nums[j], total_original_slices, CONFIDENCE_THRESHOLD\n                        )\n                        if threshold_to_use < CONFIDENCE_THRESHOLD:\n                            print(f\"    Edge slice {sub_batch_slice_nums[j]}: using threshold {threshold_to_use:.2f}\")\n                    else:\n                        threshold_to_use = CONFIDENCE_THRESHOLD\n                    \n                    if confidence >= threshold_to_use:\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                    # print(\"all_detections:\", all_detections)\n\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 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 = YOLOv10(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=4) 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            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            # 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":{"execution":{"iopub.status.busy":"2025-06-02T18:37:32.759004Z","iopub.execute_input":"2025-06-02T18:37:32.759216Z","iopub.status.idle":"2025-06-02T18:37:32.781646Z","shell.execute_reply.started":"2025-06-02T18:37:32.759197Z","shell.execute_reply":"2025-06-02T18:37:32.780870Z"},"papermill":{"duration":0.024219,"end_time":"2025-05-31T15:19:26.944774","exception":false,"start_time":"2025-05-31T15:19:26.920555","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"id":"f92c8626","cell_type":"code","source":"# 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":{"execution":{"iopub.status.busy":"2025-06-02T18:37:32.782469Z","iopub.execute_input":"2025-06-02T18:37:32.782990Z","iopub.status.idle":"2025-06-02T18:51:20.206547Z","shell.execute_reply.started":"2025-06-02T18:37:32.782960Z","shell.execute_reply":"2025-06-02T18:51:20.205810Z"},"papermill":{"duration":190.147301,"end_time":"2025-05-31T15:22:37.094832","exception":false,"start_time":"2025-05-31T15:19:26.947531","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"id":"10210dd2","cell_type":"code","source":"# tomo_00e047\t169\t546\t603\n# tomo_01a877\t147\t638\t286","metadata":{"execution":{"iopub.status.busy":"2025-06-02T18:51:20.208254Z","iopub.execute_input":"2025-06-02T18:51:20.208512Z","iopub.status.idle":"2025-06-02T18:51:20.212118Z","shell.execute_reply.started":"2025-06-02T18:51:20.208489Z","shell.execute_reply":"2025-06-02T18:51:20.211128Z"},"papermill":{"duration":0.02357,"end_time":"2025-05-31T15:22:37.137390","exception":false,"start_time":"2025-05-31T15:22:37.113820","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null}]}