{"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"},{"sourceId":297065,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":254225,"modelId":275646}],"dockerImageVersionId":30919,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# NEW DeepBac Motor Finder – Decoding Bacterial Mysteries with ResNet50 Transfer Learning\n\nWelcome to **NEW DeepBac Motor Finder**, your cutting-edge tool for the **BYU Locating Bacterial Flagellar Motors** competition on Kaggle. This notebook is designed to revolutionize the way you explore the microscopic universe by leveraging the power of **ResNet50** and state-of-the-art deep learning techniques—all executed completely offline.\n\n## 🔄 Overview\n\n### ⚡ Setup & Environment\n- **Robust Environment & Library Setup:**  \n  We initialize a high-performance computing environment, importing essential libraries and configuring the device (GPU/CPU) for optimal speed. With Kaggle’s powerful resources and local model caching, our pipeline operates seamlessly even without internet access.\n- **Offline Model Initialization:**  \n  Our system loads the ResNet50 model offline from a pre-saved `.pth` file, ensuring that all operations run reliably and swiftly.\n- **Optimized for GPU Acceleration:**  \n  For best performance, this notebook is designed to run on a **P100 GPU**. While CPU execution is possible, a GPU significantly enhances speed and efficiency.\n\n### 🔍 Data Loading & Preprocessing\n- **Data Paths & Label Management:**  \n  Define precise file paths and load training labels while filtering out any invalid entries. This ensures only high-quality data is used for training.\n- **Custom Dataset Handling:**  \n  Our custom Dataset class efficiently processes grayscale tomogram slices by converting them into 3-channel images—making them perfectly compatible with our deep learning model.\n\n### 🌀 Data Augmentation\n- **Innovative Transformations:**  \n  Experience robust data augmentation that transforms your dataset dynamically:\n  - **RandomResizedCrop:** Simulates different scales and perspectives.\n  - **RandomHorizontalFlip & RandomRotation:** Introduce variability in orientation.\n  - **ColorJitter:** Fine-tunes brightness and contrast to enhance image features.\n- These augmentations significantly improve the model’s ability to generalize and predict with higher accuracy.\n\n### 💪 Model Building & Fine-Tuning\n- **ResNet50 Fine-Tuning:**  \n  We adapt the powerful ResNet50 model by replacing its final classification layer with a custom regression head—comprising multiple fully connected layers, ReLU activations, and dropout layers. This specialized architecture is tailored to precisely predict the 3D coordinates of bacterial motors.\n- **Optimized Training:**  \n  Our model is fine-tuned using advanced optimization techniques and learning rate scheduling to ensure rapid convergence and superior performance.\n\n### 🔥 Training, Evaluation & Submission\n- **Streamlined Training Process:**  \n  Enjoy an efficient training pipeline with accelerated DataLoaders, real-time progress tracking via `tqdm`, and smart scheduling strategies.\n- **Accurate Prediction Aggregation:**  \n  Predictions are intelligently aggregated per tomogram slice, ensuring that the final submission is both precise and formatted according to Kaggle’s requirements.\n- **Competition-Ready Output:**  \n  The notebook exports a `submission.csv` file that meets the competition’s format—making your submission as polished as your model’s performance.\n\n---\n\n**Embark on a journey of discovery and innovation!** With NEW DeepBac Motor Finder, each tomogram unveils hidden secrets of bacterial motors. Whether you're pushing the boundaries of scientific exploration or honing your deep learning prowess, this notebook is your ultimate launchpad.\n\n**Let’s unlock the mysteries together—may your model shine and lead the way to victory! 🌟**","metadata":{}},{"cell_type":"markdown","source":"## Import Libraries, Set Device, and Random Seed","metadata":{}},{"cell_type":"code","source":"import random\nimport numpy as np\nimport pandas as pd\nimport torch\nimport os\nfrom torch import nn, optim\nfrom torch.utils.data import Dataset, DataLoader, ConcatDataset\nfrom torchvision import models, transforms\nfrom torchvision.models import resnet50, ResNet50_Weights\nimport cv2\nfrom PIL import Image, ImageFilter\nfrom tqdm import tqdm\nimport matplotlib.pyplot as plt\n\n# Set device (GPU if available, else CPU)\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Using device: {device}\")\n\n# Set random seed for reproducibility\nSEED = 42\nrandom.seed(SEED)\nnp.random.seed(SEED)\ntorch.manual_seed(SEED)\nif torch.cuda.is_available():\n    torch.cuda.manual_seed_all(SEED)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-03-23T03:44:20.723234Z","iopub.execute_input":"2025-03-23T03:44:20.723519Z","iopub.status.idle":"2025-03-23T03:44:27.453641Z","shell.execute_reply.started":"2025-03-23T03:44:20.723498Z","shell.execute_reply":"2025-03-23T03:44:27.452835Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Load ResNet50 Model with Offline Weights","metadata":{}},{"cell_type":"code","source":"print(\"Loading ResNet50 model with offline weights...\")\nweights_path = \"/kaggle/input/resnet50/pytorch/default/1/resnet50_weights.pth\"\nif os.path.exists(weights_path):\n    model = resnet50()  # Initialize ResNet50\n    model.load_state_dict(torch.load(weights_path, map_location=device))\n    print(\"Model loaded successfully with offline ResNet50 weights!\")\nelse:\n    print(\"Weights file not found. Loading pretrained ResNet50 from torchvision...\")\n    model = resnet50(weights=ResNet50_Weights.DEFAULT)\n    print(\"Model loaded with pretrained weights from torchvision.\")\n\n# Modify the first convolution layer if needed (keeping the original MaxPool)\nmodel.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)\nmodel = model.to(device)\nprint(\"Cell 2: Completed\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-23T03:44:27.454609Z","iopub.execute_input":"2025-03-23T03:44:27.455032Z","iopub.status.idle":"2025-03-23T03:44:29.323832Z","shell.execute_reply.started":"2025-03-23T03:44:27.455007Z","shell.execute_reply":"2025-03-23T03:44:29.323040Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Set Paths for Directories and Files","metadata":{}},{"cell_type":"code","source":"base_path = \"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025\"\ntrain_dir = os.path.join(base_path, \"train\")\ntest_dir = os.path.join(base_path, \"test\")\ntrain_labels_path = os.path.join(base_path, \"train_labels.csv\")\nsample_submission_path = os.path.join(base_path, \"sample_submission.csv\")\n\nprint(\"Base directory contents:\", os.listdir(base_path))\nprint(\"Cell 3: Completed\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-23T03:44:29.325207Z","iopub.execute_input":"2025-03-23T03:44:29.325442Z","iopub.status.idle":"2025-03-23T03:44:29.330835Z","shell.execute_reply.started":"2025-03-23T03:44:29.325420Z","shell.execute_reply":"2025-03-23T03:44:29.330056Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"##  Read and Filter Labels File","metadata":{}},{"cell_type":"code","source":"train_labels = pd.read_csv(train_labels_path)\n# Filter out rows with -1 values in motor axes\ntrain_labels = train_labels[(train_labels[\"Motor axis 0\"] != -1.0) &\n                            (train_labels[\"Motor axis 1\"] != -1.0) &\n                            (train_labels[\"Motor axis 2\"] != -1.0)]\nprint(\"Number of valid samples after filtering:\", len(train_labels))\nprint(\"Cell 4: Completed\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-23T03:44:29.331922Z","iopub.execute_input":"2025-03-23T03:44:29.332138Z","iopub.status.idle":"2025-03-23T03:44:29.369738Z","shell.execute_reply.started":"2025-03-23T03:44:29.332119Z","shell.execute_reply":"2025-03-23T03:44:29.369048Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Define the Training and Test Dataset Classes","metadata":{}},{"cell_type":"code","source":"class BacterialDataset(Dataset):\n    def __init__(self, root_dir, labels_df, transform=None):\n        self.root_dir = root_dir\n        self.transform = transform\n        # Map each tomo_id to its motor coordinates\n        self.labels_dict = {\n            str(row[\"tomo_id\"]).lower(): (row[\"Motor axis 0\"], row[\"Motor axis 1\"], row[\"Motor axis 2\"])\n            for _, row in labels_df.iterrows()\n        }\n        # Get image file paths from each valid tomo directory\n        self.filepaths = [\n            os.path.join(root, f)\n            for tomo_id in self.labels_dict.keys()\n            if os.path.exists(os.path.join(root_dir, tomo_id))\n            for root, _, files in os.walk(os.path.join(root_dir, tomo_id))\n            for f in files if f.lower().endswith('.jpg')\n        ]\n        print(\"Number of training images found:\", len(self.filepaths))\n\n    def __len__(self):\n        return len(self.filepaths)\n\n    def __getitem__(self, idx):\n        img_path = self.filepaths[idx]\n        image = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)\n        if image is None:\n            raise ValueError(f\"Failed to load image from {img_path}\")\n        image = Image.fromarray(image)\n        if self.transform:\n            image = self.transform(image)\n        # Extract tomo_id from the parent directory name (in lowercase)\n        tomo_id = os.path.basename(os.path.dirname(img_path)).lower()\n        if tomo_id not in self.labels_dict:\n            raise ValueError(f\"tomo_id {tomo_id} not found in labels dictionary.\")\n        label = torch.tensor(self.labels_dict[tomo_id], dtype=torch.float32)\n        return image, label\n\n\nclass TestDataset(Dataset):\n    def __init__(self, root_dir, transform=None):\n        self.root_dir = root_dir\n        self.transform = transform\n        # Get image file paths from all directories\n        self.filepaths = [\n            os.path.join(root, f)\n            for root, _, files in os.walk(root_dir)\n            for f in files if f.lower().endswith('.jpg')\n        ]\n        print(\"Number of test images found:\", len(self.filepaths))\n\n    def __len__(self):\n        return len(self.filepaths)\n\n    def __getitem__(self, idx):\n        img_path = self.filepaths[idx]\n        image = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)\n        if image is None:\n            raise ValueError(f\"Failed to load image from {img_path}\")\n        image = Image.fromarray(image)\n        if self.transform:\n            image = self.transform(image)\n        tomogram_id = os.path.basename(os.path.dirname(img_path))\n        return image, tomogram_id\n\nprint(\"Cell 5: Completed\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-23T03:44:29.370431Z","iopub.execute_input":"2025-03-23T03:44:29.370630Z","iopub.status.idle":"2025-03-23T03:44:29.380309Z","shell.execute_reply.started":"2025-03-23T03:44:29.370612Z","shell.execute_reply":"2025-03-23T03:44:29.379559Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Define Image Transforms","metadata":{}},{"cell_type":"code","source":"train_transform = transforms.Compose([\n    transforms.Resize((256, 256), interpolation=Image.LANCZOS),\n    transforms.Lambda(lambda x: x.filter(ImageFilter.SHARPEN)),\n    transforms.RandomResizedCrop(224, interpolation=Image.LANCZOS),\n    transforms.RandomHorizontalFlip(),\n    transforms.RandomRotation(10),\n    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1),\n    transforms.ToTensor(),\n    transforms.Lambda(lambda x: x.repeat(3, 1, 1)),\n    transforms.Normalize(mean=[0.485, 0.456, 0.406],\n                         std=[0.229, 0.224, 0.225])\n])\n\ntest_transform = transforms.Compose([\n    transforms.Resize((224, 224), interpolation=Image.LANCZOS),\n    transforms.Lambda(lambda x: x.filter(ImageFilter.SHARPEN)),\n    transforms.ToTensor(),\n    transforms.Lambda(lambda x: x.repeat(3, 1, 1)),\n    transforms.Normalize(mean=[0.485, 0.456, 0.406],\n                         std=[0.229, 0.224, 0.225])\n])\nprint(\"Cell 6: Completed\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-23T03:44:29.381100Z","iopub.execute_input":"2025-03-23T03:44:29.381424Z","iopub.status.idle":"2025-03-23T03:44:29.397618Z","shell.execute_reply.started":"2025-03-23T03:44:29.381402Z","shell.execute_reply":"2025-03-23T03:44:29.396933Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Create Datasets and DataLoaders","metadata":{}},{"cell_type":"code","source":"train_dataset = BacterialDataset(train_dir, train_labels, transform=train_transform)\ntest_dataset = TestDataset(test_dir, transform=test_transform)\n\nprint(\"Processing training dataset file paths...\")\nfor i in tqdm(range(len(train_dataset)), desc=\"Training dataset loaded\", total=len(train_dataset)):\n    pass\nprint(f\"Training dataset loaded with {len(train_dataset)} images.\")\n\nprint(\"Processing test dataset file paths...\")\nfor i in tqdm(range(len(test_dataset)), desc=\"Test dataset loaded\", total=len(test_dataset)):\n    pass\nprint(f\"Test dataset loaded with {len(test_dataset)} images.\")\n\n# Optionally, use ConcatDataset to repeat training data if needed\ntrain_dataset = ConcatDataset([train_dataset] * 1)\nprint(f\"After concatenation, training dataset has {len(train_dataset)} samples.\")\n\ntrain_loader = DataLoader(train_dataset, batch_size=8, shuffle=True, \n                          num_workers=4, pin_memory=True, prefetch_factor=2, persistent_workers=True)\ntest_loader = DataLoader(test_dataset, batch_size=1, shuffle=False, \n                         num_workers=2, pin_memory=True, prefetch_factor=2, persistent_workers=True)\nprint(\"Cell 7: Completed\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-23T03:44:29.398357Z","iopub.execute_input":"2025-03-23T03:44:29.398701Z","iopub.status.idle":"2025-03-23T03:48:05.940751Z","shell.execute_reply.started":"2025-03-23T03:44:29.398680Z","shell.execute_reply":"2025-03-23T03:48:05.939889Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Modify the Model’s Fully Connected Layers and Set Criterion, Optimizer, and Scheduler","metadata":{}},{"cell_type":"code","source":"num_ftrs = model.fc.in_features\n\n# Modify the final layers of the model\nmodel.fc = nn.Sequential(\n    nn.Linear(num_ftrs, 1024),\n    nn.ReLU(),\n    nn.Dropout(0.5),\n    nn.Linear(1024, 512),\n    nn.ReLU(),\n    nn.Dropout(0.5),\n    nn.Linear(512, 3)\n)\nmodel = model.to(device)\nprint(\"Model architecture (ResNet50):\")\nprint(model)\n\ncriterion = nn.MSELoss()\noptimizer = optim.Adam(model.parameters(), lr=0.0001)\nscheduler = optim.lr_scheduler.StepLR(optimizer, step_size=5, gamma=0.1)\nprint(\"Cell 8: Completed\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-23T03:48:05.942813Z","iopub.execute_input":"2025-03-23T03:48:05.943035Z","iopub.status.idle":"2025-03-23T03:48:05.977412Z","shell.execute_reply.started":"2025-03-23T03:48:05.943016Z","shell.execute_reply":"2025-03-23T03:48:05.976718Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Training Loop (Adjust the Number of Epochs by Changing num_epochs)","metadata":{}},{"cell_type":"code","source":"# You can increase the number of epochs to improve training; currently num_epochs is set to 1 as a test run.\nnum_epochs = 3\nprint(\"Starting training using ResNet50...\")\n\nfor epoch in range(num_epochs):\n    running_loss = 0.0\n    model.train()\n    for images, labels in tqdm(train_loader, desc=f\"Epoch {epoch+1}/{num_epochs}\"):\n        images, labels = images.to(device), labels.to(device)\n        optimizer.zero_grad()\n        outputs = model(images)\n        loss = criterion(outputs, labels)\n        loss.backward()\n        optimizer.step()\n        running_loss += loss.item()\n    scheduler.step()\n    avg_loss = running_loss / len(train_loader)\n    print(f\"Epoch {epoch+1}/{num_epochs} - Loss: {avg_loss:.4f}\")\n\nprint(\"Training complete!\")\nprint(\"Cell 9: Completed\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-23T03:48:05.978349Z","iopub.execute_input":"2025-03-23T03:48:05.978655Z","iopub.status.idle":"2025-03-23T04:38:03.416703Z","shell.execute_reply.started":"2025-03-23T03:48:05.978624Z","shell.execute_reply":"2025-03-23T04:38:03.414913Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"##  Evaluation and Submission File Creation","metadata":{}},{"cell_type":"code","source":"tomogram_predictions = {}\nmodel.eval()\nwith torch.no_grad():\n    for image, tomogram_id in tqdm(test_loader, desc=\"Testing\"):\n        image = image.to(device)\n        output = model(image)[0].tolist()\n        # Ensure tomogram_id is a string (if a list, take the first element)\n        tomogram_id = tomogram_id[0] if isinstance(tomogram_id, list) else tomogram_id\n        if tomogram_id not in tomogram_predictions:\n            tomogram_predictions[tomogram_id] = [output]\n        else:\n            tomogram_predictions[tomogram_id].append(output)\n\nfinal_predictions = []\nfor tomo_id, preds in tomogram_predictions.items():\n    preds_array = np.array(preds)\n    avg_coords = preds_array.mean(axis=0)\n    # If the average coordinates are less than -0.5, consider that no motor is present\n    if np.mean(avg_coords) < -0.5:\n        final_predictions.append([tomo_id, -1, -1, -1])\n    else:\n        final_predictions.append([tomo_id, avg_coords[0], avg_coords[1], avg_coords[2]])\n\nsubmission_df = pd.DataFrame(final_predictions, columns=[\"tomo_id\", \"Motor axis 0\", \"Motor axis 1\", \"Motor axis 2\"])\nsubmission_df = submission_df.fillna(method='ffill')\nsubmission_df = submission_df.applymap(lambda x: -1 if (isinstance(x, (int, float)) and (x < -1e5 or x > 1e5)) else x)\nsubmission_df = submission_df[[\"tomo_id\", \"Motor axis 0\", \"Motor axis 1\", \"Motor axis 2\"]]\n\nsubmission_df.to_csv(\"submission.csv\", index=False)\nprint(\"Submission file created successfully!\")\nprint(\"First 5 predictions:\")\nprint(submission_df.head())\nprint(\"Cell 9: Completed\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-23T04:38:03.421012Z","iopub.execute_input":"2025-03-23T04:38:03.421294Z","iopub.status.idle":"2025-03-23T04:38:24.119644Z","shell.execute_reply.started":"2025-03-23T04:38:03.421263Z","shell.execute_reply":"2025-03-23T04:38:24.118531Z"}},"outputs":[],"execution_count":null}]}