{"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"}],"dockerImageVersionId":30919,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install numpy scipy scikit-image","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-06T23:41:40.773227Z","iopub.execute_input":"2025-04-06T23:41:40.773535Z","iopub.status.idle":"2025-04-06T23:41:44.878212Z","shell.execute_reply.started":"2025-04-06T23:41:40.773499Z","shell.execute_reply":"2025-04-06T23:41:44.877142Z"},"collapsed":true,"jupyter":{"outputs_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport glob\nimport csv\nimport numpy as np\nimport torch\nfrom PIL import Image\nfrom tqdm import tqdm\nimport warnings\n\n# Suppress FutureWarnings (e.g., for torch.cuda.amp.autocast)\nwarnings.filterwarnings('ignore', category=FutureWarning)\n\n# --------------------------------------------------------------------\n# Load YOLO Model (using YOLOv5s pretrained on COCO)\n# --------------------------------------------------------------------\ndef load_yolo_model():\n    \"\"\"\n    Loads the YOLOv5s model using torch.hub.\n    This model is pretrained on COCO and serves as a placeholder.\n    \"\"\"\n    model = torch.hub.load('ultralytics/yolov5', 'yolov5s', pretrained=True)\n    # Move model to GPU if available\n    if torch.cuda.is_available():\n        model.to('cuda')\n    return model\n\n# --------------------------------------------------------------------\n# Process a Single 2D Slice\n# --------------------------------------------------------------------\ndef process_slice(image, model, conf_threshold=0.5):\n    \"\"\"\n    Processes a single 2D image slice using the YOLO model.\n    \n    Parameters:\n      image: PIL.Image instance (can be grayscale or RGB)\n      model: YOLO detection model\n      conf_threshold: confidence threshold to filter detections\n\n    Returns:\n      A list of detections; each detection is an array containing:\n      [xmin, ymin, xmax, ymax, confidence, class]\n    \"\"\"\n    # Convert to RGB if not already, since YOLO expects 3 channels.\n    if image.mode != 'RGB':\n        image = image.convert('RGB')\n    results = model(image)\n    # The results are stored in results.xyxy[0] as a tensor.\n    detections = results.xyxy[0].cpu().numpy()\n    # Filter out detections with confidence below the threshold.\n    detections = [det for det in detections if det[4] >= conf_threshold]\n    return detections\n\n# --------------------------------------------------------------------\n# Process an Entire Tomogram (Directory of Slices)\n# --------------------------------------------------------------------\ndef process_tomogram_yolo(tomo_dir, model):\n    \"\"\"\n    Processes all slices in a tomogram directory. For each slice, it runs object\n    detection using the YOLO model and tracks the detection with the highest\n    confidence. The best detection’s bounding box center is combined with the\n    slice index to form a 3D coordinate [z, y, x].\n    \n    Parameters:\n      tomo_dir: directory path containing JPEG slices\n      model: YOLO detection model\n      \n    Returns:\n      A list with the 3D coordinate or [-1, -1, -1] if no detection is found.\n    \"\"\"\n    slice_paths = sorted(glob.glob(os.path.join(tomo_dir, \"*.jpg\")))\n    best_detection = None\n    best_conf = -1\n    best_slice_index = -1\n\n    for i, slice_path in enumerate(slice_paths):\n        try:\n            image = Image.open(slice_path)\n        except Exception as e:\n            print(f\"Error loading {slice_path}: {e}\")\n            continue\n\n        detections = process_slice(image, model, conf_threshold=0.5)\n        for det in detections:\n            conf = det[4]\n            if conf > best_conf:\n                best_conf = conf\n                best_detection = det\n                best_slice_index = i\n\n    if best_detection is not None:\n        xmin, ymin, xmax, ymax, conf, cls = best_detection\n        # Compute the center of the bounding box.\n        x_center = (xmin + xmax) / 2\n        y_center = (ymin + ymax) / 2\n        return [best_slice_index, y_center, x_center]\n    else:\n        return [-1, -1, -1]\n\n# --------------------------------------------------------------------\n# Save Submission CSV File\n# --------------------------------------------------------------------\ndef save_submission(results, output_file='submission.csv'):\n    \"\"\"\n    Saves the list of results to a CSV file in the required submission format.\n    \n    Each row is: [tomo_id, Motor axis 0, Motor axis 1, Motor axis 2]\n    \"\"\"\n    with open(output_file, 'w', newline='') as csvfile:\n        writer = csv.writer(csvfile)\n        writer.writerow([\"tomo_id\", \"Motor axis 0\", \"Motor axis 1\", \"Motor axis 2\"])\n        writer.writerows(results)\n\n# --------------------------------------------------------------------\n# Main Function\n# --------------------------------------------------------------------\ndef main():\n    # Set the test directory path (ensure this path is correct in your Kaggle environment)\n    test_dir = \"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/test\"\n    \n    # Load the YOLO model (using YOLOv5s pretrained on COCO)\n    model = load_yolo_model()\n    results = []\n    \n    # Get a list of tomogram directories.\n    tomo_dirs = [d for d in os.listdir(test_dir) if os.path.isdir(os.path.join(test_dir, d))]\n    \n    # Process each tomogram with a progress bar.\n    for tomo_id in tqdm(tomo_dirs, desc=\"Processing tomograms\"):\n        tomo_path = os.path.join(test_dir, tomo_id)\n        coord = process_tomogram_yolo(tomo_path, model)\n        results.append([tomo_id] + coord)\n    \n    # Save the results to a CSV file.\n    save_submission(results)\n    print(\"Submission saved to submission.csv\")\n\nif __name__ == \"__main__\":\n    main()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-08T11:38:42.585624Z","iopub.execute_input":"2025-04-08T11:38:42.585874Z","iopub.status.idle":"2025-04-08T11:39:51.229237Z","shell.execute_reply.started":"2025-04-08T11:38:42.585852Z","shell.execute_reply":"2025-04-08T11:39:51.228282Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}