{"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":"import os\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nfrom tqdm import tqdm\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision import transforms\nfrom torch.optim import Adam","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-04-14T09:44:33.674966Z","iopub.execute_input":"2025-04-14T09:44:33.675279Z","iopub.status.idle":"2025-04-14T09:44:41.190523Z","shell.execute_reply.started":"2025-04-14T09:44:33.675248Z","shell.execute_reply":"2025-04-14T09:44:41.189824Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"TRAIN_DIR = \"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/train\"\nTEST_DIR = \"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/test\"\nLABELS_PATH = \"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/train_labels.csv\"\nSUBMIT_PATH = \"/kaggle/working/submission.csv\"\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-14T09:44:41.191631Z","iopub.execute_input":"2025-04-14T09:44:41.192086Z","iopub.status.idle":"2025-04-14T09:44:41.242260Z","shell.execute_reply.started":"2025-04-14T09:44:41.192057Z","shell.execute_reply":"2025-04-14T09:44:41.241168Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"transform = transforms.Compose([\n    transforms.Resize((128, 128)),\n    transforms.ToTensor(),\n    transforms.Normalize(mean=[0.5], std=[0.5])\n])\n\nclass TomogramDataset(Dataset):\n    def __init__(self, root_dir, labels_df=None, transform=None, mode='train', target_shape=(64, 128, 128)):\n        self.root_dir = root_dir\n        self.tomo_dirs = sorted(os.listdir(root_dir))\n        self.transform = transform\n        self.mode = mode\n        self.labels_df = labels_df\n        self.target_shape = target_shape\n\n        self.label_dict = {}\n        if labels_df is not None:\n            for _, row in labels_df.iterrows():\n                self.label_dict.setdefault(row['tomo_id'], []).append(\n                    (row['Motor axis 0'], row['Motor axis 1'], row['Motor axis 2'])\n                )\n\n    def __len__(self):\n        return len(self.tomo_dirs)\n\n    def __getitem__(self, idx):\n        tomo_id = self.tomo_dirs[idx]\n        tomo_path = os.path.join(self.root_dir, tomo_id)\n\n        slice_files = sorted(os.listdir(tomo_path))\n        slices = []\n        for f in slice_files:\n            try:\n                img = Image.open(os.path.join(tomo_path, f)).convert(\"L\")\n                slices.append(np.array(img, dtype=np.float32))\n            except Exception as e:\n                print(f\"Error reading slice {f} in {tomo_id}: {e}\")\n                continue\n\n        if len(slices) == 0:\n            raise ValueError(f\"No valid slices found for {tomo_id}\")\n\n        volume = np.stack(slices)\n        volume = (volume - volume.min()) / (volume.max() - volume.min() + 1e-5)\n        volume = torch.tensor(volume).unsqueeze(0)\n\n        # Resize to target_shape (C, D, H, W)\n        volume = F.interpolate(volume.unsqueeze(0), size=self.target_shape, mode='trilinear', align_corners=False).squeeze(0)\n\n        if self.mode == 'train':\n            label = self.label_dict.get(tomo_id, [(-1, -1, -1)])\n            orig_shape = (len(slice_files), volume.shape[1], volume.shape[2])\n            scale = [self.target_shape[i] / orig_shape[i] for i in range(3)]\n            target = torch.tensor([\n                label[0][0] * scale[0],\n                label[0][1] * scale[1],\n                label[0][2] * scale[2]\n            ], dtype=torch.float32)\n            return volume, target\n        else:\n            return volume, tomo_id\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-14T09:44:41.244109Z","iopub.execute_input":"2025-04-14T09:44:41.244450Z","iopub.status.idle":"2025-04-14T09:44:41.265021Z","shell.execute_reply.started":"2025-04-14T09:44:41.244416Z","shell.execute_reply":"2025-04-14T09:44:41.264225Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class MotorNet(nn.Module):\n    def __init__(self):\n        super(MotorNet, self).__init__()\n        self.conv1 = nn.Conv3d(1, 16, kernel_size=3, padding=1)\n        self.pool1 = nn.MaxPool3d(2)\n        self.conv2 = nn.Conv3d(16, 32, kernel_size=3, padding=1)\n        self.pool2 = nn.MaxPool3d(2)\n        self.conv3 = nn.Conv3d(32, 64, kernel_size=3, padding=1)\n        self.pool3 = nn.AdaptiveAvgPool3d((4, 8, 8))\n        self.fc1 = nn.Linear(64 * 4 * 8 * 8, 128)\n        self.fc2 = nn.Linear(128, 3)\n\n    def forward(self, x):\n        x = self.pool1(F.relu(self.conv1(x)))\n        x = self.pool2(F.relu(self.conv2(x)))\n        x = self.pool3(F.relu(self.conv3(x)))\n        x = x.view(x.size(0), -1)\n        x = F.relu(self.fc1(x))\n        return self.fc2(x)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-14T09:44:41.265961Z","iopub.execute_input":"2025-04-14T09:44:41.266192Z","iopub.status.idle":"2025-04-14T09:44:41.284605Z","shell.execute_reply.started":"2025-04-14T09:44:41.266164Z","shell.execute_reply":"2025-04-14T09:44:41.283745Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def main():\n    labels_df = pd.read_csv(LABELS_PATH)\n\n    train_dataset = TomogramDataset(TRAIN_DIR, labels_df, transform=transform)\n    train_loader = DataLoader(train_dataset, batch_size=2, shuffle=True, num_workers=0)\n\n    test_dataset = TomogramDataset(TEST_DIR, mode='test', transform=transform)\n    test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False, num_workers=0)\n\n    model = MotorNet().to(device)\n    optimizer = Adam(model.parameters(), lr=1e-4)\n    criterion = nn.MSELoss()\n\n    EPOCHS = 10\n    for epoch in range(EPOCHS):\n        model.train()\n        running_loss = 0\n        for volume, target in tqdm(train_loader, desc=f\"Epoch {epoch+1}\"):\n            volume, target = volume.to(device), target.to(device)\n            optimizer.zero_grad()\n            output = model(volume)\n            loss = criterion(output, target)\n            loss.backward()\n            optimizer.step()\n            running_loss += loss.item()\n        print(f\"Epoch {epoch+1}/{EPOCHS}, Loss: {running_loss / len(train_loader):.4f}\")\n\n    torch.save(model.state_dict(), \"flagellar_motors_model.pth\")\n\n    model.eval()\n    predictions = []\n    with torch.no_grad():\n        for volume, tomo_id in tqdm(test_loader, desc=\"Predicting\"):\n            volume = volume.to(device)\n            output = model(volume).cpu().numpy()[0]\n            output = [int(v) if 0 <= v <= 512 else -1 for v in output]\n            predictions.append([tomo_id[0]] + output)\n\n    submission_df = pd.DataFrame(predictions, columns=[\"tomo_id\", \"Motor axis 0\", \"Motor axis 1\", \"Motor axis 2\"])\n    submission_df.to_csv(SUBMIT_PATH, index=False)\n    print(submission_df.head())\n\nif __name__ == '__main__':\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-14T09:44:41.285554Z","iopub.execute_input":"2025-04-14T09:44:41.285918Z"}},"outputs":[],"execution_count":null}]}