{"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":11257470,"sourceType":"datasetVersion","datasetId":7035485},{"sourceId":11257570,"sourceType":"datasetVersion","datasetId":7035582},{"sourceId":11467729,"sourceType":"datasetVersion","datasetId":7186405},{"sourceId":296173,"sourceType":"modelInstanceVersion","modelInstanceId":253402,"modelId":274853}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# This Notebook based on excelent competition host example but utilizes 2xT4 and all CPU resources to preload and normalize images as fast as we can\n## This notebook uses one of my models not higly optimized for LB score (you can try different confidence and model from public kernels just export it to torchscript)","metadata":{}},{"cell_type":"markdown","source":"# Flagellar Motor Detection in Bacteria Tomograms\n\nIn this notebook, we explore the dataset for detecting flagellar motors in bacterial tomograms.\n\n**We will**:\n- Set up our environment and reproducibility settings.\n- Load and inspect the training labels.\n- Examine key statistics including the distribution of motors per tomogram, tomogram sizes, and voxel spacing.\n- Visualize sample tomogram slices to get a better feel for the data.\n- \n**Let's begin!**\n  \n# 1) Library and dependecies","metadata":{}},{"cell_type":"code","source":"import plotly.express as px\nfrom PIL import Image, ImageDraw\nimport random\nimport seaborn as sns\nfrom matplotlib.patches import Rectangle\nimport yaml\nimport json\nimport os\nimport glob\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom tqdm.notebook import tqdm\nfrom sklearn.model_selection import train_test_split\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.optim.lr_scheduler import ReduceLROnPlateau\nimport cv2\nimport threading\nimport time\nfrom contextlib import nullcontext\nfrom concurrent.futures import ThreadPoolExecutor\nimport math","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-19T02:17:52.773346Z","iopub.execute_input":"2025-04-19T02:17:52.773676Z","iopub.status.idle":"2025-04-19T02:17:52.780120Z","shell.execute_reply.started":"2025-04-19T02:17:52.773646Z","shell.execute_reply":"2025-04-19T02:17:52.779282Z"},"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Export your models to torchscript and try","metadata":{}},{"cell_type":"code","source":"# Set random seed for reproducibility\nnp.random.seed(42)\ntorch.manual_seed(42)\n\n# Define paths for the test data and submission\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\nmodel_path ='/kaggle/input/lb0780-script-ljc/best.torchscript'\n\n# Define detection and processing parameters\n# CONFIDENCE_THRESHOLD = 0.58\n\nLOW_CONF_THRESH = 0.3      # 用于初步筛选和相邻切片验证\nFINAL_CONF_THRESH = 0.58    # 最终结果阈值 \nADJACENT_IOU_THRESH = 0.8  # 相邻切片IOU阈值\n\nDISTANCE_THRESHOLD = 60\nCONCENTRATION = 1 # Process a fraction of slices for fast submission\nSIZE = (960, 960)\n\n# Set device and dynamic batch size\ndevices = ['cuda:0','cuda:1']\nBATCH_SIZE = 32","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-19T02:17:52.781191Z","iopub.execute_input":"2025-04-19T02:17:52.781397Z","iopub.status.idle":"2025-04-19T02:17:52.802408Z","shell.execute_reply.started":"2025-04-19T02:17:52.781382Z","shell.execute_reply":"2025-04-19T02:17:52.801862Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class 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\ntorch.backends.cudnn.benchmark = True\ntorch.backends.cudnn.deterministic = False\ntorch.backends.cuda.matmul.allow_tf32 = True \ntorch.backends.cudnn.allow_tf32 = True","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-19T02:17:52.803057Z","iopub.execute_input":"2025-04-19T02:17:52.803282Z","iopub.status.idle":"2025-04-19T02:17:52.821137Z","shell.execute_reply.started":"2025-04-19T02:17:52.803266Z","shell.execute_reply":"2025-04-19T02:17:52.820533Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def normalize_slice(slice_data):\n    \"\"\"\n    Normalize slice data using the 2nd and 98th percentiles.\n    \"\"\"\n    p2, p98 = np.percentile(slice_data, [2, 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 letterbox(image, size=SIZE, color=(114, 114, 114), min_pad=0):\n    \"\"\"Resizes image with unchanged aspect ratio and pads it.\"\"\"\n    new_w, new_h = size\n    h,w = image.shape[:2]  # Current shape [height, width]\n    ratio = min(new_w / w, new_h / h )  # Scale ratio (new / old)\n    \n    new_unpad_w,new_unpad_h  = (int(round(w * ratio) - 2 * min_pad), int(round(h * ratio)) - 2 * min_pad)  # Rescaled width, height\n    ratio = min(new_unpad_w / w, new_unpad_h / h)\n    dw, dh = new_w - new_unpad_w, new_h - new_unpad_h  # Padding\n\n    dw /= 2  \n    dh /= 2\n\n    # Resize image\n    image_resized = cv2.resize(image, (new_unpad_w,new_unpad_h), interpolation=cv2.INTER_LINEAR)\n    # # Normalization is a very resource-intensive operation, so we do it on an image of minimal size\n    # image_resized = normalize_slice(image_resized)\n    # Padding\n    top, bottom = int(round(dh - 0.1)), int(round(dh + 0.1))\n    left, right = int(round(dw - 0.1)), int(round(dw + 0.1))\n    image_padded = cv2.copyMakeBorder(image_resized, top, bottom, left, right, cv2.BORDER_CONSTANT, value=color)\n\n    return image_padded, ratio, left, top\n\ndef preprocess_image(image, size):\n    image, ratio, pad_w, pad_h = letterbox(image, size)\n    image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)\n    image = image.transpose(2, 0, 1)  # BGR to RGB, HWC to CHW\n    image = np.ascontiguousarray(image, dtype=np.float32) / 255.0\n    img_tensor = torch.from_numpy(image)  # Add batch dimension\n    return img_tensor, ratio, pad_w, pad_h\n\ndef yolo_batch_infer_impl(model,img_tensor,ratios,paddings, conf_thres):\n    device = next(model.parameters()).device\n    img_tensor = img_tensor.to(device).half()\n    with torch.no_grad():\n        outputs = model(img_tensor)  # Shape [B, 5, 8400]\n        if isinstance(outputs, list): #yolov9\n            outputs = outputs[0]\n\n    outputs = outputs.cpu()\n    batch_results = []\n    for i, output in enumerate(outputs):  # Loop directly over model output\n        x_center, y_center, width, height, confidence = output\n\n        mask = confidence > conf_thres\n\n        x, y, w, h, conf = x_center[mask], y_center[mask], width[mask], height[mask], confidence[mask]\n        # Convert to original image coordinates\n        ratio, pad_w, pad_h = ratios[i], paddings[0][i], paddings[1][i]\n        x1, y1, x2, y2 = (x - w / 2 - pad_w) / ratio, (y - h / 2 - pad_h) / ratio, (x + w / 2 - pad_w) / ratio, (y + h / 2 - pad_h) / ratio\n\n        batch_results.append((np.stack((x1, y1, x2, y2), axis=1), conf.numpy()))\n\n    return batch_results\n\n# def yolo_batch_infer_impl(model, img_tensor, ratios, paddings):\n#     \"\"\"修改后的推理函数，不再进行confidence过滤\"\"\"\n#     device = next(model.parameters()).device\n#     img_tensor = img_tensor.to(device).half()\n#     with torch.no_grad():\n#         outputs = model(img_tensor)  # Shape [B, 5, 8400]\n#         if isinstance(outputs, list):  # yolov9\n#             outputs = outputs[0]\n\n#     outputs = outputs.cpu()\n#     batch_results = []\n#     for i, output in enumerate(outputs):  # Loop directly over model output\n#         x_center, y_center, width, height, confidence = output\n\n#         # 不再在此处过滤confidence，保留所有预测结果\n#         ratio, pad_w, pad_h = ratios[i], paddings[0][i], paddings[1][i]\n#         x1 = (x_center - width / 2 - pad_w) / ratio\n#         y1 = (y_center - height / 2 - pad_h) / ratio\n#         x2 = (x_center + width / 2 - pad_w) / ratio\n#         y2 = (y_center + height / 2 - pad_w) / ratio\n\n#         batch_results.append((np.stack((x1, y1, x2, y2), axis=1), confidence.numpy()))\n\n#     return batch_results","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-19T02:17:52.822625Z","iopub.execute_input":"2025-04-19T02:17:52.822894Z","iopub.status.idle":"2025-04-19T02:17:52.841339Z","shell.execute_reply.started":"2025-04-19T02:17:52.822871Z","shell.execute_reply":"2025-04-19T02:17:52.840720Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class MyDataset(torch.utils.data.Dataset):\n    def __init__(self,tomo_dir, files, size):\n        self.files = files\n        self.tomo_dir = tomo_dir\n        self.size = size\n\n    def __getitem__(self, index):\n        image_path = os.path.join(self.tomo_dir, self.files[index])\n        image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)\n        if image is None:\n            image = np.array(Image.open(image_path))\n        return (*preprocess_image(image, self.size), int(self.files[index].split('_')[1].split('.')[0]))\n        \n    def __len__(self):\n        return len(self.files)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-19T02:17:52.841931Z","iopub.execute_input":"2025-04-19T02:17:52.842107Z","iopub.status.idle":"2025-04-19T02:17:52.859904Z","shell.execute_reply.started":"2025-04-19T02:17:52.842093Z","shell.execute_reply":"2025-04-19T02:17:52.859326Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def perform_3d_nms(detections, distance_threshold):\n    \"\"\"\n    Perform 3D Non-Maximum Suppression on detections to merge nearby motors.\n    \"\"\"\n    if not detections:\n        return []\n    \n    detections = sorted(detections, key=lambda x: x['confidence'], reverse=True)\n    final_detections = []\n    def distance_3d(d1, d2):\n        return np.sqrt((d1['z'] - d2['z'])**2 + (d1['y'] - d2['y'])**2 + (d1['x'] - d2['x'])**2)\n\n    while detections:\n        best_detection = detections.pop(0)\n\n        final_detections.append(best_detection)\n\n        detections = [d for d in detections if distance_3d(d, best_detection) > distance_threshold]\n    return final_detections\n\n\ndef calculate_iou(box1, box2):\n    \"\"\"计算两个2D框的IOU\"\"\"\n    x1 = max(box1[0], box2[0])\n    y1 = max(box1[1], box2[1])\n    x2 = min(box1[2], box2[2])\n    y2 = min(box1[3], box2[3])\n    \n    inter_area = max(0, x2 - x1) * max(0, y2 - y1)\n    box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1])\n    box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1])\n    \n    return inter_area / (box1_area + box2_area - inter_area)\n\n    \ndef process_tomogram(tomo_id, model, index=0, total=1, device=''):\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    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    all_detections = []\n    \n    if device.startswith('cuda'):\n        streams = [torch.cuda.Stream(device=device) for _ in range(min(4, BATCH_SIZE))]\n    else:\n        streams = [None]\n\n    dataset = MyDataset(tomo_dir, slice_files, SIZE)\n    dataloader = torch.utils.data.DataLoader(dataset, batch_size=BATCH_SIZE, num_workers = 2)\n    \n    # 用于存储所有检测结果\n    all_detections = []\n    detection_dict = {}  # 按z索引组织的检测 {z: [(box, conf), ...]}\n    \n    for batch in tqdm(dataloader, desc=f\"Processing tomogram {tomo_id} ({index}/{total}) | {len(slice_files)} out of {len(os.listdir(tomo_dir))} slices (CONCENTRATION={CONCENTRATION}): \"):\n\n        sub_batch_results = []\n        \n        images,ratios,paddings_w,paddings_h,indexes = batch\n        sub_batches = (torch.tensor_split(images, len(streams)),\n                        torch.tensor_split(ratios, len(streams)),\n                        torch.tensor_split(paddings_w, len(streams)),\n                        torch.tensor_split(paddings_h, len(streams)),\n                        torch.tensor_split(indexes, len(streams))\n                      )\n        for i, (sub_images,sub_ratios,sub_paddings_w,sub_paddings_h,sub_indexes) in enumerate(zip(*sub_batches)):\n            if len(sub_images) == 0:\n                continue\n            stream = streams[i % len(streams)]\n            with torch.cuda.stream(stream) if stream and device.startswith('cuda') else nullcontext():\n                with GPUProfiler(f\"Inference batch {i+1}/{len(dataloader)}\"):\n                    sub_results = yolo_batch_infer_impl(model, sub_images,sub_ratios,(sub_paddings_w,sub_paddings_h),LOW_CONF_THRESH)\n\n                for j, (boxes, confs) in enumerate(sub_results):\n                    z_index = sub_indexes[j].item()\n                    if z_index not in detection_dict:\n                        detection_dict[z_index] = []\n                    \n                    for box, conf in zip(boxes, confs):\n                        x1, y1, x2, y2 = box\n                        detection = {\n                            'z': z_index,\n                            'y': (y1 + y2) / 2,\n                            'x': (x1 + x2) / 2,\n                            'box': (x1, y1, x2, y2),\n                            'confidence': float(conf)\n                        }\n                        all_detections.append(detection)\n                        detection_dict[z_index].append(detection)\n        # Synchronize streams\n        if device.startswith('cuda'):\n            torch.cuda.synchronize()\n\n    if torch.cuda.is_available():\n        torch.cuda.empty_cache()\n\n    # # 第一步：先用较低的初始阈值(LOW_CONF_THRESH)保留更多候选检测\n    if not all_detections:\n        return {'tomo_id': tomo_id, 'Motor axis 0': -1, 'Motor axis 1': -1, 'Motor axis 2': -1}\n    \n    # 第二步：找到最高confidence的检测(在初始过滤结果中)\n    best_detection = max(all_detections, key=lambda x: x['confidence'])\n    best_z = best_detection['z']\n    \n    # 第三步：增强型相邻切片验证\n    adjacent_hits = {'prev': False, 'next': False}  # 记录上下帧命中状态\n    adjacent_exists = {'prev': False, 'next': False}  # 记录上下帧是否存在任何检测框\n    \n    # 检查上一帧\n    prev_z = best_z - 1\n    if prev_z in detection_dict:\n        adjacent_exists['prev'] = True\n        for adj_det in detection_dict[prev_z]:\n            if adj_det['confidence'] > LOW_CONF_THRESH:\n                iou = calculate_iou(best_detection['box'], adj_det['box'])\n                if iou > ADJACENT_IOU_THRESH:\n                    adjacent_hits['prev'] = True\n                    break\n    \n    # 检查下一帧\n    next_z = best_z + 1\n    if next_z in detection_dict:\n        adjacent_exists['next'] = True\n        for adj_det in detection_dict[next_z]:\n            if adj_det['confidence'] > LOW_CONF_THRESH:\n                iou = calculate_iou(best_detection['box'], adj_det['box'])\n                if iou > ADJACENT_IOU_THRESH:\n                    adjacent_hits['next'] = True\n                    break\n    \n    # 情况1：任意相邻帧都没有检测框 → 直接无效\n    if not adjacent_exists['prev'] and not adjacent_exists['next']:\n        return {'tomo_id': tomo_id, 'Motor axis 0': -1, 'Motor axis 1': -1, 'Motor axis 2': -1}\n    \n    # 情况2：上下帧都有有效检测框 → 增强置信度\n    if adjacent_hits['prev'] and adjacent_hits['next']:\n        best_detection['confidence'] += 0.3\n        # 确保不超过最大可能值\n        best_detection['confidence'] = min(best_detection['confidence'], 1.0)\n    \n    # 情况3：没有任何一帧有有效检测框 → 无效\n    elif not adjacent_hits['prev'] and not adjacent_hits['next']:\n        return {'tomo_id': tomo_id, 'Motor axis 0': -1, 'Motor axis 1': -1, 'Motor axis 2': -1}\n    \n    # 第四步：现在用更高的正式阈值(FINAL_CONF_THRESH)过滤\n    final_candidates = [d for d in all_detections if d['confidence'] > FINAL_CONF_THRESH]\n    \n    # 第五步：执行3D NMS\n    final_detections = perform_3d_nms(final_candidates, DISTANCE_THRESHOLD)\n    \n    if not final_detections:\n        return {'tomo_id': tomo_id, 'Motor axis 0': -1, 'Motor axis 1': -1, 'Motor axis 2': -1}\n    \n    # 返回最高confidence的结果\n    final_detections.sort(key=lambda x: x['confidence'], reverse=True)\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    }\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-19T02:17:52.930349Z","iopub.execute_input":"2025-04-19T02:17:52.930559Z","iopub.status.idle":"2025-04-19T02:17:52.949502Z","shell.execute_reply.started":"2025-04-19T02:17:52.930544Z","shell.execute_reply":"2025-04-19T02:17:52.948646Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# def perform_3d_nms(detections, distance_threshold):\n#     \"\"\"\n#     Perform 3D Non-Maximum Suppression on detections to merge nearby motors.\n#     \"\"\"\n#     if not detections:\n#         return []\n    \n#     detections = sorted(detections, key=lambda x: x['confidence'], reverse=True)\n#     final_detections = []\n#     def distance_3d(d1, d2):\n#         return np.sqrt((d1['z'] - d2['z'])**2 + (d1['y'] - d2['y'])**2 + (d1['x'] - d2['x'])**2)\n\n#     while detections:\n#         best_detection = detections.pop(0)\n\n#         final_detections.append(best_detection)\n\n#         detections = [d for d in detections if distance_3d(d, best_detection) > distance_threshold]\n#     return final_detections\n\n        \n# def process_tomogram(tomo_id, model, index=0, total=1, device=''):\n#     \"\"\"\n#     Process a single tomogram and return the most confident motor detection.\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#     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#     all_detections = []\n    \n#     if device.startswith('cuda'):\n#         streams = [torch.cuda.Stream(device=device) for _ in range(min(4, BATCH_SIZE))]\n#     else:\n#         streams = [None]\n\n#     dataset = MyDataset(tomo_dir, slice_files, SIZE)\n#     dataloader = torch.utils.data.DataLoader(dataset, batch_size=BATCH_SIZE, num_workers = 2)\n    \n#     for batch in tqdm(dataloader, desc=f\"Processing tomogram {tomo_id} ({index}/{total}) | {len(slice_files)} out of {len(os.listdir(tomo_dir))} slices (CONCENTRATION={CONCENTRATION}): \"):\n\n#         sub_batch_results = []\n        \n#         images,ratios,paddings_w,paddings_h,indexes = batch\n#         sub_batches = (torch.tensor_split(images, len(streams)),\n#                         torch.tensor_split(ratios, len(streams)),\n#                         torch.tensor_split(paddings_w, len(streams)),\n#                         torch.tensor_split(paddings_h, len(streams)),\n#                         torch.tensor_split(indexes, len(streams))\n#                       )\n\n#         for i, (sub_images,sub_ratios,sub_paddings_w,sub_paddings_h,sub_indexes) in enumerate(zip(*sub_batches)):\n#             if len(sub_images) == 0:\n#                 continue\n#             stream = streams[i % len(streams)]\n#             with torch.cuda.stream(stream) if stream and device.startswith('cuda') else nullcontext():\n#                 with GPUProfiler(f\"Inference batch {i+1}/{len(dataloader)}\"):\n#                     sub_results = yolo_batch_infer_impl(model, sub_images,sub_ratios,(sub_paddings_w,sub_paddings_h),CONFIDENCE_THRESHOLD)\n                \n#                 # Process each result in this sub-batch\n#                 for j, (boxes, confs) in enumerate(sub_results):\n#                     for box, conf in zip(boxes, confs):\n#                         x1, y1, x2, y2 = box\n#                         x_center = (x1 + x2) / 2\n#                         y_center = (y1 + y2) / 2\n#                         all_detections.append({\n#                             'z': sub_indexes[j].item(),\n#                             'y': round(y_center),\n#                             'x': round(x_center),\n#                             'confidence': float(conf)\n#                         })\n        \n#         # Synchronize streams\n#         if device.startswith('cuda'):\n#             torch.cuda.synchronize()\n\n#     if torch.cuda.is_available():\n#         torch.cuda.empty_cache()\n        \n#     final_detections = perform_3d_nms(all_detections, DISTANCE_THRESHOLD)\n    \n#     if not final_detections:\n#         return {'tomo_id': tomo_id, 'Motor axis 0': -1, 'Motor axis 1': -1, 'Motor axis 2': -1}\n\n#     final_detections.sort(key=lambda x: x['confidence'], reverse=True)\n#     best_detection = final_detections[0]\n\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#     }\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-19T02:17:52.950472Z","iopub.execute_input":"2025-04-19T02:17:52.950677Z","iopub.status.idle":"2025-04-19T02:17:52.967620Z","shell.execute_reply.started":"2025-04-19T02:17:52.950662Z","shell.execute_reply":"2025-04-19T02:17:52.966894Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 4.3) Submission Generation\n\nThe function below processes each tomogram in the test directory using parallel processing, applies 3D NMS to merge detections, and then creates a CSV submission file with the predicted motor coordinates.","metadata":{}},{"cell_type":"code","source":"from joblib import Parallel, delayed\n\ndef generate_submission(test_tomos):\n    \"\"\"\n    Main function to generate the submission file.\n    \"\"\"\n\n    total_tomos = len(test_tomos)\n    print(f\"Found {total_tomos} tomograms in test directory\")\n    \n    if torch.cuda.is_available():\n        torch.cuda.empty_cache()\n    \n    print(f\"Loading YOLO model from {model_path}\")\n    models = [torch.jit.load(model_path, map_location=device).eval().half() for device in devices]\n    \n    results = []\n    motors_found = 0\n    future_to_tomo = {}\n    \n    with ThreadPoolExecutor(max_workers=4) as executor:\n        \n        for i, tomo_id in enumerate(test_tomos):\n            future = executor.submit(process_tomogram, tomo_id, models[i%2], i, total_tomos, devices[i%2])\n            future_to_tomo[future] = tomo_id\n            \n    for future,tomo_id in future_to_tomo.items():\n        try:\n            result = future.result()\n            results.append(result)\n        except Exception as e:\n            print(e)\n            results.append({'tomo_id': tomo_id, 'Motor axis 0': -1, 'Motor axis 1': -1, 'Motor axis 2': -1})\n    \n    submission_df = pd.DataFrame(results)\n    submission_df = submission_df[['tomo_id', 'Motor axis 0', 'Motor axis 1', 'Motor axis 2']]\n\n    return submission_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-19T02:17:52.968914Z","iopub.execute_input":"2025-04-19T02:17:52.969176Z","iopub.status.idle":"2025-04-19T02:17:52.985295Z","shell.execute_reply.started":"2025-04-19T02:17:52.969161Z","shell.execute_reply":"2025-04-19T02:17:52.984666Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## 4.4) Run the Submission Pipeline\n\nFinally, we time the entire inference process and generate the submission file.","metadata":{}},{"cell_type":"code","source":"start_time = time.time()\n\ntest_tomos = sorted([d for d in os.listdir(test_dir) if os.path.isdir(os.path.join(test_dir, d))])\n\nsubmission = generate_submission(test_tomos)\nelapsed = time.time() - start_time\nprint(f\"\\nTotal execution time: {elapsed:.2f} seconds ({elapsed/60:.2f} minutes)\")\n\nprint(\"\\nSubmission preview:\")\nsubmission.to_csv(submission_path, index=False)\nprint(submission.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-19T02:17:52.985978Z","iopub.execute_input":"2025-04-19T02:17:52.986215Z","iopub.status.idle":"2025-04-19T02:18:42.432359Z","shell.execute_reply.started":"2025-04-19T02:17:52.986196Z","shell.execute_reply":"2025-04-19T02:18:42.431531Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}