{"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":[{"sourceType":"competition","sourceId":91249,"databundleVersionId":11294684},{"sourceType":"datasetVersion","sourceId":11181857,"datasetId":6821877,"databundleVersionId":11583985},{"sourceType":"modelInstanceVersion","sourceId":296173,"databundleVersionId":11519159,"modelInstanceId":253402}],"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-03-27T11:02:49.666362Z","iopub.execute_input":"2025-03-27T11:02:49.666634Z","iopub.status.idle":"2025-03-27T11:02:56.478330Z","shell.execute_reply.started":"2025-03-27T11:02:49.666605Z","shell.execute_reply":"2025-03-27T11:02:56.477464Z"},"_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/yolo-v11l/yolov9-s-960_v8_map-0.56.torchscript'\n\n# Define detection and processing parameters\nCONFIDENCE_THRESHOLD = 0.45\nMAX_DETECTIONS_PER_TOMO = 1\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-03-27T11:02:56.479673Z","iopub.execute_input":"2025-03-27T11:02:56.480286Z","iopub.status.idle":"2025-03-27T11:02:56.489150Z","shell.execute_reply.started":"2025-03-27T11:02:56.480250Z","shell.execute_reply":"2025-03-27T11:02:56.488321Z"}},"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-03-27T11:02:56.490686Z","iopub.execute_input":"2025-03-27T11:02:56.490957Z","iopub.status.idle":"2025-03-27T11:02:56.502515Z","shell.execute_reply.started":"2025-03-27T11:02:56.490938Z","shell.execute_reply":"2025-03-27T11:02:56.501739Z"}},"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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-27T11:02:56.503866Z","iopub.execute_input":"2025-03-27T11:02:56.504148Z","iopub.status.idle":"2025-03-27T11:02:56.521950Z","shell.execute_reply.started":"2025-03-27T11:02:56.504122Z","shell.execute_reply":"2025-03-27T11:02:56.521059Z"}},"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-27T11:02:56.522697Z","iopub.execute_input":"2025-03-27T11:02:56.522977Z","iopub.status.idle":"2025-03-27T11:02:56.540833Z","shell.execute_reply.started":"2025-03-27T11:02:56.522958Z","shell.execute_reply":"2025-03-27T11:02:56.540050Z"}},"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 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-03-27T11:02:56.541805Z","iopub.execute_input":"2025-03-27T11:02:56.542097Z","iopub.status.idle":"2025-03-27T11:02:56.561229Z","shell.execute_reply.started":"2025-03-27T11:02:56.542072Z","shell.execute_reply":"2025-03-27T11:02:56.560484Z"}},"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-27T11:02:56.561934Z","iopub.execute_input":"2025-03-27T11:02:56.562182Z","iopub.status.idle":"2025-03-27T11:02:56.588398Z","shell.execute_reply.started":"2025-03-27T11:02:56.562153Z","shell.execute_reply":"2025-03-27T11:02:56.587637Z"}},"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-03-27T11:02:56.589248Z","iopub.execute_input":"2025-03-27T11:02:56.589557Z","iopub.status.idle":"2025-03-27T11:03:41.052708Z","shell.execute_reply.started":"2025-03-27T11:02:56.589530Z","shell.execute_reply":"2025-03-27T11:03:41.051870Z"}},"outputs":[],"execution_count":null}]}