{"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,"sourceType":"competition"}],"dockerImageVersionId":30919,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install ultralytics","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-18T00:51:47.733118Z","iopub.execute_input":"2025-03-18T00:51:47.733524Z","iopub.status.idle":"2025-03-18T00:51:53.208711Z","shell.execute_reply.started":"2025-03-18T00:51:47.733489Z","shell.execute_reply":"2025-03-18T00:51:53.207841Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Import & Setup\n\nimport os\nimport torch\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nfrom tqdm.auto import tqdm\nfrom ultralytics import YOLO\nimport torchvision.ops as ops\nimport cv2\nimport threading\nfrom concurrent.futures import ThreadPoolExecutor\n\n# Enable GPU acceleration\ndevice = 'cuda' if torch.cuda.is_available() else 'cpu'\ntorch.backends.cudnn.benchmark = True\ntorch.backends.cuda.matmul.allow_tf32 = True\n\n# Set confidence & IOU thresholds\nCONFIDENCE_THRESHOLD = 0.45\nNMS_IOU_THRESHOLD = 0.2\n\n# Define paths\nDATA_PATH = \"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/\"\nTEST_DIR = os.path.join(DATA_PATH, \"test\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-18T00:51:56.056025Z","iopub.execute_input":"2025-03-18T00:51:56.056435Z","iopub.status.idle":"2025-03-18T00:52:00.890990Z","shell.execute_reply.started":"2025-03-18T00:51:56.056401Z","shell.execute_reply":"2025-03-18T00:52:00.890304Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def preprocess_image(img_path):\n    \"\"\" Load and normalize image for inference \"\"\"\n    img = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)\n    if img is None:\n        img = np.array(Image.open(img_path))\n    return img\n\ndef load_images_in_parallel(image_paths):\n    \"\"\" Load batch of images in parallel for efficiency \"\"\"\n    images = []\n    with ThreadPoolExecutor() as executor:\n        images = list(executor.map(preprocess_image, image_paths))\n    return images","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-18T00:52:03.578082Z","iopub.execute_input":"2025-03-18T00:52:03.578549Z","iopub.status.idle":"2025-03-18T00:52:03.583238Z","shell.execute_reply.started":"2025-03-18T00:52:03.578525Z","shell.execute_reply":"2025-03-18T00:52:03.582280Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def run_inference(batch_images):\n    \"\"\" Run YOLO inference asynchronously on a batch of images \"\"\"\n    with torch.no_grad():\n        results = model(batch_images, verbose=False)\n    return results","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-18T00:52:05.712336Z","iopub.execute_input":"2025-03-18T00:52:05.712670Z","iopub.status.idle":"2025-03-18T00:52:05.716632Z","shell.execute_reply.started":"2025-03-18T00:52:05.712642Z","shell.execute_reply":"2025-03-18T00:52:05.715651Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def perform_3d_nms(detections, iou_threshold=0.2):\n    \"\"\" Efficient 3D NMS using PyTorch's native ops.nms() \"\"\"\n    if not detections:\n        return []\n    \n    detections = sorted(detections, key=lambda x: x['confidence'], reverse=True)\n    boxes = torch.tensor([[d['x'], d['y'], d['z'], d['confidence']] for d in detections])\n    scores = boxes[:, -1]\n    \n    keep_indices = ops.nms(boxes[:, :3], scores, iou_threshold)\n    return [detections[i] for i in keep_indices.tolist()]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-18T00:52:07.659415Z","iopub.execute_input":"2025-03-18T00:52:07.659700Z","iopub.status.idle":"2025-03-18T00:52:07.664526Z","shell.execute_reply.started":"2025-03-18T00:52:07.659679Z","shell.execute_reply":"2025-03-18T00:52:07.663658Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def process_tomogram(tomo_id):\n    \"\"\" Process a tomogram, detect motors, and return best results \"\"\"\n    slice_dir = os.path.join(TEST_DIR, tomo_id)\n    slice_files = sorted([f for f in os.listdir(slice_dir) if f.endswith('.jpg')])\n    \n    detections = []\n    for i in range(0, len(slice_files), 8):  # Dynamic batch size\n        batch_files = slice_files[i:i+8]\n        batch_paths = [os.path.join(slice_dir, f) for f in batch_files]\n        \n        batch_images = load_images_in_parallel(batch_paths)\n        results = run_inference(batch_images)\n        \n        for idx, result in enumerate(results):\n            for box in result.boxes:\n                detections.append({\n                    'z': i + idx,\n                    'x': (box.xyxy[0] + box.xyxy[2]) / 2,\n                    'y': (box.xyxy[1] + box.xyxy[3]) / 2,\n                    'confidence': box.conf\n                })\n    \n    detections = perform_3d_nms(detections, NMS_IOU_THRESHOLD)\n    \n    if not detections:\n        return {'tomo_id': tomo_id, 'Motor axis 0': -1, 'Motor axis 1': -1, 'Motor axis 2': -1}\n    \n    best_detection = detections[0]\n    return {\n        'tomo_id': tomo_id,\n        'Motor axis 0': int(best_detection['z']),\n        'Motor axis 1': int(best_detection['y']),\n        'Motor axis 2': int(best_detection['x'])\n    }","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-18T00:52:09.467264Z","iopub.execute_input":"2025-03-18T00:52:09.467546Z","iopub.status.idle":"2025-03-18T00:52:09.474076Z","shell.execute_reply.started":"2025-03-18T00:52:09.467524Z","shell.execute_reply":"2025-03-18T00:52:09.473140Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def generate_submission():\n    \"\"\" Run inference on all tomograms and generate submission file \"\"\"\n    test_tomos = sorted([d for d in os.listdir(TEST_DIR) if os.path.isdir(os.path.join(TEST_DIR, d))])\n    results = []\n    \n    with ThreadPoolExecutor(max_workers=4) as executor:\n        results = list(tqdm(executor.map(process_tomogram, test_tomos), total=len(test_tomos), desc=\"Processing Tomograms\"))\n    \n    submission_df = pd.DataFrame(results)\n    submission_df.to_csv(SUBMISSION_PATH, index=False)\n    print(f\"✅ Submission saved to {SUBMISSION_PATH}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-18T00:52:11.904425Z","iopub.execute_input":"2025-03-18T00:52:11.904708Z","iopub.status.idle":"2025-03-18T00:52:11.909721Z","shell.execute_reply.started":"2025-03-18T00:52:11.904687Z","shell.execute_reply":"2025-03-18T00:52:11.908860Z"}},"outputs":[],"execution_count":null}]}