{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":73479,"databundleVersionId":8144099,"sourceType":"competition"}],"dockerImageVersionId":30674,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"**This is a sample script to help you get started. I've created a basic convolutional neural network to train on the images and their labels. Feel free to make your own copy and adjust it as needed.**","metadata":{}},{"cell_type":"code","source":"import os\nimport pandas as pd\nfrom sklearn.model_selection import train_test_split\nfrom PIL import Image\nfrom torchvision import transforms, utils\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader, Subset\nimport numpy as np\n\n# Setting the paths for our data\ntrain_images_dir = \"/kaggle/input/byu-locating-bacterial-flagellar-motors-v2/train\"\ntrain_csv_path = \"/kaggle/input/byu-locating-bacterial-flagellar-motors-v2/train.csv\"\nval_images_dir = \"/kaggle/input/byu-locating-bacterial-flagellar-motors-v2/val\"\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n# Define the custom dataset\nclass BacteriaDataset(Dataset):\n    def __init__(self, data, images_dir, transform=None):\n        self.data = data\n        self.images_dir = images_dir\n        self.transform = transform\n\n    def __len__(self):\n        return len(self.data)\n\n    def __getitem__(self, idx):\n        image_name = self.data.iloc[idx][\"image_name\"]\n        image_path = os.path.join(self.images_dir, image_name)\n        if os.path.exists(image_path):\n            image = Image.open(image_path).convert('RGB')\n        else:\n            image = Image.new('RGB', (224, 224), (0, 0, 0))\n\n        x = self.data.iloc[idx][\"x\"]\n        y = self.data.iloc[idx][\"y\"]\n        motor_visible = self.data.iloc[idx][\"motor_visible\"]\n\n        if self.transform:\n            image = self.transform(image)\n\n        # Scale the targets\n        x = x / 511\n        y = y / 511\n        motor_visible = float(motor_visible)\n\n        return image, x, y, motor_visible\n# Define the data transformations\ntransform = transforms.Compose([\n    transforms.Resize((224, 224)),\n    transforms.ToTensor(),\n    transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])\n])\n\n# Load the full dataset\nfull_dataset = BacteriaDataset(pd.read_csv(train_csv_path), train_images_dir, transform=transform)\n\n# Split the dataset into train and validation sets\ntrain_size = int(0.8 * len(full_dataset))\nval_size = len(full_dataset) - train_size\ntrain_dataset, val_dataset = torch.utils.data.random_split(full_dataset, [train_size, val_size])\n\n# DataLoaders with specified number of workers\ntrain_loader = DataLoader(train_dataset, batch_size=32, shuffle=True, num_workers=4)\nval_loader = DataLoader(val_dataset, batch_size=32, shuffle=False, num_workers=4)","metadata":{"execution":{"iopub.status.busy":"2024-04-11T18:16:30.324783Z","iopub.execute_input":"2024-04-11T18:16:30.325174Z","iopub.status.idle":"2024-04-11T18:16:30.363144Z","shell.execute_reply.started":"2024-04-11T18:16:30.325143Z","shell.execute_reply":"2024-04-11T18:16:30.362426Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class BacteriaCNN(nn.Module):\n    def __init__(self):\n        super(BacteriaCNN, self).__init__()\n        self.backbone = nn.Sequential(\n            nn.Conv2d(3, 16, kernel_size=3, stride=1, padding=1),\n            nn.ReLU(),\n            nn.MaxPool2d(kernel_size=2),\n            nn.Conv2d(16, 32, kernel_size=3, stride=1, padding=1),\n            nn.ReLU(),\n            nn.MaxPool2d(kernel_size=2),\n            nn.Flatten(),\n            nn.Linear(32 * 56 * 56, 128),\n            nn.ReLU()\n        )\n        self.motor_visible = nn.Linear(128, 1)\n        self.motor_x = nn.Linear(128, 1)\n        self.motor_y = nn.Linear(128, 1)\n\n    def forward(self, x):\n        features = self.backbone(x)\n        motor_visible = self.motor_visible(features)\n        motor_x = self.motor_x(features)\n        motor_y = self.motor_y(features)\n        return motor_visible, motor_x, motor_y","metadata":{"execution":{"iopub.status.busy":"2024-04-11T18:15:56.823499Z","iopub.status.idle":"2024-04-11T18:15:56.823855Z","shell.execute_reply.started":"2024-04-11T18:15:56.823684Z","shell.execute_reply":"2024-04-11T18:15:56.823698Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Define the loss function and optimizer\ncriterion = nn.BCEWithLogitsLoss()\nmse_criterion = nn.MSELoss()\nmodel = BacteriaCNN()\nmodel = model.to(device)\noptimizer = torch.optim.Adam(model.parameters(), lr=0.001)\n\n# Train the model\nnum_epochs = 10\nvalidate_every_n_epochs = 1\n\nfor epoch in range(num_epochs):\n    model.train()\n    running_loss = 0.0\n    for images, x, y, motor_visible in train_loader:\n        images = images.to(device)\n        motor_visible = torch.tensor(motor_visible, dtype=torch.float32, device=device)\n        x = torch.tensor(x, dtype=torch.float32, device=device)\n        y = torch.tensor(y, dtype=torch.float32, device=device)\n\n        optimizer.zero_grad()\n        motor_visible_pred, motor_x_pred, motor_y_pred = model(images)\n\n        loss_visible = criterion(motor_visible_pred.squeeze(), motor_visible)\n        loss_x = mse_criterion(motor_x_pred.squeeze(), x)\n        loss_y = mse_criterion(motor_y_pred.squeeze(), y)\n        loss = loss_visible + loss_x + loss_y\n\n        loss.backward()\n        optimizer.step()\n\n        running_loss += loss.item()\n\n    epoch_loss = running_loss / len(train_loader)\n    print(f\"Epoch {epoch + 1}, Training Loss: {epoch_loss:.4f}\")\n\n    if (epoch + 1) % validate_every_n_epochs == 0:\n        model.eval()\n        val_running_loss = 0.0\n        with torch.no_grad():\n            for images, x, y, motor_visible in val_loader:\n                images = images.to(device)\n                motor_visible = torch.tensor(motor_visible, dtype=torch.float32, device=device)\n                x = torch.tensor(x, dtype=torch.float32, device=device)\n                y = torch.tensor(y, dtype=torch.float32, device=device)\n\n                motor_visible_pred, motor_x_pred, motor_y_pred = model(images)\n                loss_visible = criterion(motor_visible_pred.squeeze(), motor_visible)\n                loss_x = mse_criterion(motor_x_pred.squeeze(), x)\n                loss_y = mse_criterion(motor_y_pred.squeeze(), y)\n                loss = loss_visible + loss_x + loss_y\n\n                val_running_loss += loss.item()\n\n        val_epoch_loss = val_running_loss / len(val_loader)\n        print(f\"Epoch {epoch + 1}, Validation Loss: {val_epoch_loss:.4f}\")\n\nprint(\"Training completed.\")\n\n# Save the model weights after training\nmodel_path = 'bacteria_cnn_weights.pth'\ntorch.save(model.state_dict(), model_path)\nprint(f'Model weights saved to {model_path}')","metadata":{"execution":{"iopub.status.busy":"2024-04-11T18:15:56.828108Z","iopub.status.idle":"2024-04-11T18:15:56.828507Z","shell.execute_reply.started":"2024-04-11T18:15:56.828332Z","shell.execute_reply":"2024-04-11T18:15:56.828351Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Create the competition submission file\ntest_dataset = BacteriaDataset(pd.read_csv(train_csv_path), val_images_dir, transform=transform)\ntest_loader = DataLoader(test_dataset, batch_size=32, shuffle=False, num_workers=4)\n\nsubmission = []\nmodel.eval()\nwith torch.no_grad():\n    for batch_idx, (images, _, _, _) in enumerate(test_loader):\n        images = images.to(device)\n        motor_visible_pred, motor_x_pred, motor_y_pred = model(images)\n\n        for i in range(images.size(0)):\n            image_name = test_dataset.data.iloc[batch_idx * test_loader.batch_size + i][\"image_name\"]\n            motor_visible_val = 1 if motor_visible_pred[i].item() > 0 else 0\n            motor_x_val = motor_x_pred[i].item()*512 if motor_visible_val == 1 else -1\n            motor_y_val = motor_y_pred[i].item()*512 if motor_visible_val == 1 else -1\n            submission.append({\n                \"image_name\": image_name,\n                \"x\": motor_x_val,\n                \"y\": motor_y_val,\n                \"motor_visible\": motor_visible_val\n            })\n\nsubmission_df = pd.DataFrame(submission)\nsubmission_df.to_csv(\"competition_submission.csv\", index=False)\nprint(\"Competition submission file created: competition_submission.csv\")","metadata":{"execution":{"iopub.status.busy":"2024-04-11T18:15:56.829657Z","iopub.status.idle":"2024-04-11T18:15:56.829968Z","shell.execute_reply.started":"2024-04-11T18:15:56.829802Z","shell.execute_reply":"2024-04-11T18:15:56.829815Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}