{"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":"none","dataSources":[{"sourceId":91249,"databundleVersionId":11294684,"sourceType":"competition"},{"sourceId":11495557,"sourceType":"datasetVersion","datasetId":7206260},{"sourceId":11768112,"sourceType":"datasetVersion","datasetId":7387959},{"sourceId":11852890,"sourceType":"datasetVersion","datasetId":7447868},{"sourceId":11905081,"sourceType":"datasetVersion","datasetId":7206328},{"sourceId":327336,"sourceType":"modelInstanceVersion","modelInstanceId":274744,"modelId":295634},{"sourceId":386302,"sourceType":"modelInstanceVersion","modelInstanceId":318564,"modelId":339142},{"sourceId":399936,"sourceType":"modelInstanceVersion","modelInstanceId":327282,"modelId":348163},{"sourceId":400014,"sourceType":"modelInstanceVersion","modelInstanceId":327345,"modelId":348227},{"sourceId":403418,"sourceType":"modelInstanceVersion","modelInstanceId":329839,"modelId":350688},{"sourceId":405315,"sourceType":"modelInstanceVersion","modelInstanceId":331197,"modelId":352093}],"dockerImageVersionId":31012,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# !pip download -d ./packages ultralytics\n# !tar cfvz archive.tar.gz ./packages\nimport kagglehub, os\nkagglehub.dataset_download('rachiteagles/yolo-pkg')\nkagglehub.dataset_download('rachiteagles/yolo-model')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-22T07:11:38.642396Z","iopub.execute_input":"2025-05-22T07:11:38.642792Z","iopub.status.idle":"2025-05-22T07:11:39.198813Z","shell.execute_reply.started":"2025-05-22T07:11:38.642761Z","shell.execute_reply":"2025-05-22T07:11:39.197792Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"os.environ['WANDB_MODE'] = 'offline'","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-22T07:11:41.494602Z","iopub.execute_input":"2025-05-22T07:11:41.494930Z","iopub.status.idle":"2025-05-22T07:11:41.499664Z","shell.execute_reply.started":"2025-05-22T07:11:41.494907Z","shell.execute_reply":"2025-05-22T07:11:41.498546Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# !tar xfvz archive.tar.gz\n# !pip install --no-index --no-deps /kaggle/input/yolo-pkg/yolo/ultralytics-8.3.112-py3-none-any.whl","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-22T07:11:42.174184Z","iopub.execute_input":"2025-05-22T07:11:42.174603Z","iopub.status.idle":"2025-05-22T07:11:42.180015Z","shell.execute_reply.started":"2025-05-22T07:11:42.174577Z","shell.execute_reply":"2025-05-22T07:11:42.178805Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!cp -r /kaggle/input/mhaf_yolo/pytorch/default/2/MHAF-YOLO-main/ultralytics /kaggle/working/","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-22T07:11:42.746130Z","iopub.execute_input":"2025-05-22T07:11:42.746527Z","iopub.status.idle":"2025-05-22T07:11:44.431547Z","shell.execute_reply.started":"2025-05-22T07:11:42.746492Z","shell.execute_reply":"2025-05-22T07:11:44.430010Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class TrainingConfig:\n def __init__(self):\n self.batch_size = 8 # Reduced from original 16 if memory issues\n self.imgsz = 640 # Image size\n self.epochs = 100\n self.patience = 20 # Early stopping patience\n self.device = [0,1] if torch.cuda.device_count() > 1 else 0 # Use all available GPUs\n self.optimizer = 'AdamW'\n self.lr0 = 1e-4 # Initial learning rate\n self.lrf = 0.1 # Final learning rate (lr0 * lrf)\n self.momentum = 0.9\n self.weight_decay = 0.0005\n self.warmup_epochs = 3.0\n self.warmup_momentum = 0.8\n self.warmup_bias_lr = 0.1\n self.box = 7.5 # box loss gain\n self.cls = 0.5 # cls loss gain\n self.dfl = 1.5 # dfl loss gain\n self.hsv_h = 0.015 # image HSV-Hue augmentation (fraction)\n self.hsv_s = 0.7 # image HSV-Saturation augmentation (fraction)\n self.hsv_v = 0.4 # image HSV-Value augmentation (fraction)\n self.degrees = 45.0 # image rotation (+/- deg)\n self.translate = 0.1 # image translation (+/- fraction)\n self.scale = 0.5 # image scale (+/- gain)\n self.shear = 0.0 # image shear (+/- deg)\n self.perspective = 0.0001 # image perspective (+/- fraction), range 0-0.001\n self.flipud = 0.5 # image flip up-down (probability)\n self.fliplr = 0.5 # image flip left-right (probability)\n self.mosaic = 1.0 # image mosaic (probability)\n self.mixup = 0.2 # image mixup (probability)\n self.copy_paste = 0.2 # segment copy-paste (probability)\n self.dropout = 0.1 # use dropout regularization","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-22T07:11:44.693857Z","iopub.execute_input":"2025-05-22T07:11:44.694200Z","iopub.status.idle":"2025-05-22T07:11:44.703068Z","shell.execute_reply.started":"2025-05-22T07:11:44.694173Z","shell.execute_reply":"2025-05-22T07:11:44.701963Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os, cv2, torch, yaml, threading, time\nimport pandas as pd\nfrom tqdm import tqdm\nimport numpy as np\nfrom PIL import Image\nimport matplotlib.pyplot as plt\nfrom torch.utils.data import Dataset, DataLoader\nimport torchvision.transforms as transforms\nfrom torchvision.transforms import Compose, Resize, ToTensor\nfrom torchvision.utils import make_grid\n# from kaggle.input.mhafyolo.pytorch.default.1.MHAF-YOLO-main.ultralytics.models import YOLOv10\nfrom ultralytics.models import YOLOv10,YOLO\nfrom concurrent.futures import ThreadPoolExecutor\n\nnp.random.seed(42)\ntorch.manual_seed(42)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-05-22T07:11:45.085213Z","iopub.execute_input":"2025-05-22T07:11:45.085659Z","iopub.status.idle":"2025-05-22T07:11:56.912020Z","shell.execute_reply.started":"2025-05-22T07:11:45.085623Z","shell.execute_reply":"2025-05-22T07:11:56.910819Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Define YOLO dataset structure\nyolo_dataset_dir = \"/kaggle/working/yolo_dataset\"\nyolo_images_train = os.path.join(yolo_dataset_dir, \"images\", \"train\")\nyolo_images_val = os.path.join(yolo_dataset_dir, \"images\", \"val\")\nyolo_labels_train = os.path.join(yolo_dataset_dir, \"labels\", \"train\")\nyolo_labels_val = os.path.join(yolo_dataset_dir, \"labels\", \"val\")\n\n# Create directories\nfor dir_path in [yolo_images_train, yolo_images_val, yolo_labels_train, yolo_labels_val]:\n os.makedirs(dir_path, exist_ok=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-22T07:11:56.914159Z","iopub.execute_input":"2025-05-22T07:11:56.914760Z","iopub.status.idle":"2025-05-22T07:11:56.923284Z","shell.execute_reply.started":"2025-05-22T07:11:56.914723Z","shell.execute_reply":"2025-05-22T07:11:56.922110Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Define global constants for dataset directories\nDATA_DIR = '/kaggle/input/byu-locating-bacterial-flagellar-motors-2025'\nTRAIN_CSV = os.path.join(DATA_DIR, 'train_labels.csv')\nTRAIN_DIR = os.path.join(DATA_DIR, 'train')\nTEST_DIR = os.path.join(DATA_DIR, 'test')\nMODEL_PATH = \"/kaggle/input/yolo-model/best.pt\"\n# MODEL_PATH = \"/kaggle/working/yolo_weights/motor_detector/weights/best.pt\"\nOUTPUT_DIR = './'\nMODEL_DIR = './models'\nTRUST = 4\nTRAIN_SPLIT = 0.8\nBOX_SIZE = 24\nDEVICE = 'cuda:0' if torch.cuda.is_available() else 'cpu'\nBATCH_SIZE = 8 \nSUBMISSION_PATH = \"/kaggle/working/submission.csv\"\n\n# Detection parameters\nCONFIDENCE_THRESHOLD = 0.45 # Lower threshold to catch more potential motors\nMAX_DETECTIONS_PER_TOMO = 1 # Keep track of top N detections per tomogram\nNMS_IOU_THRESHOLD = 0.2 # Non-maximum suppression threshold for 3D clustering\nCONCENTRATION = 1 # ONLY PROCESS 1/20 slices for fast submission","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-22T07:11:56.924397Z","iopub.execute_input":"2025-05-22T07:11:56.924919Z","iopub.status.idle":"2025-05-22T07:11:56.973136Z","shell.execute_reply.started":"2025-05-22T07:11:56.924877Z","shell.execute_reply":"2025-05-22T07:11:56.971914Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Image processing functions\ndef normalize_slice(slice_data):\n \"\"\"\n Normalize slice data using 2nd and 98th percentiles\n \"\"\"\n # Calculate percentiles\n p2 = np.percentile(slice_data, 2)\n p98 = np.percentile(slice_data, 98)\n \n # Clip the data to the percentile range\n clipped_data = np.clip(slice_data, p2, p98)\n \n # Normalize to [0, 255] range\n normalized = 255 * (clipped_data - p2) / (p98 - p2)\n \n return np.uint8(normalized)\n\n\n\ndef prepare_yolo_dataset(trust=TRUST, train_split=TRAIN_SPLIT):\n \"\"\"\n Extract slices containing motors from tomograms and save to YOLO structure with annotations\n \"\"\"\n # Load the labels CSV\n labels_df = pd.read_csv(TRAIN_CSV)\n \n # Count total number of motors\n total_motors = labels_df['Number of motors'].sum()\n print(f\"Total number of motors in the dataset: {total_motors}\")\n \n # Get unique tomograms that have motors\n tomo_df = labels_df[labels_df['Number of motors'] > 0].copy()\n unique_tomos = tomo_df['tomo_id'].unique()\n \n print(f\"Found {len(unique_tomos)} unique tomograms with motors\")\n \n # Perform the train-val split at the tomogram level (not motor level)\n # This ensures all slices from a single tomogram go to either train or val\n np.random.shuffle(unique_tomos) # Shuffle the tomograms\n split_idx = int(len(unique_tomos) * train_split)\n train_tomos = unique_tomos[:split_idx]\n val_tomos = unique_tomos[split_idx:]\n \n print(f\"Split: {len(train_tomos)} tomograms for training, {len(val_tomos)} tomograms for validation\")\n \n # Function to process a set of tomograms\n def process_tomogram_set(tomogram_ids, images_dir, labels_dir, set_name):\n motor_counts = []\n for tomo_id in tomogram_ids:\n # Get all motors for this tomogram\n tomo_motors = labels_df[labels_df['tomo_id'] == tomo_id]\n for _, motor in tomo_motors.iterrows():\n if pd.isna(motor['Motor axis 0']):\n continue\n motor_counts.append(\n (tomo_id, \n int(motor['Motor axis 0']), \n int(motor['Motor axis 1']), \n int(motor['Motor axis 2']),\n int(motor['Array shape (axis 0)']))\n )\n \n print(f\"Will process approximately {len(motor_counts) * (2 * trust + 1)} slices for {set_name}\")\n \n # Process each motor\n processed_slices = 0\n \n for tomo_id, z_center, y_center, x_center, z_max in tqdm(motor_counts, desc=f\"Processing {set_name} motors\"):\n # Calculate range of slices to include\n z_min = max(0, z_center - trust)\n z_max = min(z_max - 1, z_center + trust)\n \n # Process each slice in the range\n for z in range(z_min, z_max + 1):\n # Create slice filename\n slice_filename = f\"slice_{z:04d}.jpg\"\n \n # Source path for the slice\n src_path = os.path.join(TRAIN_DIR, tomo_id, slice_filename)\n \n if not os.path.exists(src_path):\n print(f\"Warning: {src_path} does not exist, skipping.\")\n continue\n \n # Load and normalize the slice\n img = cv2.imread(src_path)\n\n # Get image dimensions\n img_height,img_width,_ = img.shape\n\n # for angle in [0, 90, 180, 270]:\n # rotated_img, new_x, new_y = rotate_image_and_coords(\n # img, angle, x_center, y_center, img_width, img_height\n # )\n img_array = np.array(img)\n \n # Normalize the image\n normalized_img = normalize_slice(img_array)\n \n # Create destination filename (with unique identifier)\n dest_filename = f\"{tomo_id}_z{z:04d}_y{y_center:04d}_x{x_center:04d}.jpg\"\n dest_path = os.path.join(images_dir, dest_filename)\n \n # Save the normalized image\n Image.fromarray(normalized_img).save(dest_path)\n \n \n \n # Create YOLO format label\n # YOLO format: \n # Values are normalized to [0, 1]\n x_center_norm = x_center / img_width\n y_center_norm = y_center / img_height\n box_width_norm = BOX_SIZE / img_width\n box_height_norm = BOX_SIZE / img_height\n \n # Write label file\n label_path = os.path.join(labels_dir, dest_filename.replace('.jpg', '.txt'))\n with open(label_path, 'w') as f:\n f.write(f\"0 {x_center_norm} {y_center_norm} {box_width_norm} {box_height_norm}\\n\")\n \n processed_slices += 1\n \n return processed_slices, len(motor_counts)\n \n # Process training tomograms\n train_slices, train_motors = process_tomogram_set(train_tomos, yolo_images_train, yolo_labels_train, \"training\")\n \n # Process validation tomograms\n val_slices, val_motors = process_tomogram_set(val_tomos, yolo_images_val, yolo_labels_val, \"validation\")\n \n # Create YAML configuration file for YOLO\n yaml_content = {\n 'path': yolo_dataset_dir,\n 'train': 'images/train',\n 'val': 'images/val',\n 'names': {0: 'motor'}\n }\n \n with open(os.path.join(yolo_dataset_dir, 'dataset.yaml'), 'w') as f:\n yaml.dump(yaml_content, f, default_flow_style=False)\n \n print(f\"\\nProcessing Summary:\")\n print(f\"- Train set: {len(train_tomos)} tomograms, {train_motors} motors, {train_slices} slices\")\n print(f\"- Validation set: {len(val_tomos)} tomograms, {val_motors} motors, {val_slices} slices\")\n print(f\"- Total: {len(train_tomos) + len(val_tomos)} tomograms, {train_motors + val_motors} motors, {train_slices + val_slices} slices\")\n \n # Return summary info\n return {\n \"dataset_dir\": yolo_dataset_dir,\n \"yaml_path\": os.path.join(yolo_dataset_dir, 'dataset.yaml'),\n \"train_tomograms\": len(train_tomos),\n \"val_tomograms\": len(val_tomos),\n \"train_motors\": train_motors,\n \"val_motors\": val_motors,\n \"train_slices\": train_slices,\n \"val_slices\": val_slices\n }\n ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-22T07:11:56.975227Z","iopub.execute_input":"2025-05-22T07:11:56.975615Z","iopub.status.idle":"2025-05-22T07:11:56.999523Z","shell.execute_reply.started":"2025-05-22T07:11:56.975583Z","shell.execute_reply":"2025-05-22T07:11:56.998185Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# summary = prepare_yolo_dataset(TRUST)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-22T02:30:08.904446Z","iopub.execute_input":"2025-05-22T02:30:08.904664Z","iopub.status.idle":"2025-05-22T02:36:00.358297Z","shell.execute_reply.started":"2025-05-22T02:30:08.904639Z","shell.execute_reply":"2025-05-22T02:36:00.357625Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# summary","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-22T02:36:00.359952Z","iopub.execute_input":"2025-05-22T02:36:00.360210Z","iopub.status.idle":"2025-05-22T02:36:00.364823Z","shell.execute_reply.started":"2025-05-22T02:36:00.360193Z","shell.execute_reply":"2025-05-22T02:36:00.364276Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def plot_dfl_loss_curve(run_dir):\n \"\"\"\n Plot the DFL loss curves for train and validation, marking the best model\n \n Args:\n run_dir (str): Directory where the training results are stored\n \"\"\"\n # Path to the results CSV file\n results_csv = os.path.join(run_dir, 'results.csv')\n \n if not os.path.exists(results_csv):\n print(f\"Results file not found at {results_csv}\")\n return\n \n # Read results CSV\n results_df = pd.read_csv(results_csv)\n \n # Check if DFL loss columns exist\n train_dfl_col = [col for col in results_df.columns if 'train/dfl_loss' in col]\n val_dfl_col = [col for col in results_df.columns if 'val/dfl_loss' in col]\n \n if not train_dfl_col or not val_dfl_col:\n print(\"DFL loss columns not found in results CSV\")\n print(f\"Available columns: {results_df.columns.tolist()}\")\n return\n \n train_dfl_col = train_dfl_col[0]\n val_dfl_col = val_dfl_col[0]\n \n # Find the epoch with the best validation loss\n best_epoch = results_df[val_dfl_col].idxmin()\n best_val_loss = results_df.loc[best_epoch, val_dfl_col]\n \n # Create the plot\n plt.figure(figsize=(10, 6))\n \n # Plot training and validation losses\n plt.plot(results_df['epoch'], results_df[train_dfl_col], label='Train DFL Loss')\n plt.plot(results_df['epoch'], results_df[val_dfl_col], label='Validation DFL Loss')\n \n # Mark the best model with a vertical line\n plt.axvline(x=results_df.loc[best_epoch, 'epoch'], color='r', linestyle='--', \n label=f'Best Model (Epoch {int(results_df.loc[best_epoch, \"epoch\"])}, Val Loss: {best_val_loss:.4f})')\n \n # Add labels and legend\n plt.xlabel('Epoch')\n plt.ylabel('DFL Loss')\n plt.title('Training and Validation DFL Loss')\n plt.legend()\n plt.grid(True, linestyle='--', alpha=0.7)\n \n # Save the plot in the same directory as weights\n plot_path = os.path.join(run_dir, 'dfl_loss_curve.png')\n plt.savefig(plot_path)\n \n # Also save it to the working directory for easier access\n plt.savefig(os.path.join('/kaggle/working', 'dfl_loss_curve.png'))\n \n print(f\"Loss curve saved to {plot_path}\")\n plt.close()\n \n # Return the best epoch info\n return best_epoch, best_val_loss","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-22T07:11:57.000710Z","iopub.execute_input":"2025-05-22T07:11:57.001065Z","iopub.status.idle":"2025-05-22T07:11:57.028668Z","shell.execute_reply.started":"2025-05-22T07:11:57.001033Z","shell.execute_reply":"2025-05-22T07:11:57.027583Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def train_yolo_model(yaml_path, pretrained_weights_path, epochs=30, batch_size=16, img_size=640):\n \"\"\"\n Train a YOLO model on the prepared dataset\n \n Args:\n yaml_path (str): Path to the dataset YAML file\n pretrained_weights_path (str): Path to pre-downloaded weights file\n epochs (int): Number of training epochs\n batch_size (int): Batch size for training\n img_size (int): Image size for training\n \"\"\"\n print(f\"Loading pre-trained weights from: {pretrained_weights_path}\")\n \n # Load a pre-trained YOLOv8 model\n model = YOLOv10(pretrained_weights_path)\n \n # Train the model with early stopping\n results = model.train(\n data=yaml_path,\n epochs=cfg.epochs,\n batch=cfg.batch_size,\n imgsz=cfg.imgsz,\n optimizer=cfg.optimizer,\n lr0=cfg.lr0,\n lrf=cfg.lrf,\n momentum=cfg.momentum,\n weight_decay=cfg.weight_decay,\n warmup_epochs=cfg.warmup_epochs,\n warmup_momentum=cfg.warmup_momentum,\n warmup_bias_lr=cfg.warmup_bias_lr,\n box=cfg.box,\n cls=cfg.cls,\n dfl=cfg.dfl,\n hsv_h=cfg.hsv_h,\n hsv_s=cfg.hsv_s,\n hsv_v=cfg.hsv_v,\n degrees=cfg.degrees,\n translate=cfg.translate,\n scale=cfg.scale,\n shear=cfg.shear,\n perspective=cfg.perspective,\n flipud=cfg.flipud,\n fliplr=cfg.fliplr,\n mosaic=cfg.mosaic,\n mixup=cfg.mixup,\n copy_paste=cfg.copy_paste,\n dropout=cfg.dropout,\n device=cfg.device,\n cos_lr= True,\n multi_scale= True,\n amp =True,\n project=yolo_weights_dir,\n save_dir='/kaggle/working',\n name='motor_detector',\n # exist_ok=True,\n patience=cfg.patience,\n save_period=5, # Save checkpoints every 5 epochs\n val=True, # Ensure validation is performed\n verbose=True,\n \n )\n \n # Get the path to the run directory\n run_dir = os.path.join(yolo_weights_dir, 'motor_detector')\n \n # Plot and save the loss curve\n best_epoch_info = plot_dfl_loss_curve(run_dir)\n \n if best_epoch_info:\n best_epoch, best_val_loss = best_epoch_info\n print(f\"\\nBest model found at epoch {best_epoch} with validation DFL loss: {best_val_loss:.4f}\")\n \n return model, results","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-22T07:11:57.030752Z","iopub.execute_input":"2025-05-22T07:11:57.031156Z","iopub.status.idle":"2025-05-22T07:11:57.054276Z","shell.execute_reply.started":"2025-05-22T07:11:57.031117Z","shell.execute_reply":"2025-05-22T07:11:57.053255Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"yaml_path = '/kaggle/working/yolo_dataset/dataset.yaml'\nyolo_pretrained_weights = '/kaggle/input/mhaf_yolon/pytorch/default/1/MAF-YOLOv2-N.pt'\nyolo_weights_dir = \"yolo_weights\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-22T07:12:02.051750Z","iopub.execute_input":"2025-05-22T07:12:02.052734Z","iopub.status.idle":"2025-05-22T07:12:02.057015Z","shell.execute_reply.started":"2025-05-22T07:12:02.052701Z","shell.execute_reply":"2025-05-22T07:12:02.055914Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"cfg = TrainingConfig()\n# model, results = train_yolo_model(\n# yaml_path,\n# pretrained_weights_path=yolo_pretrained_weights,\n# epochs=40 # Using 30 epochs instead of 100 for faster training\n# )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-22T02:38:18.569092Z","iopub.execute_input":"2025-05-22T02:38:18.569308Z","iopub.status.idle":"2025-05-22T05:06:38.245591Z","shell.execute_reply.started":"2025-05-22T02:38:18.569294Z","shell.execute_reply":"2025-05-22T05:06:38.244802Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def debug_image_loading(tomo_id):\n \"\"\"\n Debug function to check image loading\n \"\"\"\n tomo_dir = os.path.join(TEST_DIR, tomo_id)\n slice_files = sorted([f for f in os.listdir(tomo_dir) if f.endswith('.jpg')])\n \n if not slice_files:\n print(f\"No image files found in {tomo_dir}\")\n return\n \n print(f\"Found {len(slice_files)} image files in {tomo_dir}\")\n sample_file = slice_files[len(slice_files)//2] # Middle slice\n img_path = os.path.join(tomo_dir, sample_file)\n \n # Try different loading methods\n try:\n # Method 1: PIL\n img_pil = Image.open(img_path)\n img_array_pil = np.array(img_pil)\n print(f\"PIL Image shape: {img_array_pil.shape}, dtype: {img_array_pil.dtype}\")\n \n # Method 2: OpenCV\n img_cv2 = cv2.imread(img_path, cv2.IMREAD_GRAYSCALE)\n print(f\"OpenCV Image shape: {img_cv2.shape}, dtype: {img_cv2.dtype}\")\n \n # Method 3: Convert to RGB\n img_rgb = cv2.cvtColor(cv2.imread(img_path), cv2.COLOR_BGR2RGB)\n print(f\"OpenCV RGB Image shape: {img_rgb.shape}, dtype: {img_rgb.dtype}\")\n \n print(\"Image loading successful!\")\n except Exception as e:\n print(f\"Error loading image {img_path}: {e}\")\n \n # Also test with YOLO's built-in loader\n try:\n test_model = YOLO(MODEL_PATH)\n test_results = test_model([img_path], verbose=False)\n print(\"YOLO model successfully processed the test image\")\n except Exception as e:\n print(f\"Error with YOLO processing: {e}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-22T07:12:05.332125Z","iopub.execute_input":"2025-05-22T07:12:05.332457Z","iopub.status.idle":"2025-05-22T07:12:05.340998Z","shell.execute_reply.started":"2025-05-22T07:12:05.332429Z","shell.execute_reply":"2025-05-22T07:12:05.339809Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def preload_image_batch(file_paths):\n \"\"\"Preload a batch of images to CPU memory\"\"\"\n images = []\n for path in file_paths:\n img = cv2.imread(path)\n if img is None:\n # Try with PIL as fallback\n img = np.array(Image.open(path))\n images.append(img)\n return images","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-22T07:12:06.973287Z","iopub.execute_input":"2025-05-22T07:12:06.973734Z","iopub.status.idle":"2025-05-22T07:12:06.979275Z","shell.execute_reply.started":"2025-05-22T07:12:06.973706Z","shell.execute_reply":"2025-05-22T07:12:06.978226Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# GPU profiling context manager\nclass GPUProfiler:\n def __init__(self, name):\n self.name = name\n self.start_time = None\n \n def __enter__(self):\n if torch.cuda.is_available():\n torch.cuda.synchronize()\n self.start_time = time.time()\n return self\n \n def __exit__(self, *args):\n if torch.cuda.is_available():\n torch.cuda.synchronize()\n elapsed = time.time() - self.start_time\n print(f\"[PROFILE] {self.name}: {elapsed:.3f}s\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-22T07:12:07.207209Z","iopub.execute_input":"2025-05-22T07:12:07.207583Z","iopub.status.idle":"2025-05-22T07:12:07.214056Z","shell.execute_reply.started":"2025-05-22T07:12:07.207561Z","shell.execute_reply":"2025-05-22T07:12:07.212824Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def perform_3d_nms(detections, iou_threshold):\n \"\"\"\n Perform 3D Non-Maximum Suppression on detections to merge nearby motors\n \"\"\"\n if not detections:\n return []\n \n # Sort by confidence (highest first)\n detections = sorted(detections, key=lambda x: x['confidence'], reverse=True)\n \n # List to store final detections after NMS\n final_detections = []\n \n # Define 3D distance function\n def distance_3d(d1, d2):\n return np.sqrt((d1['z'] - d2['z'])**2 + \n (d1['y'] - d2['y'])**2 + \n (d1['x'] - d2['x'])**2)\n \n # Maximum distance threshold (based on box size and slice gap)\n box_size = 24 # Same as annotation box size\n distance_threshold = box_size * iou_threshold\n \n # Process each detection\n while detections:\n # Take the detection with highest confidence\n best_detection = detections.pop(0)\n final_detections.append(best_detection)\n \n # Filter out detections that are too close to the best detection\n detections = [d for d in detections if distance_3d(d, best_detection) > distance_threshold]\n \n return final_detections","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-22T07:12:07.465262Z","iopub.execute_input":"2025-05-22T07:12:07.465631Z","iopub.status.idle":"2025-05-22T07:12:07.472981Z","shell.execute_reply.started":"2025-05-22T07:12:07.465609Z","shell.execute_reply":"2025-05-22T07:12:07.471896Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"BATCH_SIZE","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-22T05:57:19.529654Z","iopub.execute_input":"2025-05-22T05:57:19.529918Z","iopub.status.idle":"2025-05-22T05:57:19.534759Z","shell.execute_reply.started":"2025-05-22T05:57:19.529900Z","shell.execute_reply":"2025-05-22T05:57:19.534018Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def process_tomogram(tomo_id, model, index=0, total=1):\n \"\"\"\n Process a single tomogram and return the most confident motor detection\n \"\"\"\n print(f\"Processing tomogram {tomo_id} ({index}/{total})\")\n \n # Get all slice files for this tomogram\n tomo_dir = os.path.join(TEST_DIR, tomo_id)\n slice_files = sorted([f for f in os.listdir(tomo_dir) if f.endswith('.jpg')])\n \n # Apply CONCENTRATION to reduce the number of slices processed\n # This will process approximately CONCENTRATION fraction of all slices\n selected_indices = np.linspace(0, len(slice_files)-1, int(len(slice_files) * CONCENTRATION))\n selected_indices = np.round(selected_indices).astype(int)\n slice_files = [slice_files[i] for i in selected_indices]\n \n print(f\"Processing {len(slice_files)} out of {len(os.listdir(tomo_dir))} slices based on CONCENTRATION={CONCENTRATION}\")\n \n # Create a list to store all detections\n all_detections = []\n \n # Create CUDA streams for parallel processing if using GPU\n if DEVICE.startswith('cuda'):\n streams = [torch.cuda.Stream() for _ in range(min(4, BATCH_SIZE))]\n else:\n streams = [None]\n \n # Variables for preloading\n next_batch_thread = None\n next_batch_images = None\n \n # Process slices in batches\n for batch_start in range(0, len(slice_files), BATCH_SIZE):\n # Wait for previous preload thread if it exists\n if next_batch_thread is not None:\n next_batch_thread.join()\n next_batch_images = None\n \n batch_end = min(batch_start + BATCH_SIZE, len(slice_files))\n batch_files = slice_files[batch_start:batch_end]\n \n # Start preloading next batch\n next_batch_start = batch_end\n next_batch_end = min(next_batch_start + BATCH_SIZE, len(slice_files))\n next_batch_files = slice_files[next_batch_start:next_batch_end] if next_batch_start < len(slice_files) else []\n \n if next_batch_files:\n next_batch_paths = [os.path.join(tomo_dir, f) for f in next_batch_files]\n next_batch_thread = threading.Thread(target=preload_image_batch, args=(next_batch_paths,))\n next_batch_thread.start()\n else:\n next_batch_thread = None\n \n # Split batch across streams for parallel processing\n sub_batches = np.array_split(batch_files, len(streams))\n sub_batch_results = []\n \n for i, sub_batch in enumerate(sub_batches):\n if len(sub_batch) == 0:\n continue\n \n stream = streams[i % len(streams)]\n with torch.cuda.stream(stream) if stream and DEVICE.startswith('cuda') else nullcontext():\n # Process sub-batch\n sub_batch_paths = [os.path.join(tomo_dir, slice_file) for slice_file in sub_batch]\n sub_batch_slice_nums = [int(slice_file.split('_')[1].split('.')[0]) for slice_file in sub_batch]\n \n # Run inference with profiling\n with GPUProfiler(f\"Inference batch {i+1}/{len(sub_batches)}\"):\n sub_results = model(sub_batch_paths, verbose=False)\n \n # Process each result in this sub-batch\n for j, result in enumerate(sub_results):\n if len(result.boxes) > 0:\n boxes = result.boxes\n for box_idx, confidence in enumerate(boxes.conf):\n if confidence >= CONFIDENCE_THRESHOLD:\n # Get bounding box coordinates\n x1, y1, x2, y2 = boxes.xyxy[box_idx].cpu().numpy()\n \n # Calculate center coordinates\n x_center = (x1 + x2) / 2\n y_center = (y1 + y2) / 2\n \n # Store detection with 3D coordinates\n all_detections.append({\n 'z': round(sub_batch_slice_nums[j]),\n 'y': round(y_center),\n 'x': round(x_center),\n 'confidence': float(confidence)\n })\n \n # Synchronize streams\n if DEVICE.startswith('cuda'):\n torch.cuda.synchronize()\n \n # Clean up thread if still running\n if next_batch_thread is not None:\n next_batch_thread.join()\n \n # 3D Non-Maximum Suppression to merge nearby detections across slices\n final_detections = perform_3d_nms(all_detections, NMS_IOU_THRESHOLD)\n \n # Sort detections by confidence (highest first)\n final_detections.sort(key=lambda x: x['confidence'], reverse=True)\n \n # If there are no detections, return NA values\n if not final_detections:\n return {\n 'tomo_id': tomo_id,\n 'Motor axis 0': -1,\n 'Motor axis 1': -1,\n 'Motor axis 2': -1\n }\n \n # Take the detection with highest confidence\n best_detection = final_detections[0]\n \n # Return result with integer coordinates\n return {\n 'tomo_id': tomo_id,\n 'Motor axis 0': round(best_detection['z']),\n 'Motor axis 1': round(best_detection['y']),\n 'Motor axis 2': round(best_detection['x'])\n }\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-22T07:12:29.848437Z","iopub.execute_input":"2025-05-22T07:12:29.848807Z","iopub.status.idle":"2025-05-22T07:12:29.872970Z","shell.execute_reply.started":"2025-05-22T07:12:29.848785Z","shell.execute_reply":"2025-05-22T07:12:29.871385Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def generate_submission():\n \"\"\"\n Main function to generate the submission file\n \"\"\"\n # Get list of test tomograms\n test_tomos = sorted([d for d in os.listdir(TEST_DIR) if os.path.isdir(os.path.join(TEST_DIR, d))])\n total_tomos = len(test_tomos)\n \n print(f\"Found {total_tomos} tomograms in test directory\")\n \n # Debug image loading for the first tomogram\n if test_tomos:\n debug_image_loading(test_tomos[0])\n \n # Clear GPU cache before starting\n if torch.cuda.is_available():\n torch.cuda.empty_cache()\n \n # Initialize model once outside the processing loop\n print(f\"Loading YOLO model from {MODEL_PATH}\")\n model = YOLO(MODEL_PATH)\n model.to(DEVICE)\n \n # Additional optimizations for inference\n if DEVICE.startswith('cuda'):\n # Fuse conv and bn layers for faster inference\n model.fuse()\n \n # Enable model half precision (FP16) if on compatible GPU\n if torch.cuda.get_device_capability(0)[0] >= 7: # Volta or newer\n model.model.half()\n print(\"Using half precision (FP16) for inference\")\n \n # Process tomograms with parallelization\n results = []\n motors_found = 0\n \n # Using ThreadPoolExecutor with max_workers=1 since each worker uses the GPU already\n # and we're parallelizing within each tomogram processing\n with ThreadPoolExecutor(max_workers=1) as executor:\n future_to_tomo = {}\n \n # Submit all tomograms for processing\n for i, tomo_id in enumerate(test_tomos, 1):\n future = executor.submit(process_tomogram, tomo_id, model, i, total_tomos)\n future_to_tomo[future] = tomo_id\n \n # Process completed futures as they complete\n for future in future_to_tomo:\n tomo_id = future_to_tomo[future]\n try:\n # Clear CUDA cache between tomograms\n if torch.cuda.is_available():\n torch.cuda.empty_cache()\n \n result = future.result()\n results.append(result)\n \n # Update motors found count\n has_motor = not pd.isna(result['Motor axis 0'])\n if has_motor:\n motors_found += 1\n print(f\"Motor found in {tomo_id} at position: \"\n f\"z={result['Motor axis 0']}, y={result['Motor axis 1']}, x={result['Motor axis 2']}\")\n else:\n print(f\"No motor detected in {tomo_id}\")\n \n print(f\"Current detection rate: {motors_found}/{len(results)} ({motors_found/len(results)*100:.1f}%)\")\n \n except Exception as e:\n print(f\"Error processing {tomo_id}: {e}\")\n # Create a default entry for failed tomograms\n results.append({\n 'tomo_id': tomo_id,\n 'Motor axis 0': -1,\n 'Motor axis 1': -1,\n 'Motor axis 2': -1\n })\n \n # Create submission dataframe\n submission_df = pd.DataFrame(results)\n \n # Ensure proper column order\n submission_df = submission_df[['tomo_id', 'Motor axis 0', 'Motor axis 1', 'Motor axis 2']]\n \n # Save the submission file\n submission_df.to_csv(SUBMISSION_PATH, index=False)\n \n print(f\"\\nSubmission complete!\")\n print(f\"Motors detected: {motors_found}/{total_tomos} ({motors_found/total_tomos*100:.1f}%)\")\n print(f\"Submission saved to: {SUBMISSION_PATH}\")\n \n # Display first few rows of submission\n print(\"\\nSubmission preview:\")\n print(submission_df.head())\n \n return submission_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-22T07:12:31.401928Z","iopub.execute_input":"2025-05-22T07:12:31.402294Z","iopub.status.idle":"2025-05-22T07:12:31.414031Z","shell.execute_reply.started":"2025-05-22T07:12:31.402269Z","shell.execute_reply":"2025-05-22T07:12:31.412979Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"generate_submission()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-22T07:12:33.032378Z","iopub.execute_input":"2025-05-22T07:12:33.032915Z","iopub.status.idle":"2025-05-22T07:12:37.015848Z","shell.execute_reply.started":"2025-05-22T07:12:33.032883Z","shell.execute_reply":"2025-05-22T07:12:37.014356Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}