{"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":245461949,"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-train-01/edit","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-14T17:40:15.930324Z","iopub.execute_input":"2025-06-14T17:40:15.930590Z","iopub.status.idle":"2025-06-14T17:42:07.898263Z","shell.execute_reply.started":"2025-06-14T17:40:15.930569Z","shell.execute_reply":"2025-06-14T17:42:07.897243Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport random\n\nSEED = 42\nnp.random.seed(SEED)\nrandom.seed(SEED)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-14T17:42:07.899731Z","iopub.execute_input":"2025-06-14T17:42:07.900081Z","iopub.status.idle":"2025-06-14T17:42:07.904441Z","shell.execute_reply.started":"2025-06-14T17:42:07.900058Z","shell.execute_reply":"2025-06-14T17:42:07.903804Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"data_path = '/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/'\ntest_path = '/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/test'\nsample_submission_path = '/kaggle/working/submission.csv'\n\nmodel_path = '/kaggle/input/byu-yolo-train-01/yolo_weights/flagellar_motor_detector/weights/best.pt'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-14T17:42:07.904940Z","iopub.execute_input":"2025-06-14T17:42:07.905114Z","iopub.status.idle":"2025-06-14T17:42:07.917435Z","shell.execute_reply.started":"2025-06-14T17:42:07.905100Z","shell.execute_reply":"2025-06-14T17:42:07.916796Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def normalize_slice(slice_data):\n    \n    percentile_2th = np.percentile(slice_data, 2)\n    percentile_98th = np.percentile(slice_data, 98)\n    clipped_data = np.clip(slice_data, percentile_2th, percentile_98th)\n    normalized = 255 * (clipped_data - percentile_2th) / (percentile_98th - percentile_2th)\n    \n    return np.uint8(normalized)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-14T17:42:07.918669Z","iopub.execute_input":"2025-06-14T17:42:07.919204Z","iopub.status.idle":"2025-06-14T17:42:07.929426Z","shell.execute_reply.started":"2025-06-14T17:42:07.919187Z","shell.execute_reply":"2025-06-14T17:42:07.928818Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"BOX_SIZE = 24\nNMS_IOU_THRESHOLD = 0.2\nDISTANCE_THRESHOLD = BOX_SIZE * NMS_IOU_THRESHOLD\n\ndef perform_3d_nms(detections):\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(distance_1, distance_2):\n        return np.sqrt((distance_1['z'] - distance_2['z']) ** 2 + \n                       (distance_1['y'] - distance_2['y']) ** 2 + \n                       (distance_1['x'] - distance_2['x']) ** 2)\n    \n    while detections:\n        \n        best_detection = detections.pop(0)\n        final_detections.append(best_detection)\n        \n        detections = [distance for distance in detections if distance_3d(distance, best_detection) > DISTANCE_THRESHOLD]\n    \n    return final_detections","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-14T17:42:07.930130Z","iopub.execute_input":"2025-06-14T17:42:07.930296Z","iopub.status.idle":"2025-06-14T17:42:07.937802Z","shell.execute_reply.started":"2025-06-14T17:42:07.930282Z","shell.execute_reply":"2025-06-14T17:42:07.937102Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\n\nCONFIDENCE_THRESHOLD = 0.45\n\ndef generalization(tomo_id, model):\n    \n    tomo_id_path = os.path.join(test_path, tomo_id)\n    slice_jpg_files = sorted([f for f in os.listdir(tomo_id_path) if f.endswith('.jpg')])\n    \n    detections = []\n    \n    if len(slice_jpg_files) == 0:\n        pass\n\n    slice_jpg_files_paths = [os.path.join(tomo_id_path, slice_jpg_file) \\\n                             for slice_jpg_file in slice_jpg_files]\n    slice_jpg_files_nums = [int(slice_jpg_file.split('_')[1].split('.')[0]) \\\n                            for slice_jpg_file in slice_jpg_files]\n    \n    sub_results = model(slice_jpg_files_paths, verbose=False)\n        \n    for i, result in enumerate(sub_results):\n        if len(result.boxes) > 0:\n            boxes = result.boxes\n            for box_id, confidence in enumerate(boxes.conf):\n                if confidence >= CONFIDENCE_THRESHOLD:\n                    \n                    x_start, y_start, x_end, y_end = boxes.xyxy[box_id].cpu().numpy()\n                    \n                    x_center = (x_start + x_end) / 2\n                    y_center = (y_start + y_end) / 2\n                    \n                    detections.append({\n                        'z': round(slice_jpg_files_nums[i]),\n                        'y': round(y_center),\n                        'x': round(x_center),\n                        'confidence': float(confidence)\n                    })\n                            \n        \n    final_detections = perform_3d_nms(detections)\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\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-14T17:42:07.938535Z","iopub.execute_input":"2025-06-14T17:42:07.938714Z","iopub.status.idle":"2025-06-14T17:42:07.951009Z","shell.execute_reply.started":"2025-06-14T17:42:07.938700Z","shell.execute_reply":"2025-06-14T17:42:07.950415Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from ultralytics import YOLO\nimport pandas as pd\n\ntest_tomos = sorted(\n    [test_tomo_id for test_tomo_id in os.listdir(test_path) \\\n     if os.path.isdir(os.path.join(test_path, test_tomo_id))])\n    \nmodel = YOLO(model_path)\n\npredictions = []\n\nfor i, tomo_id in enumerate(test_tomos):\n    prediction = generalization(tomo_id, model)\n    predictions.append(prediction)\n    \nsubmission_df = pd.DataFrame(predictions)\nsubmission_df = submission_df[['tomo_id', 'Motor axis 0', 'Motor axis 1', 'Motor axis 2']]\nsubmission_df.to_csv(sample_submission_path, index=False)\n\nsubmission_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-06-14T17:42:07.951616Z","iopub.execute_input":"2025-06-14T17:42:07.951809Z","iopub.status.idle":"2025-06-14T17:43:16.932666Z","shell.execute_reply.started":"2025-06-14T17:42:07.951790Z","shell.execute_reply":"2025-06-14T17:43:16.932031Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}