{"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":"gpu","dataSources":[{"sourceId":91249,"databundleVersionId":11294684,"sourceType":"competition"},{"sourceId":237906738,"sourceType":"kernelVersion"}],"dockerImageVersionId":30919,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"from detectron2.engine import DefaultPredictor\nfrom detectron2.config import get_cfg\nfrom detectron2 import model_zoo\nfrom tqdm.notebook import tqdm\nimport matplotlib.patches as patches\nimport matplotlib.pyplot as plt\nimport pandas as pd\nimport numpy as np\nimport warnings\nimport random\nimport torch\nimport cv2\nimport os\n\nwarnings.filterwarnings(\"ignore\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class CFG:\n    dataset_path = \"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/\"\n    test_image_path = os.path.join(dataset_path, \"test\")\n    sample_sub_path = os.path.join(dataset_path, \"sample_submission.csv\")\n\n    checkpoint_path = \"/kaggle/input/byu-2025-faster-r-cnn-1-2-training/faster_rcnn_R_50_FPN_3x_fold_0/best_checkpoint.pth\"\n    model_name = \"COCO-Detection/faster_rcnn_R_50_FPN_3x.yaml\"\n\n    confidence_threshold = 0.8\n    nms_iou_threshold = 0.5\n    seed = 42","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"random.seed(CFG.seed)\nnp.random.seed(CFG.seed)\ntorch.manual_seed(CFG.seed)\ntorch.cuda.manual_seed_all(CFG.seed)\ntorch.backends.cudnn.deterministic = True\ntorch.backends.cudnn.benchmark = False","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"cfg = get_cfg()\ncfg.merge_from_file(model_zoo.get_config_file(CFG.model_name))\n\ncfg.MODEL.WEIGHTS = CFG.checkpoint_path\ncfg.MODEL.ROI_HEADS.NUM_CLASSES = 1","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def perform_3d_nms(detections, distance_threshold=CFG.nms_iou_threshold):\n    if not detections:\n        return []\n\n    detections = sorted(detections, key=lambda x: x['confidence'], reverse=True)\n    final_detections = []\n\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        \n    return final_detections","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def predict_motors_in_tomogram(tomo_id, predictor):\n    tomo_path = os.path.join(CFG.test_image_path, str(tomo_id))    \n    slices = sorted([f for f in os.listdir(tomo_path) if f.endswith('.jpg')])\n    if not slices:\n        return None, 0\n    \n    all_detections = []\n    for z, slice_file in enumerate(slices):\n        slice_path = os.path.join(tomo_path, slice_file)\n        img = cv2.imread(slice_path)\n        \n        outputs = predictor(img)\n        instances = outputs[\"instances\"].to(\"cpu\")\n        boxes = instances.pred_boxes.tensor.numpy() if instances.has(\"pred_boxes\") else []\n        scores = instances.scores.numpy() if instances.has(\"scores\") else []\n        \n        for box, score in zip(boxes, scores):\n            if score > CFG.confidence_threshold:\n                x_center = (box[0] + box[2]) / 2\n                y_center = (box[1] + box[3]) / 2\n                all_detections.append({\n                    'z': z,\n                    'y': y_center,\n                    'x': x_center,\n                    'confidence': score\n                })\n\n    all_detections = perform_3d_nms(all_detections)\n    if not all_detections:\n        return [-1, -1, -1], 0\n    \n    all_detections.sort(key=lambda x: x['confidence'], reverse=True)\n    best_detection = all_detections[0]\n    \n    motor_position = [\n        int(best_detection['z']),\n        int(best_detection['y']),\n        int(best_detection['x'])\n    ]\n    \n    return motor_position, float(best_detection['confidence'])","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def generate_predictions(test_dir, predictor):\n    predictions = []\n    \n    tomo_ids = [d for d in os.listdir(test_dir) if os.path.isdir(os.path.join(test_dir, d))]\n    for tomo_id in tqdm(tomo_ids):\n        motor_position, confidence = predict_motors_in_tomogram(tomo_id, predictor)        \n        predictions.append({\n            'tomo_id': tomo_id,\n            'Motor axis 0': motor_position[0],\n            'Motor axis 1': motor_position[1],\n            'Motor axis 2': motor_position[2],\n            'Confidence': confidence,\n            'Has motor': 1 if motor_position[0] != -1 else 0\n        })\n    \n    return pd.DataFrame(predictions)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"predictor = DefaultPredictor(cfg)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"%%time\ntest_preds = generate_predictions(CFG.test_image_path, predictor)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"sub = pd.read_csv(CFG.sample_sub_path)\nsub = test_preds[['tomo_id', 'Motor axis 0', 'Motor axis 1', 'Motor axis 2']]\nsub.to_csv(\"submission.csv\", index=False)\nsub.head()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if len(test_preds) == 3:\n    points = []\n    image_paths = []\n    \n    for i, r in test_preds.iterrows():\n        if (\n            r['Motor axis 0'] != -1 and \n            r['Motor axis 1'] != -1 and \n            r['Motor axis 2'] != -1\n        ):\n            slice_num = int(r['Motor axis 0'])\n            slice_str = f\"{slice_num:04d}\"\n            image_path = f\"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/test/{r.tomo_id}/slice_{slice_str}.jpg\"\n            image_paths.append(image_path)\n            points.append((r['Motor axis 2'], r['Motor axis 1'], r['Confidence']))\n\n    fig, axes = plt.subplots(1, 3, figsize=(15, 5))\n    box_size = 64\n    half_box = box_size // 2\n\n    for i, path in enumerate(image_paths):\n        img = cv2.imread(path)\n        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n        axes[i].imshow(img_rgb)\n        x, y, confidence = points[i]\n        \n        axes[i].scatter(x, y, color=\"red\")\n        \n        rect = patches.Rectangle((x - half_box, y - half_box), box_size, box_size, linewidth=2, edgecolor='lime', facecolor='none')\n        axes[i].add_patch(rect)\n        \n        axes[i].text(x - half_box, y - half_box - 10, f\"Confidence: {confidence:.4f}\", color='lime', fontsize=10, weight='bold', ha='center')\n        axes[i].set_title(path.split(\"/\")[-2] + \"/\" + path.split(\"/\")[-1])\n        axes[i].axis('off')\n\n    plt.tight_layout()\n    plt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}