{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.10.12"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":91249,"databundleVersionId":11294684,"sourceType":"competition"},{"sourceId":241170270,"sourceType":"kernelVersion"}],"dockerImageVersionId":30919,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install -U ultralytics albumentations onnxslim onnxruntime-gpu","metadata":{"trusted":true,"_kg_hide-output":true,"scrolled":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Imports and configs","metadata":{}},{"cell_type":"code","source":"from ultralytics.data.augment import Albumentations\nfrom ultralytics.utils import LOGGER, colorstr\nfrom ultralytics.utils.metrics import Metric\nfrom ultralytics import RTDETR, settings\nfrom functools import partial\nimport albumentations as A\nimport numpy as np\nimport warnings\nimport random\nimport shutil\nimport torch\nimport yaml\nimport json\nimport os\n\nwarnings.filterwarnings(\"ignore\")","metadata":{"_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","_kg_hide-output":true,"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class CFG:\n    dataset_path = \"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/\"\n    train_image_path = os.path.join(dataset_path, \"train\")\n    train_label_path = os.path.join(dataset_path, \"train_labels.csv\")\n\n    debug = False\n\n    seed = 42\n    n_fold = 5\n    current_fold = 0\n    \n    epochs = 1 if debug else 10\n    image_size = 640\n    es_patience = 7\n    device = \"0\"\n\n    yolo_dataset_path = f\"/temp/dataset/fold_{current_fold}\"\n    yolo_yaml_path = f\"/temp/dataset/fold_{current_fold}/dataset.yaml\"\n    yolo_model_name = \"rtdetr-l.pt\"\n\n    use_wandb = False\n    project = \"byu-locating-bacterial-flagellar-motors\"\n    name = yolo_model_name.split(\".\")[0] + f\"_fold_{current_fold}\"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"torch.manual_seed(CFG.seed)\nnp.random.seed(CFG.seed)\nrandom.seed(CFG.seed)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"settings.update({\n    \"runs_dir\": \"/temp/logs\", \n    \"tensorboard\": False\n})\n\nif CFG.use_wandb:\n    from kaggle_secrets import UserSecretsClient\n    import wandb\n    os.environ[\"WANDB_API_KEY\"] = UserSecretsClient().get_secret(\"WANDB_API_KEY\")\n    settings.update({\"wandb\": True})","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Data preprocessing","metadata":{}},{"cell_type":"code","source":"if os.path.exists(CFG.yolo_dataset_path):\n    shutil.rmtree(CFG.yolo_dataset_path)\n\nshutil.copytree(f\"/kaggle/input/byu-flagellar-motor-detection-1-preprocessing/dataset/fold_{CFG.current_fold}\", CFG.yolo_dataset_path)\nos.makedirs(CFG.project, exist_ok=True)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"yaml_data = {\n    \"path\": CFG.yolo_yaml_path,\n    \"train\": os.path.join(CFG.yolo_dataset_path, \"images\", \"train\"),\n    \"val\": os.path.join(CFG.yolo_dataset_path, \"images\", \"val\"),\n    \"names\": {0: \"motor\"}\n}\n\nwith open(CFG.yolo_yaml_path, \"w\") as f:\n    yaml.dump(yaml_data, f)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Training","metadata":{}},{"cell_type":"code","source":"def __init__(self, p=1.0):\n    self.p = p\n    self.transform = None\n    prefix = colorstr(\"albumentations: \")\n\n    try:\n        spatial_transforms = {\n            \"Affine\",\n            \"BBoxSafeRandomCrop\",\n            \"CenterCrop\",\n            \"CoarseDropout\",\n            \"Crop\",\n            \"CropAndPad\",\n            \"CropNonEmptyMaskIfExists\",\n            \"D4\",\n            \"ElasticTransform\",\n            \"Flip\",\n            \"GridDistortion\",\n            \"GridDropout\",\n            \"HorizontalFlip\",\n            \"Lambda\",\n            \"LongestMaxSize\",\n            \"MaskDropout\",\n            \"MixUp\",\n            \"Morphological\",\n            \"NoOp\",\n            \"OpticalDistortion\",\n            \"PadIfNeeded\",\n            \"Perspective\",\n            \"PiecewiseAffine\",\n            \"PixelDropout\",\n            \"RandomCrop\",\n            \"RandomCropFromBorders\",\n            \"RandomGridShuffle\",\n            \"RandomResizedCrop\",\n            \"RandomRotate90\",\n            \"RandomScale\",\n            \"RandomSizedBBoxSafeCrop\",\n            \"RandomSizedCrop\",\n            \"Resize\",\n            \"Rotate\",\n            \"SafeRotate\",\n            \"ShiftScaleRotate\",\n            \"SmallestMaxSize\",\n            \"Transpose\",\n            \"VerticalFlip\",\n            \"XYMasking\",\n        } \n        \n        T = [\n            A.HorizontalFlip(p=0.5),\n            A.RandomRotate90(p=0.5),\n            A.VerticalFlip(p=0.5)\n        ]\n\n        self.contains_spatial = any(transform.__class__.__name__ in spatial_transforms for transform in T)\n        \n        self.transform = (\n            A.Compose(T, bbox_params=A.BboxParams(format=\"yolo\", label_fields=[\"class_labels\"], check_each_transform=True))\n            if self.contains_spatial\n            else A.Compose(T)\n        )\n        \n        if hasattr(self.transform, \"set_random_seed\"):\n            self.transform.set_random_seed(torch.initial_seed())\n        \n        LOGGER.info(prefix + \", \".join(f\"{x}\".replace(\"always_apply=False, \", \"\") for x in T if x.p))\n        \n    except Exception as e:\n        LOGGER.info(f\"{prefix}{e}\")\n\nAlbumentations.__init__ = __init__","metadata":{"trusted":true,"_kg_hide-input":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def fitness(self):\n    return (5 * self.mp * self.mr) / (4 * self.mp + self.mr)\n\nMetric.fitness = fitness","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model = RTDETR(CFG.yolo_model_name)\n\nresults = model.train(\n    data=CFG.yolo_yaml_path,\n    epochs=CFG.epochs,\n    batch=6,\n    device=CFG.device,\n    imgsz=CFG.image_size,\n    optimizer='AdamW',\n    lr0=1e-4,\n    lrf=0.1,\n    warmup_epochs=0,\n    dropout=0.1,\n    project=CFG.project,\n    name=CFG.name,\n    exist_ok=True,\n    patience=CFG.es_patience,\n    save=True,\n    seed=CFG.seed,\n    val=True,\n    verbose=True\n)","metadata":{"trusted":true,"_kg_hide-output":true,"scrolled":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"results = model.val(verbose=False, save_json=True)\n\nprint(json.dumps({\n    \"ap\": float(results.box.ap[0]),\n    \"ap50\": float(results.box.ap50[0]),\n    \"f1\": float(results.box.f1[0]),\n    \"map\": float(results.box.map),\n    \"map50\": float(results.box.map50),\n    \"map75\": float(results.box.map75),\n    \"maps\": float(results.box.maps[0]),\n    \"mp\": float(results.box.mp),\n    \"mr\": float(results.box.mr),\n    \"p\": float(results.box.p[0]),\n    \"r\": float(results.box.r[0]),\n    \"f2\": (5 * float(results.box.mp) * float(results.box.mr)) / (4 * float(results.box.mp) + float(results.box.mr))\n}, indent=4))","metadata":{"trusted":true,"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"model.export(format='torchscript', imgsz=CFG.image_size, optimize=False, batch=8)","metadata":{"trusted":true,"_kg_hide-output":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Cleanup","metadata":{}},{"cell_type":"code","source":"shutil.rmtree(\"wandb\", ignore_errors=True)\ntry:\n    os.remove(\"yolo11n.pt\")\n    os.remove(CFG.yolo_model_name)\nexcept:\n    pass","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}