{"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"},{"sourceId":11193303,"sourceType":"datasetVersion","datasetId":6821877},{"sourceId":296173,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":253402,"modelId":274853}],"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"\n# Main notebook - https://www.kaggle.com/code/fautei/byd-yolo-optimized-submission-notebook\n## This is a example how you can utilize saved resources to make submission more complicated like my implementation of slicing aided hyper inference","metadata":{}},{"cell_type":"markdown","source":"# What is SAHI? I think this image better than any explanation\n![](https://raw.githubusercontent.com/obss/sahi/main/resources/sliced_inference.gif)","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-03-28T21:17:39.309205Z","iopub.execute_input":"2025-03-28T21:17:39.309512Z","iopub.status.idle":"2025-03-28T21:17:39.315318Z","shell.execute_reply.started":"2025-03-28T21:17:39.309488Z","shell.execute_reply":"2025-03-28T21:17:39.314419Z"},"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"torch.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-03-28T21:17:39.316363Z","iopub.execute_input":"2025-03-28T21:17:39.316668Z","iopub.status.idle":"2025-03-28T21:17:39.338959Z","shell.execute_reply.started":"2025-03-28T21:17:39.316622Z","shell.execute_reply":"2025-03-28T21:17:39.338322Z"}},"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/yolo-v11l/yolov9-s-960_v8_map-0.56.torchscript'\n\n# Define detection and processing parameters\nCONFIDENCE_THRESHOLD = 0.5\nMAX_DETECTIONS_PER_TOMO = 1\nDISTANCE_THRESHOLD = 60\nCONCENTRATION = 1 # Process a fraction of slices for fast submission\nSLICE_SIZE = (960, 960) #same as model input\nIMAGE_SIZE = (1024, 1024) #preprocess image to that size\n# Set device and dynamic batch size\ndevices = ['cuda:0','cuda:1']\nBATCH_SIZE = 16","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-28T21:17:39.340607Z","iopub.execute_input":"2025-03-28T21:17:39.340858Z","iopub.status.idle":"2025-03-28T21:17:39.355535Z","shell.execute_reply.started":"2025-03-28T21:17:39.340831Z","shell.execute_reply":"2025-03-28T21:17:39.354941Z"}},"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\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-28T21:17:39.356236Z","iopub.execute_input":"2025-03-28T21:17:39.356465Z","iopub.status.idle":"2025-03-28T21:17:39.363942Z","shell.execute_reply.started":"2025-03-28T21:17:39.356445Z","shell.execute_reply":"2025-03-28T21:17:39.363310Z"}},"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, color=(114, 114, 114), min_pad=20):\n    \"\"\"Resizes image with unchanged aspect ratio and pads it.\"\"\"\n    new_w, new_h = size\n    \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    dw, dh = new_w - new_unpad_w, new_h - new_unpad_h  # Padding\n    ratio = min(new_unpad_w / w, new_unpad_h / h )\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    \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 = normalize_slice(image)\n\n    image = cv2.cvtColor(image, cv2.COLOR_GRAY2BGR)\n    image = image.transpose(2, 0, 1)  # BGR to RGB, HWC to CHW\n    \n    # Normalize and convert to float32\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\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\ndef yolo_sahi_infer_impl(model, img_tensor, ratios_, paddings, size, conf_thres=0.5):\n    patches, paddings_w, paddings_h, ratios = [], [], [], []\n    img_num_to_quantity = {}\n\n    for idx, (image, r, pad_w, pad_h)  in enumerate(zip(img_tensor, ratios_, *paddings)):\n        _, h, w = image.shape\n        \n        x_slices = int(np.ceil(w / size[0]))\n        y_slices = int(np.ceil(h / size[1]))\n\n        unpad_x = int(np.ceil((x_slices * size[0] - w) /  max((x_slices - 1), 1)))\n        unpad_y = int(np.ceil((y_slices * size[1] - h) / max((y_slices - 1), 1)))\n\n        img_num_to_quantity[idx] = x_slices * y_slices\n\n        for i in range(y_slices):\n            for j in range(x_slices):\n                x_start, y_start = j * (size[0] - unpad_x), i * (size[1] - unpad_y)\n                x_end, y_end = min(x_start + size[0], w), min(y_start + size[1], h)\n\n                patch = image[:, y_start:y_end, x_start:x_end]\n                patches.append(patch)\n                ratios.append(r.item())\n                paddings_w.append(pad_w.item()-x_start)\n                paddings_h.append(pad_h.item()-y_start)\n\n    img_tensor = torch.stack(patches)\n    results = yolo_batch_infer_impl(model, img_tensor, ratios, (paddings_w, paddings_h), conf_thres)\n\n    results_final = []\n    current = 0\n    for i, taken in img_num_to_quantity.items():\n        res = tuple(np.concatenate(arrays) for arrays in zip(*results[current: current + taken]))\n        results_final.append(res)\n        current += taken\n    return results_final","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-28T21:17:39.386936Z","iopub.execute_input":"2025-03-28T21:17:39.387128Z","iopub.status.idle":"2025-03-28T21:17:39.401306Z","shell.execute_reply.started":"2025-03-28T21:17:39.387112Z","shell.execute_reply":"2025-03-28T21:17:39.400724Z"}},"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-03-28T21:17:39.411071Z","iopub.execute_input":"2025-03-28T21:17:39.411293Z","iopub.status.idle":"2025-03-28T21:17:39.420253Z","shell.execute_reply.started":"2025-03-28T21:17:39.411275Z","shell.execute_reply":"2025-03-28T21:17:39.419703Z"}},"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    \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 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, IMAGE_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\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_sahi_infer_impl(model, sub_images,sub_ratios,(sub_paddings_w,sub_paddings_h),SLICE_SIZE, 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-03-28T21:17:39.421308Z","iopub.execute_input":"2025-03-28T21:17:39.421560Z","iopub.status.idle":"2025-03-28T21:17:39.437902Z","shell.execute_reply.started":"2025-03-28T21:17:39.421541Z","shell.execute_reply":"2025-03-28T21:17:39.437143Z"}},"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-03-28T21:17:39.439080Z","iopub.execute_input":"2025-03-28T21:17:39.439344Z","iopub.status.idle":"2025-03-28T21:17:39.459232Z","shell.execute_reply.started":"2025-03-28T21:17:39.439325Z","shell.execute_reply":"2025-03-28T21:17:39.458664Z"}},"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":"\nstart_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)\")\nprint(\"\\nSubmission preview:\")\nsubmission.to_csv(submission_path, index=False)\nprint(submission.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-28T21:17:39.469279Z","iopub.execute_input":"2025-03-28T21:17:39.469470Z","iopub.status.idle":"2025-03-28T21:19:25.389982Z","shell.execute_reply.started":"2025-03-28T21:17:39.469453Z","shell.execute_reply":"2025-03-28T21:19:25.388959Z"}},"outputs":[],"execution_count":null}]}