{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":91249,"databundleVersionId":11294684,"sourceType":"competition"}],"dockerImageVersionId":30919,"isInternetEnabled":true,"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\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader\nfrom tqdm import tqdm\nimport matplotlib.pyplot as plt\nimport seaborn as sns\n\nimport warnings\nwarnings.filterwarnings(\"ignore\")\n\n# Define the device\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(f\"Using device: {device}\")\n\n# Define data paths\ndata_path = \"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/\"\ntest_dir = os.path.join(data_path, \"test\")\ntrain_dir = os.path.join(data_path, \"train\")\ntrain_labels_path = os.path.join(data_path, \"train_labels.csv\")\nsubmission_path = os.path.join(data_path, \"submission.csv\")\noutput_path = 'submission.csv'\n\n# Load training labels\ntry:\n    train_labels_df = pd.read_csv(train_labels_path).drop_duplicates(subset='tomo_id', keep='first').set_index('tomo_id')\n    print(\"Train Labels DataFrame:\")\n    print(train_labels_df.head())\nexcept FileNotFoundError:\n    print(f\"Error: File not found at {train_labels_path}. Please ensure the dataset is correctly attached.\")\n    train_labels_df = None  # Or handle the error appropriately","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Custom Dataset class\nclass TomogramDataset(Dataset):\n    def __init__(self, data_dir, labels_df, resize=(64, 64), fixed_depth=32, is_test=False):\n        self.data_dir = data_dir\n        self.labels_df = labels_df\n        if is_test:\n            self.tomogram_ids = [f for f in os.listdir(data_dir) if os.path.isdir(os.path.join(data_dir, f))]\n        else:\n            self.tomogram_ids = list(labels_df.index)\n        self.resize = resize\n        self.fixed_depth = fixed_depth\n        self.is_test = is_test\n\n    def __len__(self):\n        return len(self.tomogram_ids)\n\n    def __getitem__(self, idx):\n        tomo_id = self.tomogram_ids[idx]\n        tomogram_path = os.path.join(self.data_dir, tomo_id)\n        tomogram = load_tomogram(tomogram_path, resize=self.resize, fixed_depth=self.fixed_depth)\n        if tomogram is None:\n            tomogram = np.zeros((1, self.fixed_depth, self.resize[0], self.resize[1]), dtype=np.float32)\n            labels = np.array([-1.0, -1.0, -1.0], dtype=np.float32)\n            return torch.tensor(tomogram).unsqueeze(0), torch.tensor(labels), tomo_id\n        tomogram = torch.tensor(tomogram).unsqueeze(0)\n        if not self.is_test:\n            labels = self.labels_df.loc[tomo_id, ['Motor axis 0', 'Motor axis 1', 'Motor axis 2']].values.astype(np.float32)\n        else:\n            labels = np.array([-1.0, -1.0, -1.0], dtype=np.float32)\n        return tomogram, torch.tensor(labels), tomo_id","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Model Definition\nclass Simple3DCNN(nn.Module):\n    def __init__(self, fixed_depth=32):\n        super(Simple3DCNN, self).__init__()\n        self.conv1 = nn.Conv3d(1, 16, kernel_size=3, padding=1)\n        self.relu1 = nn.ReLU()\n        self.pool1 = nn.MaxPool3d(kernel_size=(2, 2, 2), stride=2)\n        self.conv2 = nn.Conv3d(16, 32, kernel_size=3, padding=1)\n        self.relu2 = nn.ReLU()\n        self.pool2 = nn.MaxPool3d(kernel_size=(2, 2, 2), stride=2)\n        self.flatten = nn.Flatten()\n        self.flattened_size = 32 * (fixed_depth // 4) * (64 // 4) * (64 // 4)\n        self.fc1 = nn.Linear(self.flattened_size, 128)\n        self.relu3 = nn.ReLU()\n        self.fc2 = nn.Linear(128, 3)\n\n    def forward(self, x):\n        x = self.pool1(self.relu1(self.conv1(x)))\n        x = self.pool2(self.relu2(self.conv2(x)))\n        x = self.flatten(x)\n        x = self.relu3(self.fc1(x))\n        return self.fc2(x)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Training\ndef train_model(model, train_loader, criterion, optimizer, num_epochs=10, display_images=False):\n    model.train()\n    for epoch in range(num_epochs):\n        running_loss = 0.0\n        for batch_idx, (tomogram, labels, _) in enumerate(tqdm(train_loader, desc=f\"Epoch {epoch + 1}/{num_epochs} - Training\")):\n            if tomogram is None:\n                continue\n            tomogram, labels = tomogram.to(device), labels.to(device)\n            optimizer.zero_grad()\n            outputs = model(tomogram)\n            loss = criterion(outputs, labels)\n            loss.backward()\n            optimizer.step()\n            running_loss += loss.item()\n\n            # --- Pair Plot ---\n            if display_images and batch_idx == 0:\n                # Create a DataFrame for the pair plot\n                df = pd.DataFrame(labels.cpu().numpy(), columns=['Motor axis 0', 'Motor axis 1', 'Motor axis 2'])\n                sns.pairplot(df)\n                plt.show()\n\n            if display_images and batch_idx == 0:\n                fig, axes = plt.subplots(1, 5, figsize=(15, 5))\n                sample_tomo = tomogram[0].cpu().squeeze().numpy()\n                depth = sample_tomo.shape[0]\n                slice_indices = np.linspace(0, depth - 1, 5, dtype=int)\n                for i, idx in enumerate(slice_indices):\n                    axes[i].imshow(sample_tomo[idx], cmap='gray')\n                    axes[i].set_title(f\"Slice {idx}\")\n                    axes[i].axis('off')\n                plt.show()\n        print(f\"Epoch {epoch + 1}/{num_epochs}: Train Loss: {running_loss / len(train_loader):.4f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Evaluation Metric Implementation\ndef calculate_fbeta_score(predictions, ground_truths, threshold=1000, beta=2):\n    tp = 0\n    fp = 0\n    fn = 0\n    for i in range(len(ground_truths)):\n        gt = ground_truths[i]\n        pred = predictions[i]\n        motor_exists_gt = not all(gt == -1)\n        motor_predicted = not all(pred == -1)\n        if motor_exists_gt and motor_predicted:\n            distance = np.linalg.norm(gt - pred)\n            if distance <= threshold:\n                tp += 1\n            else:\n                fp += 1\n                fn += 1\n        elif motor_exists_gt and not motor_predicted:\n            fn += 1\n        elif not motor_exists_gt and motor_predicted:  # Corrected line\n            fp += 1\n    if tp + fp == 0:\n        precision = 0.0\n    else:\n        precision = tp / (tp + fp)\n    if tp + fn == 0:\n        recall = 0.0\n    else:\n        recall = tp / (tp + fn)\n    if precision + recall == 0:\n        fbeta = 0.0\n    else:\n        fbeta = (1 + beta**2) * (precision * recall) / ((beta**2 * precision) + recall)\n    return fbeta, precision, recall","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if __name__ == '__main__':\n    batch_size = 4\n    learning_rate = 0.00001  # Further reduced learning rate\n    num_epochs = 10  # Increase epochs\n    fixed_depth = 32\n    resize_dims = (64, 64)\n    threshold = 1000\n    beta = 2\n    CONFIDENCE_THRESHOLD = 0.45\n    MAX_DETECTIONS_PER_TOMO = 1\n\n    train_dataset = TomogramDataset(train_dir, train_labels_df, resize=resize_dims, fixed_depth=fixed_depth)\n    test_dataset = TomogramDataset(test_dir, train_labels_df, resize=resize_dims, fixed_depth=fixed_depth, is_test=True)\n\n    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)\n    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)\n\n    model = Simple3DCNN(fixed_depth).to(device)\n    criterion = nn.HuberLoss()  # Use Huber Loss\n    optimizer = optim.Adam(model.parameters(), lr=learning_rate)\n    scheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, 'min') #Add learning rate scheduler\n\n    def modified_train_model(model, train_loader, criterion, optimizer, num_epochs=10, display_images=False):\n        model.train()\n        for epoch in range(num_epochs):\n            running_loss = 0.0\n            for batch_idx, (tomogram, labels, _) in enumerate(tqdm(train_loader, desc=f\"Epoch {epoch + 1}/{num_epochs} - Training\")):\n                if tomogram is None:\n                    continue\n                tomogram, labels = tomogram.to(device), labels.to(device)\n                optimizer.zero_grad()\n                outputs = model(tomogram)\n                loss = criterion(outputs, labels)\n                loss.backward()\n\n                # Gradient Clipping\n                torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)\n\n                optimizer.step()\n                running_loss += loss.item()\n\n                if batch_idx % 10 == 0:\n                    print(f\"Epoch {epoch + 1}/{num_epochs}, Batch {batch_idx}: Loss = {loss.item()}\")\n\n                if display_images and batch_idx == 0:\n                    df = pd.DataFrame(labels.cpu().numpy(), columns=['Motor axis 0', 'Motor axis 1', 'Motor axis 2'])\n                    sns.pairplot(df)\n                    plt.show()\n\n                    fig, axes = plt.subplots(1, 5, figsize=(15, 5))\n                    sample_tomo = tomogram[0].cpu().squeeze().numpy()\n                    depth = sample_tomo.shape[0]\n                    slice_indices = np.linspace(0, depth - 1, 5, dtype=int)\n                    for i, idx in enumerate(slice_indices):\n                        axes[i].imshow(sample_tomo[idx], cmap='gray')\n                        axes[i].set_title(f\"Slice {idx}\")\n                        axes[i].axis('off')\n                    plt.show()\n\n            epoch_loss = running_loss / len(train_loader)\n            scheduler.step(epoch_loss) #step the scheduler\n            print(f\"Epoch {epoch + 1}/{num_epochs}: Train Loss: {epoch_loss:.4f}\")\n\n    modified_train_model(model, train_loader, criterion, optimizer, num_epochs, display_images=True)\n\n    predictions = predict(model, test_loader, CONFIDENCE_THRESHOLD, MAX_DETECTIONS_PER_TOMO)\n\n    create_submission_file(predictions, submission_path)\n\n    ground_truths = train_labels_df[['Motor axis 0', 'Motor axis 1', 'Motor axis 2']].values\n    predicted_values = []\n    tomo_ids_in_order = train_labels_df.index.tolist()\n    for tomo_id in tomo_ids_in_order:\n        predicted_values.append(predictions.get(tomo_id, [-1, -1, -1]))\n    predictions_array = np.array(predicted_values)\n\n    fbeta, precision, recall = calculate_fbeta_score(predictions_array, ground_truths, threshold, beta)\n\n    print(f\"F-beta score: {fbeta:.4f}\")\n    print(f\"Precision: {precision:.4f}\")\n    print(f\"Recall: {recall:.4f}\")\n\n    precisions = [precision]\n    recalls = [recall]\n    plt.figure(figsize=(8, 6))\n    plt.plot(recalls, precisions, marker='o', linestyle='-')\n    plt.xlabel(\"Recall\")\n    plt.ylabel(\"Precision\")\n    plt.title(\"Precision-Recall Curve\")\n    plt.xlim([0.0, 1.0])\n    plt.ylim([0.0, 1.0])\n    plt.grid(True)\n    plt.show()\n\n    thresholds = np.arange(500, 2000, 100)\n    fbeta_scores = []\n    for t in thresholds:\n        fbeta_score, _, _ = calculate_fbeta_score(predictions_array, ground_truths, t, beta)\n        fbeta_scores.append(fbeta_score)\n\n    plt.figure(figsize=(8, 6))\n    plt.plot(thresholds, fbeta_scores, marker='o', linestyle='-')\n    plt.xlabel(\"Threshold\")\n    plt.ylabel(\"F-beta Score\")\n    plt.title(\"F-beta Score vs. Threshold\")\n    plt.grid(True)\n    plt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Prediction and Submission\n\ndef predict(model, test_loader, confidence_threshold=0.45, max_detections_per_tomo=1):\n    model.eval()\n    predictions = {}\n    with torch.no_grad():\n        for tomogram, _, tomo_ids in tqdm(test_loader, desc=\"Predicting\"):\n            if tomogram is None:\n                print(f\"Skipping tomo_ids {tomo_ids} due to None tomogram.\")\n                continue\n            outputs = model(tomogram.to(device)).cpu().numpy()\n            for i, tomo_id in enumerate(tomo_ids):\n                prediction = outputs[i]\n                predictions[tomo_id] = prediction\n    return predictions\n\ndef create_submission_file(predictions, submission_path, output_path='submission.csv'):\n    try:\n        submission = pd.read_csv(submission_path)\n    except FileNotFoundError:\n        print(f\"Error: Submission file not found at {submission_path}\")\n        test_tomos = sorted([d for d in os.listdir(test_dir) if os.path.isdir(os.path.join(test_dir, d))])\n        submission = pd.DataFrame({'tomo_id': test_tomos, 'Motor axis 0': -1, 'Motor axis 1': -1, 'Motor axis 2': -1})\n        submission[['Motor axis 0', 'Motor axis 1', 'Motor axis 2']] = submission[['Motor axis 0', 'Motor axis 1', 'Motor axis 2']].astype(float)\n        print(\"Created a dummy submission file because the original was not found.\")\n\n    if 'Motor axis 0' not in submission.columns:\n        submission['Motor axis 0'] = -1.0\n    if 'Motor axis 1' not in submission.columns:\n        submission['Motor axis 1'] = -1.0\n    if 'Motor axis 2' not in submission.columns:\n        submission['Motor axis 2'] = -1.0\n    submission[['Motor axis 0', 'Motor axis 1', 'Motor axis 2']] = submission[['Motor axis 0', 'Motor axis 1', 'Motor axis 2']].astype(float)\n\n    for tomo_id, coords in predictions.items():\n        try:\n            submission.loc[submission['tomo_id'] == tomo_id, ['Motor axis 0', 'Motor axis 1', 'Motor axis 2']] = coords\n        except KeyError:\n            print(f\"Warning: tomo_id {tomo_id} not found in submission file.\")\n    submission.to_csv(output_path, index=False)\n    print(f\"Submission file created/updated at {output_path}\")\n    print(\"\\nSubmission file head:\")\n    print(submission.head())\n    import sys\n    sys.stdout.flush()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Assuming you have already run the previous cells and have access to:\n# - train_labels_df\n# - predictions (dictionary of predictions)\n# - predictions_array (numpy array of predicted values)\n# - ground_truths (numpy array of ground truth values)\n# - fbeta, precision, recall (evaluation metrics)\n\n# --- Graphical Dump ---\n\n# 1. Pie Charts (Distribution of Number of Motors)\nmotor_counts = train_labels_df['Number of motors'].value_counts()\nplt.figure(figsize=(8, 6))\nplt.pie(motor_counts, labels=motor_counts.index, autopct='%1.1f%%', startangle=140)\nplt.title('Distribution of Number of Motors')\nplt.show()\n\n# 2. Pair Plot (Ground Truth Motor Coordinates)\nsns.pairplot(train_labels_df[['Motor axis 0', 'Motor axis 1', 'Motor axis 2']])\nplt.suptitle('Pair Plot of Ground Truth Motor Coordinates', y=1.02)\nplt.show()\n\n# 3. KDE Plot (Distribution of Predicted vs. Ground Truth Motor Axis 0)\nplt.figure(figsize=(10, 6))\nsns.kdeplot(ground_truths[:, 0], label='Ground Truth Axis 0', fill=True)\nsns.kdeplot(predictions_array[:, 0], label='Predicted Axis 0', fill=True)\nplt.title('KDE Plot of Predicted vs. Ground Truth Motor Axis 0')\nplt.legend()\nplt.show()\n\n# 4. Pair Grid (Comparison of Predicted vs. Ground Truth)\ndf_comparison = pd.DataFrame({\n    'Ground Truth Axis 0': ground_truths[:, 0],\n    'Predicted Axis 0': predictions_array[:, 0],\n    'Ground Truth Axis 1': ground_truths[:, 1],\n    'Predicted Axis 1': predictions_array[:, 1],\n    'Ground Truth Axis 2': ground_truths[:, 2],\n    'Predicted Axis 2': predictions_array[:, 2],\n})\ng = sns.PairGrid(df_comparison)\ng.map_diag(sns.histplot)\ng.map_offdiag(sns.scatterplot)\ng.fig.suptitle('Pair Grid: Predicted vs. Ground Truth Motor Coordinates', y=1.02)\nplt.show()\n\n# 5. Relplot (Scatter Plot of Predicted vs. Ground Truth Axis 1)\nplt.figure(figsize=(10, 6))\nsns.relplot(x=ground_truths[:, 1], y=predictions_array[:, 1], kind='scatter')\nplt.title('Relplot: Predicted vs. Ground Truth Motor Axis 1')\nplt.show()\n\n# 6. Catplot (Distribution of Voxel Spacing)\nplt.figure(figsize=(10, 6))\nsns.catplot(x='Voxel spacing', kind='count', data=train_labels_df)\nplt.title('Catplot: Distribution of Voxel Spacing')\nplt.show()\n\n# 7. Density Plot (2D Density of Predicted vs. Ground Truth Axis 2)\nplt.figure(figsize=(10, 6))\nsns.kdeplot(x=ground_truths[:, 2], y=predictions_array[:, 2], fill=True)\nplt.title('2D Density Plot: Predicted vs. Ground Truth Motor Axis 2')\nplt.show()\n\n# 8. Area Chart (Distribution of Array Shape Axis 0)\narray_shape_counts = train_labels_df['Array shape (axis 0)'].value_counts().sort_index()\nplt.figure(figsize=(10, 6))\nplt.fill_between(array_shape_counts.index, array_shape_counts.values)\nplt.title('Area Chart: Distribution of Array Shape (Axis 0)')\nplt.xlabel('Array Shape (Axis 0)')\nplt.ylabel('Count')\nplt.show()\n\n# 9. fbeta vs threshold\nthresholds = np.arange(500, 2000, 100)\nfbeta_scores = []\nbeta = 2\nfor t in thresholds:\n    fbeta_score, _, _ = calculate_fbeta_score(predictions_array, ground_truths, t, beta)\n    fbeta_scores.append(fbeta_score)\n\nplt.figure(figsize=(8, 6))\nplt.plot(thresholds, fbeta_scores, marker='o', linestyle='-')\nplt.xlabel(\"Threshold\")\nplt.ylabel(\"F-beta Score\")\nplt.title(\"F-beta Score vs. Threshold\")\nplt.grid(True)\nplt.show()\n\n# 10. Precision recall curve\nprecisions = [precision]\nrecalls = [recall]\nplt.figure(figsize=(8, 6))\nplt.plot(recalls, precisions, marker='o', linestyle='-')\nplt.xlabel(\"Recall\")\nplt.ylabel(\"Precision\")\nplt.title(\"Precision-Recall Curve\")\nplt.xlim([0.0, 1.0])\nplt.ylim([0.0, 1.0])\nplt.grid(True)\nplt.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}