{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"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"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":91249,"databundleVersionId":11294684,"sourceType":"competition"},{"sourceId":224916709,"sourceType":"kernelVersion"},{"sourceId":247096157,"sourceType":"kernelVersion"}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"https://www.kaggle.com/code/mirenaborisova/byu-yolo-infer-04 --> version_9","metadata":{}},{"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-06-24T09:49:19.125677Z","iopub.execute_input":"2025-06-24T09:49:19.126186Z","iopub.status.idle":"2025-06-24T09:51:26.691289Z","shell.execute_reply.started":"2025-06-24T09:49:19.126162Z","shell.execute_reply":"2025-06-24T09:51:26.690244Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport torch\n\nSEED = 42\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\n\ndevice = 'cuda:0' if torch.cuda.is_available() else 'cpu'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-24T09:51:26.692829Z","iopub.execute_input":"2025-06-24T09:51:26.693093Z","iopub.status.idle":"2025-06-24T09:51:29.926831Z","shell.execute_reply.started":"2025-06-24T09:51:26.693068Z","shell.execute_reply":"2025-06-24T09:51:29.926038Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_dir = '/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/test'\nsubmission_path = '/kaggle/working/submission.csv'\nmodel_path = \\\n    '/kaggle/input/byu-yolo-train-02/yolo_weights/flagellar_motor_detector/weights/best.pt'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-24T09:51:29.927657Z","iopub.execute_input":"2025-06-24T09:51:29.928063Z","iopub.status.idle":"2025-06-24T09:51:29.931627Z","shell.execute_reply.started":"2025-06-24T09:51:29.928038Z","shell.execute_reply":"2025-06-24T09:51:29.930910Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nfrom PIL import Image\nimport cv2\nfrom ultralytics import YOLO\nimport threading\nfrom contextlib import nullcontext\nfrom concurrent.futures import ThreadPoolExecutor\n\nCONFIDENCE_THRESHOLD = 0.45\nNMS_IOU_THRESHOLD = 0.2\nBATCH_SIZE = 8\n    \ntorch.backends.cudnn.benchmark = True\ntorch.backends.cudnn.deterministic = False\ntorch.backends.cuda.matmul.allow_tf32 = True\ntorch.backends.cudnn.allow_tf32 = True\n\ngpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9\n\nfree_mem = gpu_mem - torch.cuda.memory_allocated(0) / 1e9\nBATCH_SIZE = max(8, min(32, int(free_mem * 4)))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-24T09:51:29.933295Z","iopub.execute_input":"2025-06-24T09:51:29.933895Z","iopub.status.idle":"2025-06-24T09:51:33.139346Z","shell.execute_reply.started":"2025-06-24T09:51:29.933869Z","shell.execute_reply":"2025-06-24T09:51:33.138750Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def normalize_slice(slice_data):\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    \n    images = []\n    for path in file_paths:\n        img = cv2.imread(path)\n        if img is None:\n            img = np.array(Image.open(path))\n        images.append(img)\n    return images\n\ndef perform_3d_nms(detections, iou_threshold):\n    \n    if not detections:\n        return []\n    \n    detections = sorted(detections, key=lambda x: x['confidence'], reverse=True)\n    \n    final_detections = []\n    \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    \n    box_size = 24\n    distance_threshold = box_size * iou_threshold\n    \n    while detections:\n        \n        best_detection = detections.pop(0)\n        final_detections.append(best_detection)\n        \n        detections = [d for d in detections if distance_3d(d, best_detection) > distance_threshold]\n    \n    return final_detections","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-24T09:51:33.140042Z","iopub.execute_input":"2025-06-24T09:51:33.140448Z","iopub.status.idle":"2025-06-24T09:51:33.147523Z","shell.execute_reply.started":"2025-06-24T09:51:33.140404Z","shell.execute_reply":"2025-06-24T09:51:33.146843Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def process_tomogram(tomo_id, model):\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    all_detections = []\n    \n    streams = [torch.cuda.Stream() for _ in range(min(4, BATCH_SIZE))]\n    \n    next_batch_thread = None\n    \n    for batch_start in range(0, len(slice_files), BATCH_SIZE):\n        \n        if next_batch_thread is not None:\n            next_batch_thread.join()\n            \n        batch_end = min(batch_start + BATCH_SIZE, len(slice_files))\n        batch_files = slice_files[batch_start:batch_end]\n        \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        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 else nullcontext():\n                \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                sub_results = model(sub_batch_paths, verbose=False)\n                    \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                                \n                                x1, y1, x2, y2 = boxes.xyxy[box_idx].cpu().numpy()\n                                \n                                x_center = (x1 + x2) / 2\n                                y_center = (y1 + y2) / 2\n                                \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        torch.cuda.synchronize()\n        \n    if next_batch_thread is not None:\n        next_batch_thread.join()\n        \n    final_detections = perform_3d_nms(all_detections, NMS_IOU_THRESHOLD)\n    \n    final_detections.sort(key=lambda x: x['confidence'], reverse=True)\n    \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    best_detection = final_detections[0]\n    \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    }","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-24T09:51:33.148289Z","iopub.execute_input":"2025-06-24T09:51:33.148553Z","iopub.status.idle":"2025-06-24T09:51:33.830267Z","shell.execute_reply.started":"2025-06-24T09:51:33.148531Z","shell.execute_reply":"2025-06-24T09:51:33.829396Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\n\ntest_tomos = sorted([d for d in os.listdir(test_dir) if os.path.isdir(os.path.join(test_dir, d))])\n\ntorch.cuda.empty_cache()\n\nmodel = YOLO(model_path)\nmodel.to(device)\nmodel.fuse()\n\nif torch.cuda.get_device_capability(0)[0] >= 7:\n    model.model.half()\n    \nresults = []\nmotors_found = 0\n\nwith ThreadPoolExecutor(max_workers=1) as executor:\n    future_to_tomo = {}\n    \n    for i, tomo_id in enumerate(test_tomos, 1):\n        future = executor.submit(process_tomogram, tomo_id, model)\n        future_to_tomo[future] = tomo_id\n        \n    for future in future_to_tomo:\n        tomo_id = future_to_tomo[future]\n        try:\n            torch.cuda.empty_cache()\n                \n            result = future.result()\n            results.append(result)\n        \n        except Exception as e:\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            \nsubmission_df = pd.DataFrame(results)\nsubmission_df = submission_df[['tomo_id', 'Motor axis 0', 'Motor axis 1', 'Motor axis 2']]\nsubmission_df.to_csv(submission_path, index=False)\n\nsubmission_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-24T09:51:33.831054Z","iopub.execute_input":"2025-06-24T09:51:33.831336Z","iopub.status.idle":"2025-06-24T09:52:46.112992Z","shell.execute_reply.started":"2025-06-24T09:51:33.831318Z","shell.execute_reply":"2025-06-24T09:52:46.112263Z"}},"outputs":[],"execution_count":null}]}