{"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":"none","dataSources":[{"sourceId":91249,"databundleVersionId":11294684,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":11924982,"sourceType":"datasetVersion","datasetId":7497431},{"sourceId":211097053,"sourceType":"kernelVersion"},{"sourceId":408274,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":333563,"modelId":354565},{"sourceId":408399,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":333670,"modelId":354668},{"sourceId":409203,"sourceType":"modelInstanceVersion","isSourceIdPinned":false,"modelInstanceId":334345,"modelId":355367}],"dockerImageVersionId":31040,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Standard libraries\nimport os\nimport json\nimport random\nimport threading\nimport time\nfrom pathlib import Path\nfrom contextlib import nullcontext\nfrom concurrent.futures import ThreadPoolExecutor\n\n# Numerical and data handling\nimport numpy as np\nimport pandas as pd\nimport yaml\n\n# Image handling and visualization\nfrom PIL import Image, ImageDraw\nimport cv2\nimport matplotlib.pyplot as plt\nfrom matplotlib.patches import Rectangle\n\n# PyTorch and torchvision\nimport torch\nimport torch.nn as nn\nfrom torchvision import transforms\nfrom torch.utils.data import Dataset\n\n# Progress bar\nfrom tqdm.notebook import tqdm","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:23:04.568984Z","iopub.execute_input":"2025-05-24T04:23:04.569165Z","iopub.status.idle":"2025-05-24T04:23:17.943305Z","shell.execute_reply.started":"2025-05-24T04:23:04.569148Z","shell.execute_reply":"2025-05-24T04:23:17.942722Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# !pip install ultralytics","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:23:17.944600Z","iopub.execute_input":"2025-05-24T04:23:17.944991Z","iopub.status.idle":"2025-05-24T04:23:17.948359Z","shell.execute_reply.started":"2025-05-24T04:23:17.944972Z","shell.execute_reply":"2025-05-24T04:23:17.947761Z"}},"outputs":[],"execution_count":null},{"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-05-24T04:23:17.950120Z","iopub.execute_input":"2025-05-24T04:23:17.950785Z","iopub.status.idle":"2025-05-24T04:25:05.359311Z","shell.execute_reply.started":"2025-05-24T04:23:17.950757Z","shell.execute_reply":"2025-05-24T04:25:05.358159Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# YOLO model (Ultralytics)\nfrom ultralytics import YOLO","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:25:05.360596Z","iopub.execute_input":"2025-05-24T04:25:05.361066Z","iopub.status.idle":"2025-05-24T04:25:07.612545Z","shell.execute_reply.started":"2025-05-24T04:25:05.361028Z","shell.execute_reply":"2025-05-24T04:25:07.611701Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"np.random.seed(42)\nrandom.seed(42)\ntorch.manual_seed(42)\ndevice = 'cuda' if torch.cuda.is_available() else 'cpu'\ndevice","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:25:07.613420Z","iopub.execute_input":"2025-05-24T04:25:07.613666Z","iopub.status.idle":"2025-05-24T04:25:07.630783Z","shell.execute_reply.started":"2025-05-24T04:25:07.613641Z","shell.execute_reply":"2025-05-24T04:25:07.630042Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"root_dir = Path(\"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:25:07.631524Z","iopub.execute_input":"2025-05-24T04:25:07.631768Z","iopub.status.idle":"2025-05-24T04:25:09.933603Z","shell.execute_reply.started":"2025-05-24T04:25:07.631750Z","shell.execute_reply":"2025-05-24T04:25:09.932667Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# yaml_dir = \"/kaggle/working/yaml_dir\"\n# os.makedirs(yaml_dir, exist_ok=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:25:09.934592Z","iopub.execute_input":"2025-05-24T04:25:09.934903Z","iopub.status.idle":"2025-05-24T04:25:09.948591Z","shell.execute_reply.started":"2025-05-24T04:25:09.934875Z","shell.execute_reply":"2025-05-24T04:25:09.947820Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class Visualise:\n    def __init__(self, path, transform=None, box_width=24):\n        self.path = Path(path)\n        self.transform = transform\n        self.box_width = box_width\n\n    def random_tomosplits(self, n):\n        image_path = list(self.path.glob(\"*/*.jpg\"))\n        if n > len(image_path):\n            raise ValueError(f\"Requested {n} samples but only {len(image_path)} found.\")\n        image_list = random.sample(image_path, n)\n\n        rows = (n + 4) // 5\n        fig, axes = plt.subplots(rows, 5, figsize=(16, rows * 4))\n        axes = axes.flatten()\n\n        for i, img_path in enumerate(image_list):\n            image = Image.open(img_path)\n            img_size = image.size\n            axes[i].imshow(image)\n            axes[i].axis(\"off\")\n\n            tomo_id = img_path.parent.name.split('_')[1]\n            slice_id = img_path.stem.split('_')[1]\n            title = f\"tomo_id : {tomo_id}\\nslice : {slice_id}\\nsize : {img_size}\"\n            axes[i].set_title(title, fontsize=12)\n\n        for j in range(i + 1, len(axes)):\n            axes[j].axis('off')\n\n        plt.tight_layout()\n        plt.show()\n\n    def display_transform(self, n):\n        image_path = list(self.path.glob(\"*/*.jpg\"))\n        if n > len(image_path):\n            raise ValueError(f\"Requested {n} samples but only {len(image_path)} found.\")\n        image_list = random.sample(image_path, n)\n\n        fig, axes = plt.subplots(n, 2, figsize=(10, 4 * n))\n        for i, img_path in enumerate(image_list):\n            image = Image.open(img_path)\n            axes[i][0].imshow(image)\n            axes[i][0].axis(\"off\")\n            axes[i][0].set_title(f\"Original\\nsize: {image.size}\", fontsize=12)\n\n            if self.transform is not None:\n                t_image = self.transform(image)\n                if isinstance(t_image, torch.Tensor):\n                    t_image = t_image.permute(1, 2, 0).numpy()\n                axes[i][1].imshow(t_image)\n                axes[i][1].axis(\"off\")\n                axes[i][1].set_title(f\"Transformed\\nsize: {t_image.size}\", fontsize=12)\n            else:\n                axes[i][1].imshow(image)\n                axes[i][1].axis(\"off\")\n                axes[i][1].set_title(\"No Transform\", fontsize=12)\n\n        plt.tight_layout()\n        plt.show()\n\n    def display_slices(self, n):\n        tomo_dirs = [p for p in self.path.iterdir() if p.is_dir() and p.name.startswith(\"tomo_\")]\n        if not tomo_dirs:\n            raise ValueError(\"No tomo_* directories found.\")\n        tomo_path = random.choice(tomo_dirs)\n        image_list = sorted(tomo_path.glob(\"*.jpg\"))[:n]\n        if not image_list:\n            raise ValueError(f\"No .jpg files found in {tomo_path}\")\n\n        rows = (len(image_list) + 4) // 5\n        fig, axes = plt.subplots(rows, 5, figsize=(16, rows * 4))\n        axes = axes.flatten()\n\n        for i, img_path in enumerate(image_list):\n            image = Image.open(img_path)\n            axes[i].imshow(image)\n            axes[i].axis(\"off\")\n            axes[i].set_title(f\"slice shape: {image.size}\")\n\n        for j in range(i + 1, len(axes)):\n            axes[j].axis('off')\n\n        plt.suptitle(f\"Random tomo: {tomo_path.name}\", fontsize=16)\n        plt.tight_layout()\n        plt.show()\n\n    def plot_with_bounding_boxes(self, n, label_path):\n        df = pd.read_csv(label_path)\n        \n        tomos  = df[df[\"Motor axis 0\"]!=-1]\n        sampled_tomos = tomos[\"tomo_id\"].drop_duplicates().sample(n=n-1, random_state=42)\n        sampled_df = df[df[\"tomo_id\"].isin(sampled_tomos)]\n        \n        rows = int(np.ceil(n / 2))\n        cols = min(n, 2)\n        fig, axes = plt.subplots(rows, cols, figsize=(14, 5 * rows))\n        \n        if n == 1:\n            axes = np.array([axes])\n        axes = axes.flatten()\n\n        for i, (_, motor) in enumerate(sampled_df.iterrows()):\n            z = int(motor[\"Motor axis 0\"])\n            if z == -1:\n                continue\n        \n            img_path = Path(self.path) / motor[\"tomo_id\"] / f\"slice_{z:04d}.jpg\"\n            # if not img_path.exists():\n            #     continue\n        \n            image = Image.open(img_path)\n            if self.transform:\n                image = self.transform(image)\n            img_rgb = image.convert('RGB')\n        \n            img_width, img_height = img_rgb.size\n            overlay = Image.new('RGBA', img_rgb.size, (0, 0, 0, 0))\n            draw = ImageDraw.Draw(overlay)\n        \n            x_center = motor[\"Motor axis 2\"]\n            y_center = motor[\"Motor axis 1\"]\n            width = height = self.box_width  # Default box size if not defined\n        \n            x1 = max(0, int(x_center - width / 2))\n            y1 = max(0, int(y_center - height / 2))\n            x2 = min(img_width, int(x_center + width / 2))\n            y2 = min(img_height, int(y_center + height / 2))\n        \n            draw.rectangle([x1, y1, x2, y2], fill=(255, 0, 0, 64), outline=(255, 0, 0, 200))\n            draw.text((x1, max(0, y1 - 10)), \"Class 0\", fill=(255, 0, 0, 255))\n        \n            img_rgb = Image.alpha_composite(img_rgb.convert('RGBA'), overlay).convert('RGB')\n        \n            axes[i].imshow(np.array(img_rgb))\n            img_name = motor[\"tomo_id\"]+\"_\"+os.path.basename(img_path)\n            axes[i].set_title(f\"Image: {img_name}\")\n            axes[i].axis('on')\n        \n        for j in range(i + 1, len(axes)):\n            axes[j].axis('off')\n        \n        plt.tight_layout()\n        plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:25:09.949515Z","iopub.execute_input":"2025-05-24T04:25:09.949784Z","iopub.status.idle":"2025-05-24T04:25:09.972155Z","shell.execute_reply.started":"2025-05-24T04:25:09.949763Z","shell.execute_reply":"2025-05-24T04:25:09.971422Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class NormalizeByPercentile:\n    def __call__(self, img):\n        if isinstance(img, Image.Image):\n            img = np.array(img)\n\n        p2 = np.percentile(img, 2)\n        p98 = np.percentile(img, 98)\n        clipped = np.clip(img, p2, p98)\n        normalized = 255 * (clipped - p2) / (p98 - p2 + 1e-5)\n        return Image.fromarray(np.uint8(normalized))\ndata_transform = transforms.Compose([\n    NormalizeByPercentile(),\n])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:25:09.974582Z","iopub.execute_input":"2025-05-24T04:25:09.975005Z","iopub.status.idle":"2025-05-24T04:25:09.992296Z","shell.execute_reply.started":"2025-05-24T04:25:09.974985Z","shell.execute_reply":"2025-05-24T04:25:09.991522Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# visualizer = Visualise(path=Path(root_dir/\"train\"), transform=data_transform)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:25:09.993208Z","iopub.execute_input":"2025-05-24T04:25:09.993469Z","iopub.status.idle":"2025-05-24T04:25:10.008279Z","shell.execute_reply.started":"2025-05-24T04:25:09.993450Z","shell.execute_reply":"2025-05-24T04:25:10.007293Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# visualizer.plot_with_bounding_boxes(9,\"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/train_labels.csv\" )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:25:10.009181Z","iopub.execute_input":"2025-05-24T04:25:10.009470Z","iopub.status.idle":"2025-05-24T04:25:10.024762Z","shell.execute_reply.started":"2025-05-24T04:25:10.009444Z","shell.execute_reply":"2025-05-24T04:25:10.023836Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# visualizer.random_tomosplits(n=10)","metadata":{"trusted":true,"scrolled":true,"execution":{"iopub.status.busy":"2025-05-24T04:25:10.025705Z","iopub.execute_input":"2025-05-24T04:25:10.026043Z","iopub.status.idle":"2025-05-24T04:25:10.040214Z","shell.execute_reply.started":"2025-05-24T04:25:10.026014Z","shell.execute_reply":"2025-05-24T04:25:10.039399Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# visualizer.display_transform(n=5)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:25:10.041132Z","iopub.execute_input":"2025-05-24T04:25:10.041907Z","iopub.status.idle":"2025-05-24T04:25:10.059121Z","shell.execute_reply.started":"2025-05-24T04:25:10.041882Z","shell.execute_reply":"2025-05-24T04:25:10.058387Z"},"scrolled":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# visualizer.display_slices(n=5)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:25:10.059947Z","iopub.execute_input":"2025-05-24T04:25:10.060243Z","iopub.status.idle":"2025-05-24T04:25:10.072592Z","shell.execute_reply.started":"2025-05-24T04:25:10.060222Z","shell.execute_reply":"2025-05-24T04:25:10.071780Z"},"scrolled":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class BYUCustomDatasetPreparer:\n    def __init__(self, root, yaml_dir, transform=None, target_transform=True, window_size=24, split_ratio=0.8, trust=2, mode=\"train\", neg_include=False):\n        self.root = Path(root)\n        self.yaml_dir = Path(yaml_dir)\n        self.transform = transform\n        self.target_transform = target_transform\n        self.window_size = window_size\n        self.mode = mode\n        self.split_ratio = split_ratio\n        self.trust = trust\n        self.data = []\n        self.neg_include = neg_include\n\n        self.image_dir = self.root / mode\n        self.paths = sorted(list(self.image_dir.glob(\"*/*.jpg\")))\n\n        if self.mode == \"train\":\n            label_path = self.root / \"train_labels.csv\"\n            self.labels = pd.read_csv(label_path)\n            self.tomos = self.labels[\"tomo_id\"].tolist()\n            self.unique_tomos = list(set(self.tomos))\n            print(f\"No. of total motors present in dataset is {len(self.tomos)} out of {len(self.unique_tomos)} tomograms\")\n            self._init_folders()\n        else:\n            self._init_folders()\n\n    def _init_folders(self):\n        if self.mode == \"train\":\n            self.yolo_images_train = self.yaml_dir / \"images/train\"\n            self.yolo_images_val = self.yaml_dir / \"images/val\"\n            self.yolo_labels_train = self.yaml_dir / \"labels/train\"\n            self.yolo_labels_val = self.yaml_dir / \"labels/val\"\n            paths = [self.yolo_images_train, self.yolo_images_val, self.yolo_labels_train, self.yolo_labels_val]\n        else:\n            self.yolo_images_test = self.yaml_dir / \"images/test\"\n            self.yolo_labels_test = self.yaml_dir / \"labels/test\"\n            paths = [self.yolo_images_test, self.yolo_labels_test]\n        for p in paths:\n            os.makedirs(p, exist_ok=True)\n\n    def _split_train_val(self):\n        train_tomos = random.sample(self.unique_tomos, int(self.split_ratio * len(self.unique_tomos)))\n        val_tomos = [t for t in self.unique_tomos if t not in train_tomos]\n        return train_tomos, val_tomos\n\n    def _extract_unique_tomos(self):\n        return list(set(self.tomos))\n\n    def _process_split(self, tomo_list, images_dir, labels_dir):\n        total_motor_count = 0\n        total_image_count = 0\n        total_unprocessed_tomos = 0\n        desc = f\"preparing training set\" if self.mode == \"train\" else f\"preparing validating set\"\n        for tomo in tqdm(tomo_list, desc=desc):\n            motors = self.labels[self.labels[\"tomo_id\"] == tomo]\n            total_slices = int(motors[\"Array shape (axis 0)\"].iloc[0])\n            present_slices = motors[\"Motor axis 0\"].dropna().astype(int).tolist()\n            present_slices_set = set(present_slices)\n    \n            if self.neg_include:\n                all_slices = list(range(total_slices))\n                negative_slices = [i for i in all_slices if i not in present_slices_set]\n                for i in negative_slices:\n                    label = {\"label_id\": 1, \"x\": None, \"y\": None, \"w\": None, \"h\": None}  # dummy\n                    self._save_image_and_label(tomo, i, label, images_dir, labels_dir)\n                total_image_count += len(negative_slices)\n    \n            m_count, i_count, u_count = self._extract_labels(tomo, motors, total_slices, images_dir, labels_dir)\n            total_motor_count += m_count\n            total_image_count += i_count\n            total_unprocessed_tomos += u_count if not self.neg_include else 0\n\n    \n        print(f\"\\n📊 [Summary] Under {desc} Motors: {total_motor_count}, Images: {total_image_count}, Unprocessed Motors: {total_unprocessed_tomos}\\n\")\n\n    def _extract_labels(self, tomo, motors, total_slices, images_dir, labels_dir):\n        motor_count = 0\n        image_count = 0\n        unprocessod_tomos = 0\n        for _, motor in motors.iterrows():\n            slice_idx = motor[\"Motor axis 0\"]\n            if int(motor[\"Motor axis 0\"]) == -1:\n                unprocessod_tomos += 1\n                continue\n            motor_count += 1\n            slice_idx = int(slice_idx)\n            label = {\"label_id\": 0}\n            if self.target_transform:\n                try:\n                    label[\"y_orig\"] = motor[\"Motor axis 1\"]\n                    label[\"x_orig\"] = motor[\"Motor axis 2\"]\n                    label[\"y\"] = motor[\"Motor axis 1\"] / motor[\"Array shape (axis 1)\"]\n                    label[\"x\"] = motor[\"Motor axis 2\"] / motor[\"Array shape (axis 2)\"]\n                    label[\"h\"] = self.window_size / motor[\"Array shape (axis 1)\"]\n                    label[\"w\"] = self.window_size / motor[\"Array shape (axis 2)\"]\n                except:\n                    continue\n            else:\n                label[\"y\"] = motor[\"Motor axis 1\"]\n                label[\"x\"] = motor[\"Motor axis 2\"]\n                label[\"h\"] = label[\"w\"] = self.window_size\n\n            z_min = max(0, slice_idx - self.trust)\n            z_max = min(total_slices - 1, slice_idx + self.trust)\n            for z in range(z_min, z_max + 1):\n                self._save_image_and_label(tomo, z, label, images_dir, labels_dir)\n                image_count += 1\n        return motor_count, image_count, unprocessod_tomos\n\n    def _save_image_and_label(self, tomo, z, label, images_dir, labels_dir):\n        src_path = self.root / self.mode / tomo / f\"slice_{z:04d}.jpg\"\n        if not src_path.exists():\n            return\n        image = Image.open(src_path)\n        if self.transform:\n            image = self.transform(image)\n        img_name = f\"{tomo}_z{z:04d}_y{int(label['y_orig']):04d}_x{int(label['x_orig']):04d}.jpg\"\n        img_path = images_dir / img_name\n        image.save(img_path)\n        label_path = labels_dir / img_name.replace(\".jpg\", \".txt\")\n        with open(label_path, \"w\") as f:\n            f.write(f\"{int(label['label_id'])} {label['x']} {label['y']} {label['w']} {label['h']}\")\n\n    def prepare(self):\n        if self.mode == \"train\":\n            train_tomos, val_tomos = self._split_train_val()\n            print(f\"🟢 Train tomos: {len(train_tomos)}, 🔵 Val tomos: {len(val_tomos)}\")\n            self._process_split(train_tomos, self.yolo_images_train, self.yolo_labels_train)\n            self._process_split(val_tomos, self.yolo_images_val, self.yolo_labels_val)\n        else:\n            unique_tomos = self._extract_unique_tomos()\n            print(f\"🟠 Test tomos: {len(unique_tomos)}\")\n            self._process_split(unique_tomos, self.yolo_images_test, self.yolo_labels_test)\n\n    def create_yaml(self):\n        if self.neg_include:\n            yaml_content = {\n                'path': str(self.yaml_dir),\n                'train': 'images/train',\n                'val': 'images/val',\n                'names': {0: 'motor', 1: 'no_motor'}\n            }\n        else:\n            yaml_content = {\n                'path': str(self.yaml_dir),\n                'train': 'images/train',\n                'val': 'images/val',\n                'names': {0: 'motor'}\n            }\n    \n        with open(self.yaml_dir / 'dataset.yaml', 'w') as f:\n            yaml.dump(yaml_content, f, default_flow_style=False)\n        print(\"✅ dataset.yaml created.\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:25:10.073539Z","iopub.execute_input":"2025-05-24T04:25:10.073812Z","iopub.status.idle":"2025-05-24T04:25:10.096948Z","shell.execute_reply.started":"2025-05-24T04:25:10.073792Z","shell.execute_reply":"2025-05-24T04:25:10.096066Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# data = BYUCustomDatasetPreparer(root_dir, yaml_dir, transform = data_transform, trust = 4)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:26:29.981796Z","iopub.execute_input":"2025-05-24T04:26:29.982117Z","iopub.status.idle":"2025-05-24T04:26:29.985968Z","shell.execute_reply.started":"2025-05-24T04:26:29.982095Z","shell.execute_reply":"2025-05-24T04:26:29.985134Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# data.prepare()\n# data.create_yaml()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:26:30.263241Z","iopub.execute_input":"2025-05-24T04:26:30.263666Z","iopub.status.idle":"2025-05-24T04:26:30.267368Z","shell.execute_reply.started":"2025-05-24T04:26:30.263642Z","shell.execute_reply":"2025-05-24T04:26:30.266551Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class YOLOv1(nn.Module):\n    def __init__(self, S=7, B=2, C=20):\n        super(YOLOv1, self).__init__()\n        self.S, self.B, self.C = S, B, C\n        self.conv_layers = nn.Sequential(\n            nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3),\n            nn.LeakyReLU(0.1),\n            nn.MaxPool2d(2, 2),\n            nn.Conv2d(64, 192, kernel_size=3, padding=1),\n            nn.LeakyReLU(0.1),\n            nn.MaxPool2d(2, 2),\n        )\n        self.fc_layers = nn.Sequential(\n            nn.Flatten(),\n            nn.Linear(192 * (S // 4) * (S // 4), 4096),\n            nn.LeakyReLU(0.1),\n            nn.Linear(4096, S * S * (B * 5 + C))\n        )\n\n    def forward(self, x):\n        x = self.conv_layers(x)\n        x = self.fc_layers(x)\n        return x.view(-1, self.S, self.S, self.B * 5 + self.C)\n\n\nclass YOLOModelLoader:\n    def __init__(self, model_version: str, num_classes: int, model_path: str = None,\n                 device='cuda' if torch.cuda.is_available() else 'cpu'):\n        self.model_version = model_version.lower()\n        self.num_classes = num_classes\n        self.model_path = model_path\n        self.device = device\n        self.model = self._load_model()\n\n    def _load_model(self):\n        if self.model_version == 'yolov1':\n            return YOLOv1(S=7, B=2, C=self.num_classes).to(self.device)\n\n        model_source = self.model_path \n        if self.model_version.startswith('yolov') or self.model_version.startswith('yolo'):\n            return YOLO(Path(model_source+f\"/{self.model_version}.pt\"))\n\n        elif self.model_version.startswith(\"rtdetr\"):\n            return RTDETR(Path(model_source+f\"/{self.model_version}.pt\"))\n\n        else:\n            raise ValueError(f\"Unsupported model version: {self.model_version}\")\n\n    def get_model(self):\n        return self.model\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:26:30.580783Z","iopub.execute_input":"2025-05-24T04:26:30.581582Z","iopub.status.idle":"2025-05-24T04:26:30.590682Z","shell.execute_reply.started":"2025-05-24T04:26:30.581555Z","shell.execute_reply":"2025-05-24T04:26:30.589824Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# loader = YOLOModelLoader(model_version='yolov8n', num_classes=1, model_path='/kaggle/input/yolo/pytorch/default/1')\n# model = loader.get_model()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:26:30.881051Z","iopub.execute_input":"2025-05-24T04:26:30.881381Z","iopub.status.idle":"2025-05-24T04:26:30.885295Z","shell.execute_reply.started":"2025-05-24T04:26:30.881326Z","shell.execute_reply":"2025-05-24T04:26:30.884435Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class evaluate_data(Dataset):\n    def __init__(self, root, evaluate_dir, transform=None, window_size=24, mode=\"evaluate\"):\n        self.root = Path(root)\n        self.transform = transform\n        self.mode = mode\n        self.window_size = window_size\n        self.evaluate_dir= evaluate_dir\n\n        if self.mode == \"evaluate\":\n            self.image_dir = Path(self.root / \"train\")\n            label_path = self.root / \"train_labels.csv\"\n            self.labels = pd.read_csv(label_path)\n            self.paths = []\n            for _, motor in self.labels.iterrows():\n                if motor[\"Motor axis 0\"] == -1:\n                    continue\n                self.paths.append(self.image_dir / motor[\"tomo_id\"] / f'slice_{int(motor[\"Motor axis 0\"]):04d}.jpg')\n        else:\n            self.image_dir = Path(self.root / \"test\")\n            self.paths = sorted(list(self.image_dir.glob(\"*/*.jpg\")))\n\n        self.slice_info = [\n            {\n                \"path\": p,\n                \"tomo_id\": p.parent.name,\n                \"slice_idx\": int(p.stem.split(\"_\")[1])\n            }\n            for p in self.paths\n        ]\n\n    def __len__(self):\n        return len(self.slice_info)\n\n    def __getitem__(self, idx):\n        file = self.slice_info[idx]\n        src_image = Image.open(file[\"path\"])\n        image = self.transform(src_image) if self.transform else src_image\n\n        tomo = file[\"tomo_id\"]\n        slice_idx = file[\"slice_idx\"]\n        dest_path = os.path.join(self.evaluate_dir, tomo, os.path.basename(file[\"path\"]))\n        os.makedirs(os.path.dirname(dest_path), exist_ok=True)\n        image.save(dest_path)\n\n        if self.mode == \"evaluate\":\n            slice_matches = self.labels[(self.labels[\"tomo_id\"] == tomo) & (self.labels[\"Motor axis 0\"] == slice_idx)]\n            row = slice_matches.iloc[0]\n            y = row[\"Motor axis 1\"]\n            x = row[\"Motor axis 2\"]\n            h = w =  self.window_size \n            return {\n                \"tomo\": tomo,\n                \"slice\": slice_idx,\n                \"image\": image,\n                \"path\": dest_path,\n                \"label\": {\n                    \"x\": x,\n                    \"y\": y,\n                    \"w\": w,\n                    \"h\": h\n                }\n            }\n        return {\n                \"tomo\": tomo,\n                \"slice\": slice_idx,\n                \"image\": image,\n                \"path\": dest_path}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:26:31.160663Z","iopub.execute_input":"2025-05-24T04:26:31.160968Z","iopub.status.idle":"2025-05-24T04:26:31.170893Z","shell.execute_reply.started":"2025-05-24T04:26:31.160936Z","shell.execute_reply":"2025-05-24T04:26:31.170158Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class Evaluate:\n    def __init__(self, dataset):\n        self.data = dataset\n\n    def plot_loss_curve(self, run_dir):\n        results_csv = os.path.join(run_dir, 'results.csv')  \n        if not os.path.exists(results_csv):\n            print(f\"Results file not found at {results_csv}\")\n            return\n        \n        results_df = pd.read_csv(results_csv)\n        train_dfl_col = [col for col in results_df.columns if 'train/dfl_loss' in col]\n        val_dfl_col = [col for col in results_df.columns if 'val/dfl_loss' in col]\n        \n        if not train_dfl_col or not val_dfl_col:\n            print(\"DFL loss columns not found in results CSV\")\n            print(f\"Available columns: {results_df.columns.tolist()}\")\n            return\n        \n        train_dfl_col = train_dfl_col[0]\n        val_dfl_col = val_dfl_col[0]\n        \n        best_epoch = results_df[val_dfl_col].idxmin()\n        best_val_loss = results_df.loc[best_epoch, val_dfl_col]\n        \n        plt.figure(figsize=(10, 6))\n        plt.plot(results_df['epoch'], results_df[train_dfl_col], label='Train DFL Loss')\n        plt.plot(results_df['epoch'], results_df[val_dfl_col], label='Validation DFL Loss')\n        plt.axvline(x=results_df.loc[best_epoch, 'epoch'], color='r', linestyle='--', \n                    label=f'Best Model (Epoch {int(results_df.loc[best_epoch, \"epoch\"])}, Val Loss: {best_val_loss:.4f})')\n        \n        plt.xlabel('Epoch')\n        plt.ylabel('DFL Loss')\n        plt.title('Training and Validation DFL Loss')\n        plt.legend()\n        plt.grid(True, linestyle='--', alpha=0.7)\n        plt.show()\n        plt.close()\n        return best_epoch, best_val_loss\n\n    def predict_on_samples(self, model, num_samples=4):\n        num_samples = min(num_samples, len(self.data))\n        indices = random.sample(range(len(self.data)), num_samples)\n        samples = [self.data[i] for i in indices]\n    \n        cols = 2\n        rows = int(np.ceil(num_samples / cols))\n        \n        fig, axes = plt.subplots(rows, cols, figsize=(4 * cols, 4 * rows))\n        axes = axes.flatten() if rows > 1 else np.array(axes).reshape(-1)\n    \n        for i, file in enumerate(samples):\n            img_path = file[\"path\"]\n            results = model.predict(img_path, conf=0.25)[0]\n            img = Image.open(img_path)\n            axes[i].imshow(np.array(img), cmap='gray')\n    \n            \n            x_gt = file[\"label\"][\"x\"]\n            y_gt = file[\"label\"][\"y\"]\n            w_gt = file[\"label\"][\"w\"]\n            h_gt = file[\"label\"][\"h\"]\n            rect_gt = Rectangle((x_gt - w_gt / 2, y_gt - h_gt / 2),\n                                w_gt, h_gt, linewidth=1, edgecolor='g', facecolor='none')\n            axes[i].add_patch(rect_gt)\n    \n            if hasattr(results, 'boxes') and len(results.boxes) > 0:\n                boxes = results.boxes.xyxy.cpu().numpy()\n                confs = results.boxes.conf.cpu().numpy()\n                for box, conf in zip(boxes, confs):\n                    x1, y1, x2, y2 = box\n                    rect_pred = Rectangle((x1, y1), x2 - x1, y2 - y1,\n                                          linewidth=1, edgecolor='r', facecolor='none')\n                    axes[i].add_patch(rect_pred)\n                    axes[i].text(x1, y1 - 5, f'{conf:.2f}', color='red')\n    \n            axes[i].set_title(f\"Ground Truth (green) vs Prediction (red)\")\n        for j in range(i + 1, len(axes)):\n            axes[j].axis('off')\n        plt.tight_layout()\n        plt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:26:35.560679Z","iopub.execute_input":"2025-05-24T04:26:35.560994Z","iopub.status.idle":"2025-05-24T04:26:35.575229Z","shell.execute_reply.started":"2025-05-24T04:26:35.560972Z","shell.execute_reply":"2025-05-24T04:26:35.574316Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def train_yolo_model(yaml_path,yolo_weights_dir, model, epochs=30, batch_size=16, img_size=640):\n    print(f\"Loading model\")\n    results = model.train(\n        data=yaml_path,\n        epochs=epochs,\n        batch=batch_size,\n        imgsz=img_size,\n        project=yolo_weights_dir,\n        name='motor_detector',\n        exist_ok=True,\n        patience=5,              # Early stopping if no improvement for 5 epochs\n        save_period=5,           # Save checkpoints every 5 epochs\n        val=True,                # Ensure validation is performed\n        verbose=True             # Show detailed output during training\n    )\n    \n    # Get the path to the run directory\n    run_dir = os.path.join(yolo_weights_dir, 'motor_detector')    \n    best_epoch_info = evaluate.plot_loss_curve(run_dir)\n    \n    if best_epoch_info:\n        best_epoch, best_val_loss = best_epoch_info\n        print(f\"\\nBest model found at epoch {best_epoch} with validation DFL loss: {best_val_loss:.4f}\")\n    \n    return model, results\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:26:36.180730Z","iopub.execute_input":"2025-05-24T04:26:36.181000Z","iopub.status.idle":"2025-05-24T04:26:36.186508Z","shell.execute_reply.started":"2025-05-24T04:26:36.180982Z","shell.execute_reply":"2025-05-24T04:26:36.185729Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# yolo_weights_dir = \"/kaggle/working/yolo_weights\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:26:38.680696Z","iopub.execute_input":"2025-05-24T04:26:38.681006Z","iopub.status.idle":"2025-05-24T04:26:38.684636Z","shell.execute_reply.started":"2025-05-24T04:26:38.680981Z","shell.execute_reply":"2025-05-24T04:26:38.683818Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# print(\"Starting YOLO training process...\")\n# yaml_path = Path(\"/kaggle/working/yaml_dir/dataset.yaml\")\n# print(f\"Using YAML file: {yaml_path}\")\n\n# evaluate_dir = \"/kaggle/working/evaluate\"\n# os.makedirs(evaluate_dir, exist_ok=True)\n# data = evaluate_data(root_dir, evaluate_dir, transform=data_transform)\n# evaluate = Evaluate(data)\n\n\n# print(\"\\nStarting YOLO training...\")\n# model, results = train_yolo_model(\n#     yaml_path=yaml_path,\n#     yolo_weights_dir=yolo_weights_dir,\n#     model=model,\n#     epochs=30  # Using 30 epochs instead of 100 for faster training\n# )\n\n# print(\"\\nTraining complete!\")\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:26:39.295734Z","iopub.execute_input":"2025-05-24T04:26:39.296009Z","iopub.status.idle":"2025-05-24T04:26:39.300847Z","shell.execute_reply.started":"2025-05-24T04:26:39.295990Z","shell.execute_reply":"2025-05-24T04:26:39.299712Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# data = evaluate_data(root_dir, evaluate_dir, transform=data_transform)\n# evaluate = Evaluate(data)\n# print(\"\\nRunning predictions on sample images...\")\n# evaluate.predict_on_samples(model, num_samples=8)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:26:42.345881Z","iopub.execute_input":"2025-05-24T04:26:42.346166Z","iopub.status.idle":"2025-05-24T04:26:42.349730Z","shell.execute_reply.started":"2025-05-24T04:26:42.346146Z","shell.execute_reply":"2025-05-24T04:26:42.348802Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"np.random.seed(42)\ntorch.manual_seed(42)\n\ndata_path = root_dir\ntest_dir = os.path.join(data_path, \"test\")\nsubmission_path = \"/kaggle/working/submission.csv\"\n\nmodel_path = \"/kaggle/input/v2/pytorch/default/1/best (2).pt\"\n\nCONFIDENCE_THRESHOLD = 0.45\nMAX_DETECTIONS_PER_TOMO = 3\nNMS_IOU_THRESHOLD = 0.2\nCONCENTRATION = 1\n\nclass 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\ndevice = 'cuda:0' if torch.cuda.is_available() else 'cpu'\nBATCH_SIZE = 8\n\nif device.startswith('cuda'):\n    torch.backends.cudnn.benchmark = True\n    torch.backends.cudnn.deterministic = False\n    torch.backends.cuda.matmul.allow_tf32 = True\n    torch.backends.cudnn.allow_tf32 = True\n\n    gpu_name = torch.cuda.get_device_name(0)\n    gpu_mem = torch.cuda.get_device_properties(0).total_memory / 1e9\n    print(f\"Using GPU: {gpu_name} with {gpu_mem:.2f} GB memory\")\n\n    free_mem = gpu_mem - torch.cuda.memory_allocated(0) / 1e9\n    BATCH_SIZE = max(8, min(32, int(free_mem * 4)))\n    print(f\"Dynamic batch size set to {BATCH_SIZE} based on {free_mem:.2f}GB free memory\")\nelse:\n    print(\"GPU not available, using CPU\")\n    BATCH_SIZE = 4\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:26:43.381446Z","iopub.execute_input":"2025-05-24T04:26:43.381732Z","iopub.status.idle":"2025-05-24T04:26:43.392557Z","shell.execute_reply.started":"2025-05-24T04:26:43.381711Z","shell.execute_reply":"2025-05-24T04:26:43.391828Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\nclass Test:\n    def __init__(self, model, device, test_dir, batch_size, confidence_threshold,\n             concentration, submission_path, nms_iou_threshold, GPUProfiler=None,\n             MAX_DETECTIONS_PER_TOMO=1):\n        self.MAX_DETECTIONS_PER_TOMO = MAX_DETECTIONS_PER_TOMO\n        self.model = YOLO(model)\n        self.device = device\n        self.test_dir = test_dir \n        self.batch_size = batch_size\n        self.confidence_threshold = confidence_threshold\n        self.concentration = concentration\n        self.submission_path = submission_path\n        self.nms_iou_threshold = nms_iou_threshold\n        self.GPUProfiler = GPUProfiler\n\n        self.model.to(self.device)\n        if self.device.startswith('cuda'):\n            self.model.fuse()\n            if torch.cuda.get_device_capability(0)[0] >= 7:\n                self.model.model.half()\n\n    def preload_image_batch(self, file_paths):\n        images = []\n        for path in file_paths:\n            img = cv2.imread(path)\n            if img is None:\n                img = np.array(Image.open(path))\n            images.append(img)\n        return images\n\n    def process_tomogram(self, tomo_id, index=0, total=1):\n        print(f\"Processing tomogram {tomo_id} ({index}/{total})\")\n        tomo_dir = os.path.join(self.test_dir, tomo_id)\n        slice_files = sorted([f for f in os.listdir(tomo_dir) if f.endswith('.jpg')])\n        selected_indices = np.linspace(0, len(slice_files)-1, int(len(slice_files) * self.concentration))\n        selected_indices = np.round(selected_indices).astype(int)\n        slice_files = [slice_files[i] for i in selected_indices]\n\n        print(f\"Processing {len(slice_files)} out of {len(os.listdir(tomo_dir))} slices based on CONCENTRATION={self.concentration}\")\n\n        all_detections = []\n        streams = [torch.cuda.Stream() for _ in range(min(4, self.batch_size))] if self.device.startswith('cuda') else [None]\n        next_batch_thread = None\n        next_batch_images = None\n\n        for batch_start in range(0, len(slice_files), self.batch_size):\n            if next_batch_thread is not None:\n                next_batch_thread.join()\n                next_batch_images = None\n\n            batch_end = min(batch_start + self.batch_size, len(slice_files))\n            batch_files = slice_files[batch_start:batch_end]\n\n            next_batch_start = batch_end\n            next_batch_end = min(next_batch_start + self.batch_size, len(slice_files))\n            next_batch_files = slice_files[next_batch_start:next_batch_end] if next_batch_start < len(slice_files) else []\n\n            if next_batch_files:\n                next_batch_paths = [os.path.join(tomo_dir, f) for f in next_batch_files]\n                next_batch_thread = threading.Thread(target=self.preload_image_batch, args=(next_batch_paths,))\n                next_batch_thread.start()\n            else:\n                next_batch_thread = None\n\n            sub_batches = np.array_split(batch_files, len(streams))\n            sub_batch_results = []\n\n            for i, sub_batch in enumerate(sub_batches):\n                if len(sub_batch) == 0:\n                    continue\n                stream = streams[i % len(streams)]\n                with torch.cuda.stream(stream) if stream and self.device.startswith('cuda') else nullcontext():\n                    sub_batch_paths = [os.path.join(tomo_dir, slice_file) for slice_file in sub_batch]\n                    sub_batch_slice_nums = [int(slice_file.split('_')[1].split('.')[0]) for slice_file in sub_batch]\n\n                    if self.GPUProfiler:\n                        with self.GPUProfiler(f\"Inference batch {i+1}/{len(sub_batches)}\"):\n                            sub_results = self.model(sub_batch_paths, verbose=False)\n                    else:\n                        sub_results = self.model(sub_batch_paths, verbose=False)\n\n                    for j, result in enumerate(sub_results):\n                        if len(result.boxes) > 0:\n                            boxes = result.boxes\n                            for box_idx, confidence in enumerate(boxes.conf):\n                                if confidence >= self.confidence_threshold:\n                                    x1, y1, x2, y2 = boxes.xyxy[box_idx].cpu().numpy()\n                                    x_center = (x1 + x2) / 2\n                                    y_center = (y1 + y2) / 2\n                                    all_detections.append({\n                                        'z': round(sub_batch_slice_nums[j]),\n                                        'y': round(y_center),\n                                        'x': round(x_center),\n                                        'confidence': float(confidence)\n                                    })\n\n            if self.device.startswith('cuda'):\n                torch.cuda.synchronize()\n\n        if next_batch_thread is not None:\n            next_batch_thread.join()\n\n        final_detections = self.perform_3d_nms(all_detections, self.nms_iou_threshold)\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        top_detections = final_detections[:self.MAX_DETECTIONS_PER_TOMO]\n        return [{\n            'tomo_id': tomo_id,\n            'Motor axis 0': round(det['z']),\n            'Motor axis 1': round(det['y']),\n            'Motor axis 2': round(det['x'])\n        } for det in top_detections]\n\n\n    def perform_3d_nms(self, detections, iou_threshold):\n        if not detections:\n            return []\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        box_size = 24\n        distance_threshold = box_size * iou_threshold\n\n        while detections:\n            best_detection = detections.pop(0)\n            final_detections.append(best_detection)\n            detections = [d for d in detections if distance_3d(d, best_detection) > distance_threshold]\n\n        return final_detections\n\n    def generate_submission(self):\n        test_tomos = sorted([d for d in os.listdir(self.test_dir) if os.path.isdir(os.path.join(self.test_dir, d))])\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        results = []\n        motors_found = 0\n\n        with ThreadPoolExecutor(max_workers=1) as executor:\n            future_to_tomo = {}\n            for i, tomo_id in enumerate(test_tomos, 1):\n                future = executor.submit(self.process_tomogram, tomo_id, i, total_tomos)\n                future_to_tomo[future] = tomo_id\n\n            for future in future_to_tomo:\n                tomo_id = future_to_tomo[future]\n                try:\n                    if torch.cuda.is_available():\n                        torch.cuda.empty_cache()\n\n                    result_list = future.result()\n                    results.extend(result_list)\n\n                    motors_in_tomo = [r for r in result_list if r['Motor axis 0'] != -1]\n                    if motors_in_tomo:\n                        motors_found += len(motors_in_tomo)\n                        for r in motors_in_tomo:\n                            print(f\"Motor found in {r['tomo_id']} at position: \"\n                                  f\"z={r['Motor axis 0']}, y={r['Motor axis 1']}, x={r['Motor axis 2']}\")\n                    else:\n                        print(f\"No motor detected in {tomo_id}\")\n                    \n                    print(f\"Current detection count: {motors_found} motors in {len(results)} entries \"\n                          f\"({motors_found / len(results) * 100:.1f}%)\")\n\n                except Exception as e:\n                    print(f\"Error processing {tomo_id}: {e}\")\n                    results.append({\n                        'tomo_id': tomo_id,\n                        'Motor axis 0': -1,\n                        'Motor axis 1': -1,\n                        'Motor axis 2': -1\n                    })\n\n        submission_df = pd.DataFrame(results)\n        submission_df = submission_df[['tomo_id', 'Motor axis 0', 'Motor axis 1', 'Motor axis 2']]\n        submission_df.to_csv(self.submission_path, index=False)\n\n        print(f\"\\nSubmission complete!\")\n        print(f\"Motors detected: {motors_found}/{total_tomos} ({motors_found/total_tomos*100:.1f}%)\")\n        print(f\"Submission saved to: {self.submission_path}\")\n        print(\"\\nSubmission preview:\")\n        print(submission_df.head())\n\n        return submission_df\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:26:43.821534Z","iopub.execute_input":"2025-05-24T04:26:43.821854Z","iopub.status.idle":"2025-05-24T04:26:43.847494Z","shell.execute_reply.started":"2025-05-24T04:26:43.821832Z","shell.execute_reply":"2025-05-24T04:26:43.846680Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test = Test(\n    model=model_path,\n    device='cuda' if torch.cuda.is_available() else 'cpu',\n    test_dir=test_dir,\n    batch_size=BATCH_SIZE,\n    confidence_threshold=CONFIDENCE_THRESHOLD,\n    concentration=CONCENTRATION,\n    submission_path='submission.csv',\n    nms_iou_threshold=NMS_IOU_THRESHOLD,\n    GPUProfiler=GPUProfiler,\n    MAX_DETECTIONS_PER_TOMO=MAX_DETECTIONS_PER_TOMO\n)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:26:48.715654Z","iopub.execute_input":"2025-05-24T04:26:48.715962Z","iopub.status.idle":"2025-05-24T04:26:48.859236Z","shell.execute_reply.started":"2025-05-24T04:26:48.715940Z","shell.execute_reply":"2025-05-24T04:26:48.858638Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if __name__ == \"__main__\":\n    # Time entire process\n    start_time = time.time()\n    \n    # Generate submission\n    submission = test.generate_submission()\n    \n    # Print total execution time\n    elapsed = time.time() - start_time\n    print(f\"\\nTotal execution time: {elapsed:.2f} seconds ({elapsed/60:.2f} minutes)\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-24T04:26:51.260568Z","iopub.execute_input":"2025-05-24T04:26:51.261193Z","iopub.status.idle":"2025-05-24T04:29:35.703624Z","shell.execute_reply.started":"2025-05-24T04:26:51.261170Z","shell.execute_reply":"2025-05-24T04:29:35.702603Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}