{"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":"nvidiaTeslaT4","dataSources":[{"sourceId":91249,"databundleVersionId":11294684,"sourceType":"competition"}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Install required packages\n!pip install retrying monai scipy scikit-image wandb imageio gcsfs\n\n# Import libraries\nimport os\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nfrom monai.networks.nets import UNet\nfrom monai.losses import DiceLoss\nimport torch.optim as optim\nfrom scipy.ndimage import gaussian_filter, center_of_mass\nfrom scipy.signal import find_peaks\nimport sklearn.metrics\nfrom sklearn.model_selection import train_test_split\nimport matplotlib.pyplot as plt\nfrom tqdm import tqdm\nimport glob\nfrom IPython.display import Video, display\nimport wandb\nimport time\nimport shutil\nimport gcsfs  # Dataset is too big for kaggle - After multiple attempts of uploading, I failed miserably. \nimport gc\nimport threading\nimport logging\nimport sys\nfrom concurrent.futures import ThreadPoolExecutor, as_completed\nimport queue\nfrom retrying import retry\nfrom concurrent.futures import ThreadPoolExecutor, as_completed\nimport torch.nn.functional as F  # Add this import\n# Set random seed for reproducibility\ntorch.manual_seed(42)\nnp.random.seed(42)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Ensure the output directory exists\nos.makedirs('/kaggle/working', exist_ok=True)\n\n# Configure root logger\nlogger = logging.getLogger('')\nlogger.setLevel(logging.INFO)  # Capture INFO and above\n\n# Clear any existing handlers to avoid conflicts\nlogger.handlers = []\n\n# File handler for detailed logs (INFO, WARNING, ERROR)\nfile_handler = logging.FileHandler(\"/kaggle/working/training.log\", mode='w')\nfile_handler.setLevel(logging.INFO)\nfile_handler.setFormatter(logging.Formatter('%(asctime)s - %(levelname)s - %(message)s'))\nlogger.addHandler(file_handler)\n\n# Console handler for errors only\nconsole_handler = logging.StreamHandler(sys.stdout)\nconsole_handler.setLevel(logging.ERROR)\nconsole_handler.setFormatter(logging.Formatter('%(levelname)s - %(message)s'))\nlogger.addHandler(console_handler)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Initialize wandb\nfrom kaggle_secrets import UserSecretsClient\nuser_secrets = UserSecretsClient()\nwb_token = user_secrets.get_secret(\"WANDB\")\nwandb.login(key=wb_token)\nwandb.init(\n    project=\"byu-bacterial-flagellar-motors\",\n    config={\n        \"learning_rate\": 1e-3,\n        \"epochs\": 50,\n        \"batch_size\": 4,\n        \"patch_size\": (128, 128, 128),\n        \"gaussian_sigma\": 5,\n        \"architecture\": \"3D U-Net\",\n        \"optimizer\": \"Adam\",\n        \"loss_function\": \"DiceLoss\",\n        \"beta\": 2\n    }\n)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"fs = gcsfs.GCSFileSystem(token=\"anon\") # Initialize GCS filesystem\nlogger.info(\"Initialized GCS filesystem with anonymous access\")\n\n# Define GCS path and local directory\ngcs_precomputed_path = \"gs://nb153/precomputedmasks\"\ngcs_preprocessed_path = \"gs://nb153/preprocessed\"\nlocal_dir = \"/kaggle/working/data\"\nos.makedirs(local_dir, exist_ok=True)\n\n# Device configuration\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 1. Analyze dataset to identify tomograms with motors\ndef identify_motor_tomograms(labels_df):\n    \"\"\"Identify tomograms with valid motor annotations.\"\"\"\n    motor_tomograms = []\n    for tomo_id in labels_df[\"tomo_id\"].unique():\n        tomo_labels = labels_df[labels_df[\"tomo_id\"] == tomo_id].iloc[0]\n        if tomo_labels[\"Number of motors\"] > 0 and tomo_labels[\"Motor axis 0\"] != -1:\n            motor_tomograms.append(tomo_id)\n    logger.info(f\"Found {len(motor_tomograms)} tomograms with motors\")\n    return motor_tomograms\n\n# Load labels and identify motor tomograms\nlabels_df = pd.read_csv(\"/kaggle/input/byu-locating-bacterial-flagellar-motors-2025/train_labels.csv\")\ntomo_ids = sorted(labels_df[\"tomo_id\"].unique())\nmotor_tomo_ids = identify_motor_tomograms(labels_df)\nlogger.info(f\"Total tomograms: {len(tomo_ids)}, Motor tomograms: {len(motor_tomo_ids)}\")\n\n# Split into train/val/test (80/10/10) using only motor tomograms for training\ntrain_val_ids, _ = train_test_split(tomo_ids, test_size=0.1, random_state=42)\ntrain_ids, val_ids = train_test_split(train_val_ids, test_size=0.1111, random_state=42)\ntrain_ids = [tid for tid in train_ids if tid in motor_tomo_ids]\ntest_ids = [\"tomo_00e047\", \"tomo_01a877\"]  # tomo_003acc Restored all test tomograms\nlogger.info(f\"Train IDs: {len(train_ids)}, Val IDs: {len(val_ids)}, Test IDs: {len(test_ids)}\")\nlogger.info(f\"Test IDs: {test_ids}\")\n\n# Analyze tomograms\ndef analyze_tomograms(labels_df):\n    total_tomograms = len(labels_df[\"tomo_id\"].unique())\n    motor_tomograms = len(labels_df[labels_df[\"Number of motors\"] > 0])\n    non_motor_tomograms = total_tomograms - motor_tomograms\n    logger.info(f\"Total tomograms: {total_tomograms}\")\n    logger.info(f\"Tomograms with motors: {motor_tomograms}\")\n    logger.info(f\"Tomograms without motors: {non_motor_tomograms}\")\n    return total_tomograms, motor_tomograms, non_motor_tomograms\n\n# Plot tomogram distribution\n# def plot_tomogram_distribution(labels_df):\n#     total_tomograms = len(labels_df[\"tomo_id\"].unique())\n#     motor_tomograms = len(labels_df[labels_df[\"Number of motors\"] > 0])\n#     non_motor_tomograms = total_tomograms - motor_tomograms\n#     plt.figure(figsize=(8, 5))\n#     plt.bar([\"Total\", \"With Motors\", \"Without Motors\"], [total_tomograms, motor_tomograms, non_motor_tomograms], color=[\"blue\", \"green\", \"red\"])\n#     plt.title(\"Tomogram Distribution\")\n#     plt.ylabel(\"Count\")\n#     for i, count in enumerate([total_tomograms, motor_tomograms, non_motor_tomograms]):\n#         plt.text(i, count + 0.5, str(count), ha=\"center\")\n#     plt.show()\n#     plt.close()\n\n# Download functions with simultaneous downloading\n\n# Optimized download function with robust error handling\n@retry(stop_max_attempt_number=3, wait_fixed=2000)\ndef download_file(gcs_path, local_path):\n    if not os.path.exists(local_path):\n        try:\n            fs.get(gcs_path, local_path)\n            logger.info(f\"Downloaded {gcs_path} to {local_path}\")\n        except Exception as e:\n            if \"404\" in str(e):\n                logger.warning(f\"GCS file not found: {gcs_path}\")\n                raise FileNotFoundError(f\"GCS file not found: {gcs_path}\")\n            else:\n                raise\n    return local_path\n\n\ndef download_npy_and_mask(tomo_id, gcs_preprocessed_path, gcs_precomputed_path, split, local_dir):\n    \"\"\"Download tomogram and mask in parallel with retry.\"\"\"\n    tomo_path = os.path.join(local_dir, f\"{tomo_id}.npy\")\n    mask_path = os.path.join(local_dir, f\"{tomo_id}_mask.npy\")\n    try:\n        with ThreadPoolExecutor(max_workers=2) as executor:\n            npy_future = executor.submit(\n                download_file,\n                f\"{gcs_preprocessed_path}/{split}/{tomo_id}/{tomo_id}.npy\",\n                tomo_path\n            )\n            mask_future = executor.submit(\n                download_file,\n                f\"{gcs_precomputed_path}/{split}/{tomo_id}_mask.npy\",\n                mask_path\n            )\n            tomo_path = npy_future.result()\n            mask_path = mask_future.result()\n        logger.info(f\"Completed parallel download for tomogram {tomo_id}\")\n        return tomo_path, mask_path\n    except FileNotFoundError as e:\n        logger.error(f\"Skipping tomogram {tomo_id}: {str(e)}\")\n        return None, None\n    except Exception as e:\n        logger.error(f\"Error downloading tomogram {tomo_id}: {str(e)}\")\n        raise\n\n\n# Combined download and preprocess function\ndef download_and_preprocess(tomo_id, gcs_preprocessed_path, gcs_precomputed_path, split, local_dir, labels_df, patches_per_volume):\n    try:\n        tomo_path, mask_path = download_npy_and_mask(tomo_id, gcs_preprocessed_path, gcs_precomputed_path, split, local_dir)\n        volume = np.load(tomo_path)\n        mask = np.load(mask_path)\n        tomo_labels = labels_df[labels_df[\"tomo_id\"] == tomo_id].iloc[0]\n        num_motors = tomo_labels[\"Number of motors\"]\n        logger.info(f\"Tomogram {tomo_id} has {num_motors} motors\")\n        patches, mask_patches = sample_patches(tomo_id, volume, mask, labels_df, patches_per_volume=patches_per_volume)\n        del volume, mask\n        logger.info(f\"Prepared patches for tomogram {tomo_id}\")\n        return tomo_id, patches, mask_patches\n    except Exception as e:\n        logger.error(f\"Error processing tomogram {tomo_id}: {str(e)}\")\n        raise\n\n# Sample patches\ndef sample_patches(tomo_id, volume, mask, labels_df, patch_size=(128, 128, 128), patches_per_volume=128):  # Increased patches\n    shape = volume.shape\n    patches = []\n    mask_patches = []\n    tomo_labels = labels_df[labels_df[\"tomo_id\"] == tomo_id]\n    motor_coords = []\n    for _, row in tomo_labels.iterrows():\n        if row[\"Number of motors\"] > 0 and row[\"Motor axis 0\"] != -1:\n            z, y, x = int(row[\"Motor axis 0\"]), int(row[\"Motor axis 1\"]), int(row[\"Motor axis 2\"])\n            if 0 <= z < shape[0] and 0 <= y < shape[1] and 0 <= x < shape[2]:\n                motor_coords.append((z, y, x))\n    \n    for _ in range(patches_per_volume // 2):\n        if motor_coords:\n            zc, yc, xc = motor_coords[np.random.randint(len(motor_coords))]\n            z = np.clip(zc - patch_size[0]//2 + np.random.randint(-32, 32), 0, shape[0] - patch_size[0])\n            y = np.clip(yc - patch_size[1]//2 + np.random.randint(-32, 32), 0, shape[1] - patch_size[1])\n            x = np.clip(xc - patch_size[2]//2 + np.random.randint(-32, 32), 0, shape[2] - patch_size[2])\n        else:\n            z = np.random.randint(0, max(1, shape[0] - patch_size[0]))\n            y = np.random.randint(0, max(1, shape[1] - patch_size[1]))\n            x = np.random.randint(0, max(1, shape[2] - patch_size[2]))\n        patch = volume[z:z+patch_size[0], y:y+patch_size[1], x:x+patch_size[2]][np.newaxis, ...]\n        mask_patch = mask[z:z+patch_size[0], y:y+patch_size[1], x:x+patch_size[2]][np.newaxis, ...]\n        patches.append(patch)\n        mask_patches.append(mask_patch)\n    \n    for _ in range(patches_per_volume // 2):\n        z = np.random.randint(0, max(1, shape[0] - patch_size[0]))\n        y = np.random.randint(0, max(1, shape[1] - patch_size[1]))\n        x = np.random.randint(0, max(1, shape[2] - patch_size[2]))\n        patch = volume[z:z+patch_size[0], y:y+patch_size[1], x:x+patch_size[2]][np.newaxis, ...]\n        mask_patch = mask[z:z+patch_size[0], y:y+patch_size[1], x:x+patch_size[2]][np.newaxis, ...]\n        patches.append(patch)\n        mask_patches.append(mask_patch)\n    \n    logger.info(f\"Sampled {len(patches)} patches for tomogram {tomo_id}\")\n    return np.array(patches), np.array(mask_patches)\n\n# Patch dataset\nclass PatchDataset(Dataset):\n    def __init__(self, patches, mask_patches):\n        self.patches = patches\n        self.mask_patches = mask_patches\n    \n    def __len__(self):\n        return len(self.patches)\n    \n    def __getitem__(self, idx):\n        patch = self.patches[idx]\n        mask_patch = self.mask_patches[idx]\n        return torch.tensor(patch, dtype=torch.float32), torch.tensor(mask_patch, dtype=torch.float32)\n\n# Tomogram dataset\nclass TomogramDataset(Dataset):\n    def __init__(self, tomo_id, gcs_preprocessed_path, local_dir, mode=\"test\"):\n        self.tomo_id = tomo_id\n        self.gcs_preprocessed_path = gcs_preprocessed_path\n        self.gcs_precomputed_path = gcs_preprocessed_path.replace(\"preprocessed\", \"precomputedmasks\")\n        self.local_dir = local_dir\n        self.mode = mode\n        self.volume = None\n    \n    def load(self):\n        tomo_path, mask_path = download_npy_and_mask(\n            self.tomo_id, \n            self.gcs_preprocessed_path, \n            self.gcs_precomputed_path, \n            \"train\" if self.mode == \"val\" else self.mode, \n            self.local_dir\n        )\n        self.volume = np.load(tomo_path)\n        logger.info(f\"Loaded tomogram {self.tomo_id} (shape: {self.volume.shape})\")\n        if os.path.exists(mask_path):\n            self.mask = np.load(mask_path)\n            logger.info(f\"Loaded mask {self.tomo_id} (min/max: {self.mask.min()}/{self.mask.max()})\")\n        else: \n            self.mask = np.zeros_like(self.volume)\n            logger.info(f\"No mask found for {self.tomo_id}, using zeros (shape: {self.mask.shape})\")\n    \n    def clear(self):\n        tomo_path = os.path.join(self.local_dir, f\"{self.tomo_id}.npy\")\n        mask_path = os.path.join(self.local_dir, f\"{self.tomo_id}_mask.npy\")\n        for path in [tomo_path, mask_path]:\n            if os.path.exists(path):\n                os.remove(path)\n                logger.info(f\"Deleted {path}\")\n        if self.volume is not None:\n            del self.volume\n            self.volume = None\n        if self.mask is not None:\n            del self.mask\n            self.mask = None\n        gc.collect()\n    \n    def __len__(self):\n        return 1\n    \n    def __getitem__(self, idx):\n        if self.volume is None:\n            self.load()\n        volume = torch.from_numpy(self.volume).float()\n        if len(volume.shape) == 3:  # (z, y, x) -> (1, 1, z, y, x)\n            volume = volume.unsqueeze(0).unsqueeze(0)\n        elif len(volume.shape) != 5:\n            raise ValueError(f\"Unexpected volume shape: {volume.shape}\")\n        mask = torch.from_numpy(self.mask).float() if self.mask is not None else torch.zeros_like(volume)\n        if len(mask.shape) == 3:\n            mask = mask.unsqueeze(0).unsqueeze(0)\n        elif len(mask.shape) != 5 and self.mask is not None:\n            raise ValueError(f\"Unexpected mask shape: {mask.shape}\")\n        return volume, mask","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Initialize model\nmodel = UNet(\n    spatial_dims=3,\n    in_channels=1,\n    out_channels=1,\n    channels=(16, 32, 64, 128, 256),\n    strides=(2, 2, 2, 2),\n    num_res_units=2,\n)\nif torch.cuda.device_count() > 1:\n    logger.info(f\"Using {torch.cuda.device_count()} GPUs with DataParallel\")\n    model = torch.nn.DataParallel(model)\nmodel = model.to(device)\nlogger.info(f\"Model device: {next(model.parameters()).device}\")\nlogger.info(f\"GPU count: {torch.cuda.device_count()}, Device: {device}\")\n\n# Loss and optimizer\ncriterion = DiceLoss(sigmoid=True)\noptimizer = optim.Adam(model.parameters(), lr=1e-3)\nscheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode=\"min\", factor=0.1, patience=5)\n\n# Load checkpoint if available\nstart_epoch = 0\nbest_val_loss = float(\"inf\")\ntrain_losses = []\nval_losses = []\nif os.path.exists(\"checkpoint.pth\"):\n    checkpoint = torch.load(\"checkpoint.pth\", map_location=device)\n    model.load_state_dict(checkpoint['model_state_dict'])\n    optimizer.load_state_dict(checkpoint['optimizer_state_dict'])\n    scheduler.load_state_dict(checkpoint['scheduler_state_dict'])\n    best_val_loss = checkpoint['best_val_loss']\n    start_epoch = checkpoint['epoch'] + 1\n    logger.info(f\"Resumed from epoch {start_epoch}, best validation loss: {best_val_loss}\")\nelse:\n    logger.info(\"No checkpoint found, starting from scratch.\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# TRAINING FUNCTION LOOP \n","metadata":{}},{"cell_type":"code","source":"# Clean memory after processing tomograms\ndef clean_memory(tomo_ids, local_dir):\n    \"\"\"Delete tomogram and mask files in parallel and clear GPU memory.\"\"\"\n    def delete_file(path):\n        \"\"\"Helper function to delete a single file.\"\"\"\n        try:\n            if os.path.exists(path):\n                os.remove(path)\n                logger.info(f\"Deleted {path}\")\n        except Exception as e:\n            logger.error(f\"Error deleting {path}: {str(e)}\")\n\n    file_paths = []\n    for tomo_id in tomo_ids:\n        tomo_path = os.path.join(local_dir, f\"{tomo_id}.npy\")\n        mask_path = os.path.join(local_dir, f\"{tomo_id}_mask.npy\")\n        file_paths.extend([tomo_path, mask_path])\n\n    with ThreadPoolExecutor(max_workers=4) as executor:\n        executor.map(delete_file, file_paths)\n\n    gc.collect()\n    torch.cuda.empty_cache()\n    logger.info(\"Cleared memory and GPU cache\")\n\n# Training function\ndef train_epoch(model, loader, criterion, optimizer, epoch, start_epoch, tomo_id):\n    model.train()\n    epoch_loss = 0.0\n    for i, (inputs, targets) in enumerate(tqdm(loader, desc=f\"Training tomo {tomo_id}\", file=sys.stdout, disable=True)):\n        inputs = inputs.to(device)\n        targets = targets.to(device)\n        optimizer.zero_grad()\n        outputs = model(inputs)\n        loss = criterion(outputs, targets)\n        loss.backward()\n        optimizer.step()\n        epoch_loss += loss.item()\n    avg_loss = epoch_loss / len(loader)\n    logger.info(f\"Epoch {epoch+1}, Tomo {tomo_id} completed, Average Loss: {avg_loss:.4f}\")\n    return avg_loss\n\n# Cache validation tomograms\ndef cache_validation_tomograms(val_ids, gcs_preprocessed_path, gcs_precomputed_path, local_dir):\n    for tomo_id in val_ids:  # Use all val_ids\n        tomo_path, mask_path = download_npy_and_mask(tomo_id, gcs_preprocessed_path, gcs_precomputed_path, \"train\", local_dir)\n        logger.info(f\"Cached {tomo_id} at {tomo_path}, {mask_path}\")\n    return val_ids\n\n# Parallel data loading for validation\ndef load_tomogram_for_validation(tomo_id, gcs_preprocessed_path, gcs_precomputed_path, local_dir, labels_df, patches_per_volume):\n    try:\n        tomo_path = os.path.join(local_dir, f\"{tomo_id}.npy\")\n        mask_path = os.path.join(local_dir, f\"{tomo_id}_mask.npy\")\n        if not (os.path.exists(tomo_path) and os.path.exists(mask_path)):\n            tomo_path, mask_path = download_npy_and_mask(tomo_id, gcs_preprocessed_path, gcs_precomputed_path, \"train\", local_dir)\n        volume = np.load(tomo_path)\n        mask = np.load(mask_path)\n        patches, mask_patches = sample_patches(tomo_id, volume, mask, labels_df, patches_per_volume=patches_per_volume)\n        logger.info(f\"Loaded and sampled patches for validation tomogram {tomo_id}\")\n        return tomo_id, patches, mask_patches, [tomo_path, mask_path]\n    except Exception as e:\n        logger.error(f\"Error loading tomogram {tomo_id}: {str(e)}\")\n        raise\n\n# Optimized validate function\ndef validate(model, val_ids, gcs_preprocessed_path, gcs_precomputed_path, local_dir, labels_df, criterion, patch_size):\n    model.eval()\n    epoch_loss = 0.0\n    patches_per_volume = 128  # Increased to match sample_patches\n    selected_val_ids = cache_validation_tomograms(val_ids, gcs_preprocessed_path, gcs_precomputed_path, local_dir)\n    \n    max_workers = min(os.cpu_count(), 8)\n    logger.info(f\"Using {max_workers} workers for validation data loading\")\n    \n    # Preload all tomograms in parallel\n    with ThreadPoolExecutor(max_workers=max_workers) as executor:\n        future_to_tomo = {\n            executor.submit(\n                load_tomogram_for_validation,\n                tomo_id,\n                gcs_preprocessed_path,\n                gcs_precomputed_path,\n                local_dir,\n                labels_df,\n                patches_per_volume\n            ): tomo_id for tomo_id in selected_val_ids\n        }\n        \n        results = []\n        for future in tqdm(as_completed(future_to_tomo), total=len(selected_val_ids), desc=\"Loading validation tomograms\"):\n            tomo_id = future_to_tomo[future]\n            try:\n                results.append(future.result())\n            except Exception as e:\n                logger.error(f\"Error loading tomogram {tomo_id}: {str(e)}\")\n                raise\n\n    # Process loaded tomograms\n    tomo_ids_processed = []\n    for tomo_id, patches, mask_patches, paths in tqdm(results, total=len(selected_val_ids), desc=\"Validating tomograms\"):\n        logger.info(f\"Validating tomogram {tomo_id}\")\n        try:\n            tomo_labels = labels_df[labels_df[\"tomo_id\"] == tomo_id].iloc[0]\n            num_motors = tomo_labels[\"Number of motors\"]\n            logger.info(f\"Tomogram {tomo_id} has {num_motors} motors\")\n            \n            dataset = PatchDataset(patches, mask_patches)\n            loader = DataLoader(dataset, batch_size=8, shuffle=False, num_workers=max_workers, pin_memory=True)\n            \n            with torch.no_grad():\n                for i, (inputs, targets) in enumerate(loader):\n                    inputs = inputs.to(device)\n                    targets = targets.to(device)\n                    outputs = model(inputs)\n                    loss = criterion(outputs, targets)\n                    epoch_loss += loss.item()\n                    logger.info(f\"Validation tomo {tomo_id}, Batch {i+1}, Loss: {loss.item():.4f}\")\n            \n            del dataset, loader, patches, mask_patches\n            tomo_ids_processed.append(tomo_id)\n            clean_memory([tomo_id], local_dir)\n        except Exception as e:\n            logger.error(f\"Error validating tomogram {tomo_id}: {str(e)}\")\n            raise\n    \n    # # Clean up all tomograms at once\n    # if tomo_ids_processed:\n    #     logger.info(f\"Cleaning memory for tomograms: {tomo_ids_processed}\")\n    #     clean_memory(tomo_ids_processed, local_dir)\n    \n    return epoch_loss / (len(selected_val_ids) * patches_per_volume)\n\n\n# Prefetch thread for producer-consumer queue\ndata_queue = queue.Queue(maxsize=2)\ndef prefetch_batches(train_ids, batch_size, gcs_preprocessed_path, gcs_precomputed_path, local_dir, labels_df, patches_per_volume):\n    for batch_start in range(0, len(train_ids), batch_size):\n        batch_tomo_ids = train_ids[batch_start:batch_start + batch_size]\n        logger.info(f\"Prefetching batch of tomograms: {batch_tomo_ids}\")\n        batch_data = []\n        download_start_time = time.time()\n        try:\n            with ThreadPoolExecutor(max_workers=batch_size) as executor:\n                future_to_tomo = {\n                    executor.submit(\n                        download_and_preprocess,\n                        tomo_id,\n                        gcs_preprocessed_path,\n                        gcs_precomputed_path,\n                        \"train\",\n                        local_dir,\n                        labels_df,\n                        patches_per_volume\n                    ): tomo_id for tomo_id in batch_tomo_ids\n                }\n                for future in as_completed(future_to_tomo):\n                    batch_data.append(future.result())\n            logger.info(f\"Batch prefetch time for {batch_tomo_ids}: {time.time() - download_start_time:.2f} seconds\")\n            data_queue.put((batch_tomo_ids, batch_data))\n        except Exception as e:\n            logger.error(f\"Error prefetching batch {batch_tomo_ids}: {str(e)}\")\n            raise\n\n\n# Training loop with batch processing\nnum_epochs = 10  # Increased for convergence\npatience = 10\ntrigger_times = 0\npatches_per_volume_train = 128  # Increased to match sample_patches\nbatch_size = 3  # Kept as is\nmetrics_log = []\nbest_val_loss = float(\"inf\")\ntrain_losses = []\nval_losses = []\n\n# Analyze tomograms (no plotting)\ntotal_tomograms, motor_tomograms, non_motor_tomograms = analyze_tomograms(labels_df)\n\nfor epoch in range(start_epoch, num_epochs):\n    logger.info(f\"Starting training - Epoch {epoch+1}/{num_epochs}\")\n    epoch_train_loss = 0.0\n    processed_tomograms = 0\n    total_train_tomograms = len(train_ids)  # Use all train_ids\n\n    # Start prefetch thread\n    prefetch_thread = threading.Thread(\n        target=prefetch_batches,\n        args=(train_ids, batch_size, gcs_preprocessed_path, gcs_precomputed_path, local_dir, labels_df, patches_per_volume_train)\n    )\n    prefetch_thread.start()\n\n    # Process tomograms in batches\n    for batch_start in range(0, len(train_ids), batch_size):\n        batch_tomo_ids, batch_data = data_queue.get()\n        logger.info(f\"Processing batch of tomograms: {batch_tomo_ids}\")\n\n        # Train on batch\n        for tomo_id, patches, mask_patches in batch_data:\n            try:\n                dataset = PatchDataset(patches, mask_patches)\n                loader = DataLoader(dataset, batch_size=4, shuffle=True, num_workers=2, pin_memory=True)\n                logger.info(f\"Training on tomogram {tomo_id}\")\n                batch_loss = train_epoch(model, loader, criterion, optimizer, epoch, start_epoch, tomo_id)\n                epoch_train_loss += batch_loss\n                processed_tomograms += 1\n                logger.info(f\"Completed training on tomogram {tomo_id}, Loss: {batch_loss:.4f}\")\n                logger.info(f\"Processed {processed_tomograms}/{total_train_tomograms} tomograms\")\n                del dataset, loader, patches, mask_patches\n                clean_memory([tomo_id], local_dir)  # Clean after each tomogram\n            except Exception as e:\n                logger.error(f\"Error training tomogram {tomo_id}: {str(e)}\")\n                raise\n\n    prefetch_thread.join()\n\n    train_loss = epoch_train_loss / max(1, processed_tomograms)\n    train_losses.append(train_loss)\n\n    # Validation\n    logger.info(\"Starting validation\")\n    patch_size = (128, 128, 128)\n    val_loss = validate(model, val_ids, gcs_preprocessed_path, gcs_precomputed_path, local_dir, labels_df, criterion, patch_size)\n    val_losses.append(val_loss)\n    \n    # Log metrics\n    logger.info(f\"Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}\")\n    metrics_log.append({\n        \"epoch\": epoch + 1,\n        \"train_loss\": train_loss,\n        \"val_loss\": val_loss\n    })\n    \n    # Save metrics to file\n    pd.DataFrame(metrics_log).to_csv(\"/kaggle/working/training_metrics.csv\", index=False)\n    logger.info(\"Saved training metrics to training_metrics.csv\")\n\n    scheduler.step(val_loss)\n\n    # Checkpointing and Early stopping\n    if val_loss < best_val_loss:\n        best_val_loss = val_loss\n        trigger_times = 0\n        torch.save({\n            'epoch': epoch,\n            'model_state_dict': model.state_dict(),\n            'optimizer_state_dict': optimizer.state_dict(),\n            'scheduler_state_dict': scheduler.state_dict(),\n            'best_val_loss': best_val_loss\n        }, \"checkpoint.pth\")\n        torch.save(model.state_dict(), \"/kaggle/working/best_model.pth\")\n        logger.info(\"Saved checkpoint and best model\")\n    else:\n        trigger_times += 1\n        if trigger_times >= patience:\n            logger.info(\"Early stopping triggered!\")\n            break","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 1 Hyperparameter Tuning (Peak Detection Threshold)\n# 2 Predict on the validation set,\n# 3 Extract motor locations using peak detection\n# 4 Tune the threshold to maximize the Fβ-score (β=2).\n# 5 Using competition's metric\n\n# Metric implementation\ndef distance_metric(solution, submission, thresh_ratio, min_radius):\n    coordinate_cols = ['Motor axis 0', 'Motor axis 1', 'Motor axis 2']\n    label_tensor = solution[coordinate_cols].values.reshape(len(solution), -1, len(coordinate_cols))\n    predicted_tensor = submission[coordinate_cols].values.reshape(len(submission), -1, len(coordinate_cols))\n    solution['distance'] = np.linalg.norm(label_tensor - predicted_tensor, axis=2).min(axis=1)\n    solution['thresholds'] = solution['Voxel spacing'].apply(lambda x: (min_radius * thresh_ratio) / x)\n    solution['predictions'] = submission['Has motor'].values\n    solution.loc[(solution['distance'] > solution['thresholds']) & (solution['Has motor'] == 1) & (submission['Has motor'] == 1), 'predictions'] = 0\n    return solution['predictions'].values\n\ndef score(solution, submission, min_radius, beta):\n    solution = solution.sort_values('tomo_id').reset_index(drop=True)\n    submission = submission.sort_values('tomo_id').reset_index(drop=True)\n    if not solution['tomo_id'].eq(submission['tomo_id']).all():\n        raise ValueError('Submitted tomo_id values do not match')\n    submission['Has motor'] = 1\n    select = (submission[['Motor axis 0', 'Motor axis 1', 'Motor axis 2']] == -1).any(axis='columns')\n    submission.loc[select, 'Has motor'] = 0\n    predictions = distance_metric(solution, submission, thresh_ratio=1.0, min_radius=min_radius)\n    return sklearn.metrics.fbeta_score(solution['Has motor'].values, predictions, beta=beta)\n\n# Predict full volume\ndef predict_full_volume(model, volume, patch_size=(128, 128, 128), stride=32, batch_size=8):  # Reduced stride\n    model.eval()\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    model.to(device)\n    \n    # Convert volume to tensor if needed and keep on CPU\n    if isinstance(volume, np.ndarray):\n        volume = torch.from_numpy(volume).float()\n    if len(volume.shape) == 3:  # (z, y, x) -> (1, 1, z, y, x)\n        volume = volume.unsqueeze(0).unsqueeze(0)\n    elif len(volume.shape) == 4:  # (channels, z, y, x) -> (1, channels, z, y, x)\n        volume = volume.unsqueeze(0)\n    elif len(volume.shape) != 5:\n        raise ValueError(f\"Expected volume with 3-5 dimensions, got shape {volume.shape}\")\n    \n    batch, channels, z_size, y_size, x_size = volume.shape\n    pz, py, px = patch_size\n    output = torch.zeros_like(volume, device='cpu')\n    counts = torch.zeros_like(volume, device='cpu')\n    \n    # Collect patch coordinates\n    patches = []\n    coords = []\n    for z in range(0, z_size, stride):\n        for y in range(0, y_size, stride):\n            for x in range(0, x_size, stride):\n                z_end, y_end, x_end = z + pz, y + py, x + px\n                coords.append((z, y, x, min(z_end, z_size), min(y_end, y_size), min(x_end, x_size)))\n                patch = volume[:, :, z:min(z_end, z_size), y:min(y_end, y_size), x:min(x_end, x_size)]\n                pad_z, pad_y, pad_x = max(0, z_end - z_size), max(0, y_end - y_size), max(0, x_end - x_size)\n                if pad_z > 0 or pad_y > 0 or pad_x > 0:\n                    patch = F.pad(patch, (0, pad_x, 0, pad_y, 0, pad_z))\n                patches.append(patch)\n    \n    logger.info(f\"Processing {len(patches)} patches for volume {volume.shape}\")\n    \n    # Process patches in batches\n    with torch.no_grad():\n        for i in tqdm(range(0, len(patches), batch_size), desc=\"Processing patches\"):\n            batch_patches = torch.cat(patches[i:i+batch_size], dim=0).to(device, non_blocking=True)\n            batch_out = torch.sigmoid(model(batch_patches))\n            batch_out = batch_out.cpu()\n            for j, (z, y, x, z_end, y_end, x_end) in enumerate(coords[i:i+batch_size]):\n                output[:, :, z:z_end, y:y_end, x:x_end] += batch_out[j, :, :z_end-z, :y_end-y, :x_end-x]\n                counts[:, :, z:z_end, y:y_end, x:x_end] += 1\n            del batch_patches, batch_out\n            torch.cuda.empty_cache()\n    \n    output = output / (counts + 1e-8)\n    logger.info(f\"Predicted volume shape: {output.shape}, patches processed: {len(patches)}\")\n    return output.numpy()\n\n\n# Extract motor location\ndef extract_motor_location(mask, threshold):\n    mask = mask.squeeze()\n    max_val = mask.max()\n    logger.info(f\"Extracting motor location, mask max: {max_val:.4f}, threshold: {threshold:.2f}\")\n    if max_val < threshold:\n        return -1, -1, -1, 0\n    z, y, x = np.unravel_index(np.argmax(mask), mask.shape)\n    region = mask[max(0, z-5):z+6, max(0, y-5):y+6, max(0, x-5):x+6]\n    if region.size == 0:\n        logger.info(f\"Motor at ({z}, {y}, {x}), region empty\")\n        return z, y, x, 1\n    z_offset, y_offset, x_offset = center_of_mass(region)\n    z, y, x = z + z_offset - 5, y + y_offset - 5, x + x_offset - 5\n    logger.info(f\"Motor at ({z:.2f}, {y:.2f}, {x:.2f})\")\n    return z, y, x, 1\n\n\n# Tune peak detection threshold\ndef tune_threshold(model, val_ids, gcs_preprocessed_path, gcs_precomputed_path, local_dir, labels_df, thresholds=np.linspace(0.3, 0.7, 3)):  # Increased thresholds\n    best_model_path = \"/kaggle/working/best_model.pth\"\n    device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n    if os.path.exists(best_model_path):\n        model.load_state_dict(torch.load(best_model_path, weights_only=True))\n        logger.info(\"Loaded best_model.pth\")\n    else:\n        logger.warning(\"best_model.pth not found, using current model state\")\n    model.to(device).eval()\n    best_threshold = 0.5\n    best_fbeta = 0.0\n    thresholds_list = []\n    fbeta_scores = [] \n\n    # Check tomo_ab804d motor count\n    tomo_labels = labels_df[labels_df[\"tomo_id\"] == \"tomo_ab804d\"]\n    if not tomo_labels.empty:\n        num_motors = tomo_labels.iloc[0][\"Number of motors\"]\n        logger.info(f\"tomo_ab804d has {num_motors} motors\")\n    \n    for threshold in thresholds:\n        predictions = []\n        for tomo_id in tqdm(val_ids, desc=f\"Tuning threshold {threshold:.2f}\"):  # Use all val_ids\n            dataset = TomogramDataset(tomo_id, gcs_preprocessed_path, local_dir, mode=\"val\")\n            try:\n                dataset.load()\n                volume, _ = dataset[0]\n                volume = volume.to(device, non_blocking=True)\n                logger.info(f\"Volume shape for {tomo_id}: {volume.shape}\")\n                pred_mask = predict_full_volume(model, volume, patch_size=(128, 128, 128), stride=32, batch_size=8)\n                z, y, x, has_motor = extract_motor_location(pred_mask, threshold)\n                predictions.append({\"tomo_id\": tomo_id, \"Motor axis 0\": z, \"Motor axis 1\": y, \"Motor axis 2\": x, \"Has motor\": has_motor})\n            except Exception as e:\n                logger.error(f\"Error processing tomogram {tomo_id}: {str(e)}\")\n                predictions.append({\"tomo_id\": tomo_id, \"Motor axis 0\": -1, \"Motor axis 1\": -1, \"Motor axis 2\": -1, \"Has motor\": 0})\n            finally:\n                dataset.clear()\n                clean_memory([tomo_id], local_dir)\n                torch.cuda.empty_cache()\n                logger.info(f\"Cleared GPU memory after {tomo_id}, usage: {torch.cuda.memory_allocated() / 1024**3:.2f} GB\")\n        \n        # Create submission and solution DataFrames\n        submission_df = pd.DataFrame(predictions)\n        solution_data = []\n        for tomo_id in val_ids:\n            tomo_labels = labels_df[labels_df[\"tomo_id\"] == tomo_id].iloc[0]\n            solution_data.append({\n                \"tomo_id\": tomo_id,\n                \"Motor axis 0\": tomo_labels[\"Motor axis 0\"],\n                \"Motor axis 1\": tomo_labels[\"Motor axis 1\"],\n                \"Motor axis 2\": tomo_labels[\"Motor axis 2\"],\n                \"Voxel spacing\": tomo_labels[\"Voxel spacing\"],\n                \"Has motor\": 1 if tomo_labels[\"Number of motors\"] > 0 else 0\n            })\n        solution_df = pd.DataFrame(solution_data)\n        fbeta = score(solution_df, submission_df, min_radius=1000, beta=2)\n        logger.info(f\"Threshold {threshold:.2f}, Fβ-score: {fbeta:.4f}, TP: {((solution_df['Has motor'] == 1) & (submission_df['Has motor'] == 1)).sum()}, \"\n                    f\"FP: {((solution_df['Has motor'] == 0) & (submission_df['Has motor'] == 1)).sum()}, \"\n                    f\"FN: {((solution_df['Has motor'] == 1) & (submission_df['Has motor'] == 0)).sum()}\")\n        thresholds_list.append(threshold)\n        fbeta_scores.append(fbeta)\n        \n        if fbeta > best_fbeta:\n            best_fbeta = fbeta\n            best_threshold = threshold\n    \n    logger.info(f\"Best threshold: {best_threshold:.2f}, Best Fβ-score: {best_fbeta:.4f}\")\n    return best_threshold\n    \n\n# Predict test\ndef predict_test(model, test_ids, gcs_preprocessed_path, local_dir, threshold):\n    if os.path.exists(\"best_model.pth\"):\n        model.load_state_dict(torch.load(\"best_model.pth\", map_location=device))\n        logger.info(\"Loaded best_model.pth\")\n    else:\n        logger.warning(\"best_model.pth not found, using current model state\")\n    model.eval()\n    predictions = []\n    \n    for tomo_id in tqdm(test_ids, desc=\"Predicting test set\"):\n        dataset = TomogramDataset(tomo_id, gcs_preprocessed_path, local_dir, mode=\"test\")\n        try:\n            dataset.load()\n            volume, _ = dataset[0]\n            pred_mask = predict_full_volume(model, volume, patch_size=(128, 128, 128), stride=32, batch_size=8)\n            z, y, x, has_motor = extract_motor_location(pred_mask, threshold)\n            predictions.append({\"tomo_id\": tomo_id, \"Motor axis 0\": z, \"Motor axis 1\": y, \"Motor axis 2\": x})\n        except Exception as e:\n            logger.error(f\"Error predicting tomogram {tomo_id}: {str(e)}\")\n            predictions.append({\"tomo_id\": tomo_id, \"Motor axis 0\": -1, \"Motor axis 1\": -1, \"Motor axis 2\": -1})\n        finally:\n            dataset.clear()\n            clean_memory([tomo_id], local_dir)\n            torch.cuda.empty_cache()\n            logger.info(f\"Cleared GPU memory after {tomo_id}, usage: {torch.cuda.memory_allocated() / 1024**3:.2f} GB\")\n    \n    return pd.DataFrame(predictions)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Run pipeline\nlogger.info(\"Starting hyperparameter tuning...\")\nbest_threshold = tune_threshold(model, val_ids, gcs_preprocessed_path, gcs_precomputed_path, local_dir, labels_df)  # Use all val_ids\nlogger.info(\"Generating test predictions...\")\nsubmission_df = predict_test(model, test_ids, gcs_preprocessed_path, local_dir, best_threshold)\nsubmission_df.to_csv(\"submission.csv\", index=False)\nlogger.info(\"Submission file created: submission.csv\")\n# Finish the wandb run\n#wandb.finish()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}