{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":84969,"databundleVersionId":10033515,"sourceType":"competition"}],"dockerImageVersionId":30918,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install -q zarr \n\n!pip install -q copick","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:48:29.454637Z","iopub.execute_input":"2025-04-03T07:48:29.454958Z","iopub.status.idle":"2025-04-03T07:48:43.017545Z","shell.execute_reply.started":"2025-04-03T07:48:29.454934Z","shell.execute_reply":"2025-04-03T07:48:43.016493Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Zarr Data Structure Exploration","metadata":{}},{"cell_type":"code","source":"import os\nimport glob\nimport zarr\nimport numpy as np\nimport json\nimport matplotlib.pyplot as plt\nimport warnings\nwarnings.filterwarnings(\"ignore\")\n\n# Set paths\nbase_dir = '/kaggle/input/czii-cryo-et-object-identification'\ntrain_dir = os.path.join(base_dir, 'train')\ntest_dir = os.path.join(base_dir, 'test')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:48:53.918409Z","iopub.execute_input":"2025-04-03T07:48:53.918715Z","iopub.status.idle":"2025-04-03T07:48:54.014837Z","shell.execute_reply.started":"2025-04-03T07:48:53.918672Z","shell.execute_reply":"2025-04-03T07:48:54.013993Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to explore directory structure\ndef explore_directory(directory, max_depth=3):\n \"\"\"\n Recursively explore a directory structure up to a specified depth\n \"\"\"\n result = []\n \n def _explore(dir_path, current_depth=0):\n if current_depth > max_depth:\n return\n \n for item in os.listdir(dir_path):\n path = os.path.join(dir_path, item)\n prefix = \" \" * current_depth\n \n if os.path.isdir(path):\n result.append(f\"{prefix}📁 {item}/\")\n _explore(path, current_depth + 1)\n else:\n result.append(f\"{prefix}📄 {item}\")\n \n try:\n _explore(directory)\n except Exception as e:\n result.append(f\"Error exploring {directory}: {str(e)}\")\n \n return result","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:49:02.978159Z","iopub.execute_input":"2025-04-03T07:49:02.978845Z","iopub.status.idle":"2025-04-03T07:49:02.984480Z","shell.execute_reply.started":"2025-04-03T07:49:02.978815Z","shell.execute_reply":"2025-04-03T07:49:02.983589Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to visualize data\ndef visualize_data(data, title=\"Data Visualization\", n_slices=3):\n \"\"\"\n Visualize slices of a 3D or 4D array\n \"\"\"\n if data is None:\n print(\"No data to visualize\")\n return\n \n # Handle different data dimensions\n if len(data.shape) == 4:\n print(f\"4D data of shape {data.shape}, taking first channel/time point\")\n data = data[0] # Take first channel/time point\n \n if len(data.shape) != 3:\n print(f\"Cannot visualize data of shape {data.shape}\")\n return\n \n # Get dimensions\n depth, height, width = data.shape\n print(f\"Data dimensions: {depth} x {height} x {width}\")\n \n # Choose slice indices\n slice_indices = np.linspace(depth // 4, 3 * depth // 4, n_slices).astype(int)\n \n # Create figure\n fig, axes = plt.subplots(1, n_slices, figsize=(5 * n_slices, 5))\n if n_slices == 1:\n axes = [axes]\n \n # Plot each slice\n for i, slice_idx in enumerate(slice_indices):\n im = axes[i].imshow(data[slice_idx], cmap='gray')\n axes[i].set_title(f'Slice {slice_idx}/{depth}')\n axes[i].axis('off')\n fig.colorbar(im, ax=axes[i], fraction=0.046, pad=0.04)\n \n plt.suptitle(title)\n plt.tight_layout()\n plt.show()\n \n # Histogram of values\n plt.figure(figsize=(10, 6))\n plt.hist(data.flatten(), bins=100, alpha=0.7, color='blue')\n plt.title(f\"Value Distribution - {title}\")\n plt.xlabel(\"Value\")\n plt.ylabel(\"Frequency\")\n plt.grid(True, alpha=0.3)\n plt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:49:08.275590Z","iopub.execute_input":"2025-04-03T07:49:08.275931Z","iopub.status.idle":"2025-04-03T07:49:08.284406Z","shell.execute_reply.started":"2025-04-03T07:49:08.275906Z","shell.execute_reply":"2025-04-03T07:49:08.283723Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to explore a zarr file's structure\ndef explore_zarr(zarr_path, max_items=5):\n \"\"\"\n Explore a zarr file's structure and metadata\n \"\"\"\n print(f\"Exploring zarr at: {zarr_path}\")\n \n # First, explore the file system structure\n print(\"\\nFile system structure:\")\n structure = explore_directory(zarr_path, max_depth=4)\n for line in structure[:30]: # Limit to first 30 lines\n print(line)\n if len(structure) > 30:\n print(\"... (truncated)\")\n \n # Check for .zattrs file which contains metadata\n zattrs_path = os.path.join(zarr_path, '.zattrs')\n if os.path.exists(zattrs_path):\n print(\"\\nFound .zattrs file. Contents:\")\n try:\n with open(zattrs_path, 'r') as f:\n zattrs = json.load(f)\n print(json.dumps(zattrs, indent=2)[:1000] + \"...\" if len(json.dumps(zattrs)) > 1000 else \"\")\n \n # Check for multiscales metadata\n if 'multiscales' in zattrs:\n print(\"\\nMultiscales metadata:\")\n multiscales = zattrs['multiscales']\n print(json.dumps(multiscales, indent=2)[:1000] + \"...\" if len(json.dumps(multiscales)) > 1000 else \"\")\n \n # Extract dataset paths\n if multiscales and 'datasets' in multiscales[0]:\n datasets = multiscales[0]['datasets']\n print(f\"\\nDataset paths: {[d.get('path') for d in datasets]}\")\n except Exception as e:\n print(f\"Error reading .zattrs: {str(e)}\")\n \n # Look for .zarray files which define arrays\n zarray_files = glob.glob(os.path.join(zarr_path, \"**\", \".zarray\"), recursive=True)\n if zarray_files:\n print(f\"\\nFound {len(zarray_files)} .zarray files:\")\n for i, zarray_file in enumerate(zarray_files[:max_items]):\n print(f\" {i+1}. {os.path.relpath(zarray_file, zarr_path)}\")\n try:\n with open(zarray_file, 'r') as f:\n zarray = json.load(f)\n print(f\" Shape: {zarray.get('shape')}\")\n print(f\" Chunks: {zarray.get('chunks')}\")\n print(f\" Data type: {zarray.get('dtype')}\")\n except Exception as e:\n print(f\" Error reading .zarray: {str(e)}\")\n \n # Try to open the zarr file\n try:\n z = zarr.open(zarr_path, mode='r')\n \n # Check if it's a Group\n if isinstance(z, zarr.Group):\n print(\"\\nZarr opened as a Group\")\n print(f\"Keys: {list(z.keys())}\")\n \n # Explore each key\n for key in list(z.keys())[:max_items]:\n print(f\"\\nExploring group '{key}':\")\n try:\n item = z[key]\n if isinstance(item, zarr.Array):\n print(f\" Array shape: {item.shape}\")\n print(f\" Data type: {item.dtype}\")\n print(f\" Chunks: {item.chunks}\")\n elif isinstance(item, zarr.Group):\n print(f\" Subgroup with keys: {list(item.keys())}\")\n except Exception as e:\n print(f\" Error exploring group: {str(e)}\")\n \n # Check if it's an Array\n elif isinstance(z, zarr.Array):\n print(\"\\nZarr opened as an Array\")\n print(f\"Shape: {z.shape}\")\n print(f\"Data type: {z.dtype}\")\n print(f\"Chunks: {z.chunks}\")\n \n # Print some sample data if small enough\n if np.prod(z.shape) < 100:\n print(f\"Data sample: {z[:]}\")\n else:\n print(\"Data too large to display sample\")\n \n else:\n print(f\"\\nZarr opened as unknown type: {type(z)}\")\n \n except Exception as e:\n print(f\"\\nError opening zarr: {str(e)}\")\n \n return\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:49:13.443098Z","iopub.execute_input":"2025-04-03T07:49:13.443389Z","iopub.status.idle":"2025-04-03T07:49:13.500184Z","shell.execute_reply.started":"2025-04-03T07:49:13.443368Z","shell.execute_reply":"2025-04-03T07:49:13.499188Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to try loading and visualizing data from a zarr file\ndef try_load_visualize(zarr_path):\n \"\"\"\n Try different approaches to load and visualize data from a zarr file\n \"\"\"\n print(f\"Attempting to load and visualize data from: {zarr_path}\")\n \n # Try approach 1: Direct array loading if zarr_path points to a zarr array\n try:\n print(\"\\nApproach 1: Direct array loading\")\n z1 = zarr.open(zarr_path, mode='r')\n if isinstance(z1, zarr.Array):\n print(f\"Found array of shape: {z1.shape}\")\n data1 = z1[:]\n visualize_data(data1, title=\"Approach 1: Direct array\")\n return data1\n except Exception as e:\n print(f\"Approach 1 failed: {str(e)}\")\n \n # Try approach 2: Using multiscales metadata if available\n try:\n print(\"\\nApproach 2: Using multiscales metadata\")\n zattrs_path = os.path.join(zarr_path, '.zattrs')\n if os.path.exists(zattrs_path):\n with open(zattrs_path, 'r') as f:\n zattrs = json.load(f)\n \n if 'multiscales' in zattrs and 'datasets' in zattrs['multiscales'][0]:\n datasets = zattrs['multiscales'][0]['datasets']\n if datasets:\n dataset_path = datasets[0].get('path')\n if dataset_path:\n full_path = os.path.join(zarr_path, dataset_path)\n print(f\"Loading dataset from: {full_path}\")\n z2 = zarr.open(full_path, mode='r')\n if isinstance(z2, zarr.Array):\n print(f\"Found array of shape: {z2.shape}\")\n data2 = z2[:]\n visualize_data(data2, title=\"Approach 2: Multiscales\")\n return data2\n except Exception as e:\n print(f\"Approach 2 failed: {str(e)}\")\n \n # Try approach 3: Find .zarray files and try to load data from their parent directories\n try:\n print(\"\\nApproach 3: Find .zarray files and load from their parent directories\")\n zarray_files = glob.glob(os.path.join(zarr_path, \"**\", \".zarray\"), recursive=True)\n if zarray_files:\n for i, zarray_file in enumerate(zarray_files[:3]): # Try first 3 .zarray files\n data_dir = os.path.dirname(zarray_file)\n print(f\"Trying to load from: {data_dir}\")\n try:\n z3 = zarr.open(data_dir, mode='r')\n if isinstance(z3, zarr.Array):\n print(f\"Found array of shape: {z3.shape}\")\n data3 = z3[:]\n visualize_data(data3, title=f\"Approach 3: Array {i+1}\")\n return data3\n except Exception as e:\n print(f\"Failed to load from {data_dir}: {str(e)}\")\n except Exception as e:\n print(f\"Approach 3 failed: {str(e)}\")\n \n # Try approach 4: Traverse numeric subdirectories to find data\n try:\n print(\"\\nApproach 4: Traverse numeric subdirectories\")\n # Open zarr root\n z4 = zarr.open(zarr_path, mode='r')\n \n # Check if it's a Group\n if isinstance(z4, zarr.Group):\n # Look for numeric keys that might contain the data\n numeric_keys = [k for k in z4.keys() if k.isdigit()]\n if numeric_keys:\n print(f\"Found numeric keys: {numeric_keys}\")\n \n # Try to traverse the structure through a chain of numeric keys\n current_group = z4\n traversal_path = []\n \n # Try up to depth 5\n for _ in range(5):\n numeric_keys = [k for k in current_group.keys() if k.isdigit()]\n if not numeric_keys:\n break\n \n next_key = numeric_keys[0] # Take first numeric key\n traversal_path.append(next_key)\n \n current_group = current_group[next_key]\n if isinstance(current_group, zarr.Array):\n print(f\"Found array at path: {'/'.join(traversal_path)}\")\n print(f\"Array shape: {current_group.shape}\")\n data4 = current_group[:]\n visualize_data(data4, title=\"Approach 4: Numeric traversal\")\n return data4\n except Exception as e:\n print(f\"Approach 4 failed: {str(e)}\")\n \n print(\"\\nAll approaches failed to load data\")\n return None\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:49:20.298876Z","iopub.execute_input":"2025-04-03T07:49:20.299179Z","iopub.status.idle":"2025-04-03T07:49:20.310065Z","shell.execute_reply.started":"2025-04-03T07:49:20.299155Z","shell.execute_reply":"2025-04-03T07:49:20.309260Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Main exploration\nprint(\"Exploring zarr file structure in the dataset\")\n\n# List all denoised.zarr files in training data\ndenoised_zarrs = glob.glob(os.path.join(train_dir, 'static/ExperimentRuns/*/VoxelSpacing10.000/denoised.zarr'))\nprint(f\"Found {len(denoised_zarrs)} denoised.zarr files in training data:\")\nfor i, zarr_path in enumerate(denoised_zarrs):\n parts = zarr_path.split('/')\n experiment = parts[-3]\n print(f\" {i+1}. {experiment}\")\n\n# Choose first experiment for exploration\nif denoised_zarrs:\n sample_zarr_path = denoised_zarrs[0]\n print(f\"\\nExploring sample zarr file: {sample_zarr_path}\")\n \n # Explore zarr structure\n explore_zarr(sample_zarr_path)\n \n # Try to load and visualize data\n print(\"\\nAttempting to load and visualize data...\")\n try_load_visualize(sample_zarr_path)\nelse:\n print(\"No denoised.zarr files found in training data\")\n\n# Explore a test zarr file as well\ntest_zarrs = glob.glob(os.path.join(test_dir, 'static/ExperimentRuns/*/VoxelSpacing10.000/denoised.zarr'))\nif test_zarrs:\n test_zarr_path = test_zarrs[0]\n print(f\"\\nExploring test zarr file: {test_zarr_path}\")\n \n # Just do a basic exploration\n explore_zarr(test_zarr_path)\nelse:\n print(\"\\nNo denoised.zarr files found in test data\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:49:33.515654Z","iopub.execute_input":"2025-04-03T07:49:33.515955Z","iopub.status.idle":"2025-04-03T07:49:38.477065Z","shell.execute_reply.started":"2025-04-03T07:49:33.515935Z","shell.execute_reply":"2025-04-03T07:49:38.476311Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# JSON Particle Data Exploration","metadata":{}},{"cell_type":"code","source":"import os\nimport glob\nimport json\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport numpy as np\nfrom collections import defaultdict\n\n# Set paths\nbase_dir = '/kaggle/input/czii-cryo-et-object-identification'\ntrain_dir = os.path.join(base_dir, 'train')\n\n# Particle types and their properties\nparticle_types = {\n 'apo-ferritin': {'difficulty': 'easy', 'weight': 1},\n 'beta-amylase': {'difficulty': 'impossible', 'weight': 0},\n 'beta-galactosidase': {'difficulty': 'hard', 'weight': 2},\n 'ribosome': {'difficulty': 'easy', 'weight': 1},\n 'thyroglobulin': {'difficulty': 'hard', 'weight': 2},\n 'virus-like-particle': {'difficulty': 'easy', 'weight': 1}\n}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:50:14.594527Z","iopub.execute_input":"2025-04-03T07:50:14.594916Z","iopub.status.idle":"2025-04-03T07:50:15.556433Z","shell.execute_reply.started":"2025-04-03T07:50:14.594887Z","shell.execute_reply":"2025-04-03T07:50:15.555801Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Find all JSON files with particle annotations\nprint(\"Looking for particle annotation JSON files...\")\nparticle_jsons = glob.glob(os.path.join(train_dir, 'overlay/ExperimentRuns/*/Picks/*.json'))\nprint(f\"Found {len(particle_jsons)} JSON files\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:50:27.585245Z","iopub.execute_input":"2025-04-03T07:50:27.585868Z","iopub.status.idle":"2025-04-03T07:50:27.649399Z","shell.execute_reply.started":"2025-04-03T07:50:27.585838Z","shell.execute_reply":"2025-04-03T07:50:27.648773Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# List the first few files\nfor i, json_path in enumerate(particle_jsons[:10]):\n # Extract experiment and particle type from path\n parts = json_path.split('/')\n experiment = parts[-3]\n particle_type = os.path.splitext(os.path.basename(json_path))[0]\n print(f\" {i+1}. {experiment} - {particle_type}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:50:37.489199Z","iopub.execute_input":"2025-04-03T07:50:37.489479Z","iopub.status.idle":"2025-04-03T07:50:37.496179Z","shell.execute_reply.started":"2025-04-03T07:50:37.489459Z","shell.execute_reply":"2025-04-03T07:50:37.495405Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to explore a JSON file's structure\ndef explore_json(json_path):\n \"\"\"\n Explore the structure of a JSON file containing particle annotations\n \"\"\"\n print(f\"Exploring JSON file: {json_path}\")\n try:\n with open(json_path, 'r') as f:\n data = json.load(f)\n \n # Print overall structure\n print(\"\\nTop level keys:\")\n for key in data.keys():\n if isinstance(data[key], list):\n print(f\" {key}: list with {len(data[key])} items\")\n elif isinstance(data[key], dict):\n print(f\" {key}: dict with {len(data[key])} keys\")\n else:\n print(f\" {key}: {type(data[key]).__name__}\")\n \n # Check for 'picks' which should contain particle coordinates\n if 'picks' in data:\n picks = data['picks']\n print(f\"\\nFound {len(picks)} picks\")\n \n # Look at the first few picks\n if picks:\n print(\"\\nSample pick entries:\")\n for i, pick in enumerate(picks[:3]):\n print(f\" Pick {i+1}:\")\n print(json.dumps(pick, indent=4))\n \n # Extract a sample of the coordinates\n coords = []\n for pick in picks[:100]: # Limit to first 100 picks\n if 'position' in pick:\n pos = pick['position']\n if all(k in pos for k in ['x', 'y', 'z']):\n coords.append((pos['x'], pos['y'], pos['z']))\n \n if coords:\n print(f\"\\nExtracted {len(coords)} coordinates\")\n \n # Basic statistics\n coords_array = np.array(coords)\n print(\"\\nCoordinate statistics:\")\n print(f\" X: min={coords_array[:, 0].min():.2f}, max={coords_array[:, 0].max():.2f}, mean={coords_array[:, 0].mean():.2f}\")\n print(f\" Y: min={coords_array[:, 1].min():.2f}, max={coords_array[:, 1].max():.2f}, mean={coords_array[:, 1].mean():.2f}\")\n print(f\" Z: min={coords_array[:, 2].min():.2f}, max={coords_array[:, 2].max():.2f}, mean={coords_array[:, 2].mean():.2f}\")\n \n return True, coords\n else:\n print(\"Could not extract valid coordinates from pick entries\")\n else:\n print(\"No pick entries found\")\n else:\n print(\"No 'picks' key found in JSON\")\n \n except Exception as e:\n print(f\"Error exploring JSON: {str(e)}\")\n \n return False, []","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:50:43.526030Z","iopub.execute_input":"2025-04-03T07:50:43.526310Z","iopub.status.idle":"2025-04-03T07:50:43.534405Z","shell.execute_reply.started":"2025-04-03T07:50:43.526290Z","shell.execute_reply":"2025-04-03T07:50:43.533631Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Collect all particle coordinates in a DataFrame for analysis\nall_coords = []\n\n# Try to process a few JSON files\nfor json_path in particle_jsons[:10]: # Process first 10 files\n # Extract experiment and particle type from path\n parts = json_path.split('/')\n experiment = parts[-3]\n particle_type = os.path.splitext(os.path.basename(json_path))[0]\n \n print(f\"\\n{'='*80}\\nProcessing {experiment} - {particle_type}\\n{'='*80}\")\n success, coords = explore_json(json_path)\n \n if success and coords:\n # Add experiment and particle type info to each coordinate\n for x, y, z in coords:\n all_coords.append({\n 'experiment': experiment,\n 'particle_type': particle_type,\n 'x': x,\n 'y': y,\n 'z': z,\n 'difficulty': particle_types[particle_type]['difficulty'],\n 'weight': particle_types[particle_type]['weight']\n })\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:50:51.212280Z","iopub.execute_input":"2025-04-03T07:50:51.212557Z","iopub.status.idle":"2025-04-03T07:50:51.286131Z","shell.execute_reply.started":"2025-04-03T07:50:51.212538Z","shell.execute_reply":"2025-04-03T07:50:51.285390Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Convert to DataFrame\ncoords_df = pd.DataFrame(all_coords)\n\nif not coords_df.empty:\n print(f\"\\nCollected {len(coords_df)} particle coordinates\")\n print(\"\\nCoordinates DataFrame sample:\")\n print(coords_df.head())\n \n # Count particles by type and experiment\n particle_counts = coords_df.groupby(['experiment', 'particle_type']).size().reset_index(name='count')\n print(\"\\nParticle counts by experiment and type:\")\n print(particle_counts)\n \n # Visualize particle counts by type\n plt.figure(figsize=(12, 6))\n type_counts = coords_df['particle_type'].value_counts()\n plt.bar(type_counts.index, type_counts.values)\n plt.title('Number of Particles by Type')\n plt.xticks(rotation=45, ha='right')\n plt.tight_layout()\n plt.show()\n \n # Plot coordinates in 3D for a single experiment\n experiments = coords_df['experiment'].unique()\n if len(experiments) > 0:\n exp_to_plot = experiments[0]\n exp_coords = coords_df[coords_df['experiment'] == exp_to_plot]\n \n if len(exp_coords) > 0:\n fig = plt.figure(figsize=(10, 8))\n ax = fig.add_subplot(111, projection='3d')\n \n # Plot each particle type with a different color\n for p_type, group in exp_coords.groupby('particle_type'):\n ax.scatter(group['x'], group['y'], group['z'], \n label=p_type, \n alpha=0.7, s=10)\n \n ax.set_xlabel('X')\n ax.set_ylabel('Y')\n ax.set_zlabel('Z')\n ax.set_title(f'3D Distribution of Particles in Experiment {exp_to_plot}')\n plt.legend()\n plt.tight_layout()\n plt.show()\nelse:\n print(\"\\nFailed to collect any valid particle coordinates\")\n \n # Try a more direct approach\n print(\"\\nTrying a direct approach to read one JSON file...\")\n if particle_jsons:\n first_json = particle_jsons[0]\n print(f\"Reading: {first_json}\")\n \n try:\n with open(first_json, 'r') as f:\n content = f.read()\n \n print(f\"File size: {len(content)} bytes\")\n print(\"First 1000 characters:\")\n print(content[:1000])\n \n # Try parsing the JSON\n try:\n data = json.loads(content)\n print(\"\\nSuccessfully parsed JSON\")\n print(f\"Top level keys: {list(data.keys())}\")\n \n # Fully print a small part of the structure\n if 'picks' in data and data['picks']:\n first_pick = data['picks'][0]\n print(\"\\nFirst pick entry:\")\n print(json.dumps(first_pick, indent=2))\n except json.JSONDecodeError as e:\n print(f\"\\nJSON parsing error: {str(e)}\")\n except Exception as e:\n print(f\"Error reading file: {str(e)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:51:08.724541Z","iopub.execute_input":"2025-04-03T07:51:08.724855Z","iopub.status.idle":"2025-04-03T07:51:08.749830Z","shell.execute_reply.started":"2025-04-03T07:51:08.724824Z","shell.execute_reply":"2025-04-03T07:51:08.748942Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Submission Format Analysis","metadata":{}},{"cell_type":"code","source":"import os\nimport pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\n\n# Set paths\nbase_dir = '/kaggle/input/czii-cryo-et-object-identification'\n\n# Load the sample submission file\nsample_submission_path = os.path.join(base_dir, 'sample_submission.csv')\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:51:18.567068Z","iopub.execute_input":"2025-04-03T07:51:18.567352Z","iopub.status.idle":"2025-04-03T07:51:18.571349Z","shell.execute_reply.started":"2025-04-03T07:51:18.567332Z","shell.execute_reply":"2025-04-03T07:51:18.570632Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if os.path.exists(sample_submission_path):\n print(f\"Loading sample submission from: {sample_submission_path}\")\n sample_submission = pd.read_csv(sample_submission_path)\n \n # Print basic info\n print(f\"\\nSample submission shape: {sample_submission.shape}\")\n print(\"\\nSample submission columns:\")\n for col in sample_submission.columns:\n print(f\" - {col}\")\n \n print(\"\\nFirst 10 rows of the sample submission:\")\n print(sample_submission.head(10))\n \n # Analyze the unique values in each column\n print(\"\\nUnique values in each column:\")\n for col in sample_submission.columns:\n unique_vals = sample_submission[col].nunique()\n print(f\" - {col}: {unique_vals} unique values\")\n \n # For columns with few unique values, print them\n if unique_vals < 10 and col != 'id':\n print(f\" Values: {sorted(sample_submission[col].unique())}\")\n \n # Check for missing values\n missing = sample_submission.isnull().sum()\n if missing.sum() > 0:\n print(\"\\nMissing values in sample submission:\")\n for col, count in missing.items():\n if count > 0:\n print(f\" - {col}: {count} missing values\")\n else:\n print(\"\\nNo missing values in the sample submission\")\n \n # Analyze coordinates\n print(\"\\nCoordinate statistics:\")\n for col in ['x', 'y', 'z']:\n if col in sample_submission.columns:\n print(f\" - {col}:\")\n print(f\" Min: {sample_submission[col].min()}\")\n print(f\" Max: {sample_submission[col].max()}\")\n print(f\" Mean: {sample_submission[col].mean()}\")\n print(f\" Std: {sample_submission[col].std()}\")\n \n # Count rows per experiment and particle type\n if 'experiment' in sample_submission.columns and 'particle_type' in sample_submission.columns:\n exp_counts = sample_submission.groupby('experiment').size()\n print(\"\\nSubmission entries per experiment:\")\n for exp, count in exp_counts.items():\n print(f\" - {exp}: {count} entries\")\n \n type_counts = sample_submission.groupby('particle_type').size()\n print(\"\\nSubmission entries per particle type:\")\n for p_type, count in type_counts.items():\n print(f\" - {p_type}: {count} entries\")\n \n # Cross-tabulation of experiment and particle type\n cross_tab = pd.crosstab(sample_submission['experiment'], sample_submission['particle_type'])\n print(\"\\nCross-tabulation of experiment and particle type:\")\n print(cross_tab)\n \n # Visualize particle types in sample submission\n plt.figure(figsize=(10, 6))\n type_counts.plot(kind='bar')\n plt.title('Particle Types in Sample Submission')\n plt.xlabel('Particle Type')\n plt.ylabel('Count')\n plt.xticks(rotation=45, ha='right')\n plt.tight_layout()\n plt.show()\nelse:\n print(f\"Sample submission file not found at: {sample_submission_path}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:51:25.881958Z","iopub.execute_input":"2025-04-03T07:51:25.882235Z","iopub.status.idle":"2025-04-03T07:51:26.174018Z","shell.execute_reply.started":"2025-04-03T07:51:25.882214Z","shell.execute_reply":"2025-04-03T07:51:26.173175Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# EDA","metadata":{}},{"cell_type":"markdown","source":"# Dataset Structure","metadata":{}},{"cell_type":"code","source":"import os\nimport glob\nimport pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom mpl_toolkits.mplot3d import Axes3D\nimport zarr\nimport json\nfrom tqdm.notebook import tqdm\nimport seaborn as sns","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:51:41.608223Z","iopub.execute_input":"2025-04-03T07:51:41.608510Z","iopub.status.idle":"2025-04-03T07:51:42.950740Z","shell.execute_reply.started":"2025-04-03T07:51:41.608489Z","shell.execute_reply":"2025-04-03T07:51:42.950083Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Set paths to the data directories\nbase_dir = '/kaggle/input/czii-cryo-et-object-identification' \ntrain_dir = os.path.join(base_dir, 'train')\ntest_dir = os.path.join(base_dir, 'test')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:51:47.368382Z","iopub.execute_input":"2025-04-03T07:51:47.368941Z","iopub.status.idle":"2025-04-03T07:51:47.372953Z","shell.execute_reply.started":"2025-04-03T07:51:47.368913Z","shell.execute_reply":"2025-04-03T07:51:47.371971Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to explore directory structure\ndef explore_directory(directory, max_depth=3, current_depth=0):\n if current_depth > max_depth:\n return []\n \n result = []\n try:\n for item in os.listdir(directory):\n path = os.path.join(directory, item)\n if os.path.isdir(path):\n result.append(f\"{' ' * current_depth}📁 {item}/\")\n result.extend(explore_directory(path, max_depth, current_depth + 1))\n else:\n result.append(f\"{' ' * current_depth}📄 {item}\")\n except Exception as e:\n result.append(f\"{' ' * current_depth}❌ Error: {str(e)}\")\n \n return result\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:51:51.110870Z","iopub.execute_input":"2025-04-03T07:51:51.111156Z","iopub.status.idle":"2025-04-03T07:51:51.116404Z","shell.execute_reply.started":"2025-04-03T07:51:51.111136Z","shell.execute_reply":"2025-04-03T07:51:51.115423Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Explore train directory structure\nprint(\"Train Directory Structure:\")\ntrain_structure = explore_directory(train_dir, max_depth=4)\nfor line in train_structure[:50]: # Limit output to first 50 lines\n print(line)\nprint(\"...\" if len(train_structure) > 50 else \"\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:51:56.076014Z","iopub.execute_input":"2025-04-03T07:51:56.076329Z","iopub.status.idle":"2025-04-03T07:51:56.208662Z","shell.execute_reply.started":"2025-04-03T07:51:56.076302Z","shell.execute_reply":"2025-04-03T07:51:56.207915Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Explore test directory structure\nprint(\"\\nTest Directory Structure:\")\ntest_structure = explore_directory(test_dir, max_depth=4)\nfor line in test_structure[:50]: # Limit output to first 50 lines\n print(line)\nprint(\"...\" if len(test_structure) > 50 else \"\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:52:04.620153Z","iopub.execute_input":"2025-04-03T07:52:04.620440Z","iopub.status.idle":"2025-04-03T07:52:04.653448Z","shell.execute_reply.started":"2025-04-03T07:52:04.620420Z","shell.execute_reply":"2025-04-03T07:52:04.652570Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# List all experiments in train and test\ntrain_experiments = [os.path.basename(p) for p in glob.glob(os.path.join(train_dir, 'static/ExperimentRuns/*'))]\ntest_experiments = [os.path.basename(p) for p in glob.glob(os.path.join(test_dir, 'static/ExperimentRuns/*'))]\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:52:13.544308Z","iopub.execute_input":"2025-04-03T07:52:13.544661Z","iopub.status.idle":"2025-04-03T07:52:13.550947Z","shell.execute_reply.started":"2025-04-03T07:52:13.544634Z","shell.execute_reply":"2025-04-03T07:52:13.550206Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(f\"\\nNumber of experiments in training data: {len(train_experiments)}\")\nprint(f\"Number of experiments in test data: {len(test_experiments)}\")\n\nprint(\"\\nSample of training experiments:\")\nfor exp in train_experiments[:5]:\n print(f\" - {exp}\")\n\nprint(\"\\nSample of test experiments:\")\nfor exp in test_experiments[:5]:\n print(f\" - {exp}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:52:18.341775Z","iopub.execute_input":"2025-04-03T07:52:18.342054Z","iopub.status.idle":"2025-04-03T07:52:18.348738Z","shell.execute_reply.started":"2025-04-03T07:52:18.342034Z","shell.execute_reply":"2025-04-03T07:52:18.348019Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Check sample submission format\nsample_submission_path = os.path.join(base_dir, 'sample_submission.csv')\nif os.path.exists(sample_submission_path):\n sample_submission = pd.read_csv(sample_submission_path)\n print(\"\\nSample Submission Format:\")\n print(sample_submission.head())\n print(f\"Sample submission shape: {sample_submission.shape}\")\n print(f\"Sample submission columns: {sample_submission.columns.tolist()}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:52:23.432557Z","iopub.execute_input":"2025-04-03T07:52:23.432908Z","iopub.status.idle":"2025-04-03T07:52:23.445111Z","shell.execute_reply.started":"2025-04-03T07:52:23.432882Z","shell.execute_reply":"2025-04-03T07:52:23.444255Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Particle Coordinate Extraction","metadata":{}},{"cell_type":"code","source":"import os\nimport glob\nimport json\nimport pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nfrom mpl_toolkits.mplot3d import Axes3D\nimport seaborn as sns\nfrom collections import Counter, defaultdict","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:52:33.144976Z","iopub.execute_input":"2025-04-03T07:52:33.145351Z","iopub.status.idle":"2025-04-03T07:52:33.149530Z","shell.execute_reply.started":"2025-04-03T07:52:33.145324Z","shell.execute_reply":"2025-04-03T07:52:33.148643Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Set paths\nbase_dir = '/kaggle/input/czii-cryo-et-object-identification'\ntrain_dir = os.path.join(base_dir, 'train')\n\n# Particle types and their properties\nparticle_types = {\n 'apo-ferritin': {'difficulty': 'easy', 'weight': 1, 'color': 'red'},\n 'beta-amylase': {'difficulty': 'impossible', 'weight': 0, 'color': 'gray'},\n 'beta-galactosidase': {'difficulty': 'hard', 'weight': 2, 'color': 'blue'},\n 'ribosome': {'difficulty': 'easy', 'weight': 1, 'color': 'green'},\n 'thyroglobulin': {'difficulty': 'hard', 'weight': 2, 'color': 'purple'},\n 'virus-like-particle': {'difficulty': 'easy', 'weight': 1, 'color': 'orange'}\n}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:52:40.246161Z","iopub.execute_input":"2025-04-03T07:52:40.246482Z","iopub.status.idle":"2025-04-03T07:52:40.251120Z","shell.execute_reply.started":"2025-04-03T07:52:40.246457Z","shell.execute_reply":"2025-04-03T07:52:40.250238Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to load particle coordinates from JSON with the correct structure\ndef load_particle_coords(json_path):\n \"\"\"\n Load particle coordinates from JSON file\n Uses 'points' key and 'location' field which is the actual structure in the data\n \"\"\"\n try:\n with open(json_path, 'r') as f:\n data = json.load(f)\n \n # Extract coordinates from points\n coords = []\n if 'points' in data:\n for point in data['points']:\n if 'location' in point:\n loc = point['location']\n coords.append((loc.get('x', 0), loc.get('y', 0), loc.get('z', 0)))\n \n print(f\"Extracted {len(coords)} coordinates from {os.path.basename(json_path)}\")\n return coords\n else:\n print(f\"No 'points' key found in {json_path}\")\n return []\n except Exception as e:\n print(f\"Error loading {json_path}: {str(e)}\")\n return []\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:52:46.745127Z","iopub.execute_input":"2025-04-03T07:52:46.745449Z","iopub.status.idle":"2025-04-03T07:52:46.750723Z","shell.execute_reply.started":"2025-04-03T07:52:46.745423Z","shell.execute_reply":"2025-04-03T07:52:46.749959Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Collect particle coordinates data\nall_particle_data = []\nexperiment_counts = defaultdict(lambda: defaultdict(int))\nparticle_counts = Counter()\nexperiments = []\n\n# Find all JSON files with particle annotations\njson_files = glob.glob(os.path.join(train_dir, 'overlay/ExperimentRuns/*/Picks/*.json'))\nprint(f\"Found {len(json_files)} JSON files\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:52:51.260824Z","iopub.execute_input":"2025-04-03T07:52:51.261123Z","iopub.status.idle":"2025-04-03T07:52:51.276238Z","shell.execute_reply.started":"2025-04-03T07:52:51.261103Z","shell.execute_reply":"2025-04-03T07:52:51.275296Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Process all JSON files\nfor json_path in json_files:\n # Extract experiment and particle type from path\n parts = json_path.split('/')\n experiment = parts[-3]\n particle_type = os.path.splitext(os.path.basename(json_path))[0]\n \n # Load coordinates\n coords = load_particle_coords(json_path)\n count = len(coords)\n \n # Update counts\n if experiment not in experiments:\n experiments.append(experiment)\n \n experiment_counts[experiment][particle_type] = count\n particle_counts[particle_type] += count\n \n # Add coordinates to dataset\n for x, y, z in coords:\n all_particle_data.append({\n 'experiment': experiment,\n 'particle_type': particle_type,\n 'x': x,\n 'y': y,\n 'z': z,\n 'difficulty': particle_types[particle_type]['difficulty'],\n 'weight': particle_types[particle_type]['weight']\n })\n\n# Convert to DataFrame for easier analysis\nparticle_df = pd.DataFrame(all_particle_data)\nexperiment_df = pd.DataFrame(experiment_counts).T.fillna(0)\nexperiment_df = experiment_df.astype(int)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:53:08.352672Z","iopub.execute_input":"2025-04-03T07:53:08.353016Z","iopub.status.idle":"2025-04-03T07:53:08.433768Z","shell.execute_reply.started":"2025-04-03T07:53:08.352994Z","shell.execute_reply":"2025-04-03T07:53:08.432943Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Print overall statistics\nprint(f\"\\nTotal number of experiments: {len(experiments)}\")\nprint(f\"Total number of particles found: {len(particle_df)}\")\nprint(\"\\nParticle distribution:\")\nfor particle, count in particle_counts.most_common():\n difficulty = particle_types.get(particle, {}).get('difficulty', 'unknown')\n weight = particle_types.get(particle, {}).get('weight', 'unknown')\n print(f\" - {particle}: {count} particles (Difficulty: {difficulty}, Weight: {weight})\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:53:15.346259Z","iopub.execute_input":"2025-04-03T07:53:15.346562Z","iopub.status.idle":"2025-04-03T07:53:15.352835Z","shell.execute_reply.started":"2025-04-03T07:53:15.346541Z","shell.execute_reply":"2025-04-03T07:53:15.351989Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if not particle_df.empty:\n # Plot particle distribution\n plt.figure(figsize=(12, 6))\n type_counts = particle_df['particle_type'].value_counts()\n bars = plt.bar(type_counts.index, type_counts.values, \n color=[particle_types[p]['color'] for p in type_counts.index])\n \n # Add difficulty labels\n for i, particle in enumerate(type_counts.index):\n difficulty = particle_types.get(particle, {}).get('difficulty', 'unknown')\n weight = particle_types.get(particle, {}).get('weight', 'unknown')\n plt.text(i, type_counts[particle] + 5, f\"{difficulty}\\nweight: {weight}\", \n ha='center', va='bottom', fontweight='bold')\n \n plt.title('Particle Distribution in Training Data')\n plt.ylabel('Number of Particles')\n plt.xticks(rotation=45, ha='right')\n plt.tight_layout()\n plt.show()\n \n # Particle counts per experiment\n experiment_df['total'] = experiment_df.sum(axis=1)\n \n print(\"\\nParticle counts per experiment:\")\n print(experiment_df)\n \n print(\"\\nParticle count statistics per experiment:\")\n print(experiment_df.describe())\n \n # Plot total particles per experiment\n plt.figure(figsize=(10, 6))\n experiment_df['total'].plot(kind='bar', color='teal')\n plt.title('Total Particles per Experiment')\n plt.xlabel('Experiment')\n plt.ylabel('Number of Particles')\n plt.xticks(rotation=45, ha='right')\n plt.grid(axis='y', alpha=0.3)\n plt.tight_layout()\n plt.show()\n \n # Plot distribution of particles by type and experiment\n plt.figure(figsize=(14, 8))\n experiment_df_subset = experiment_df.drop(columns=['total'])\n experiment_df_subset.plot(kind='bar', \n color=[particle_types[col]['color'] for col in experiment_df_subset.columns])\n plt.title('Particle Distribution by Type and Experiment')\n plt.xlabel('Experiment')\n plt.ylabel('Number of Particles')\n plt.grid(axis='y', alpha=0.3)\n plt.legend(title='Particle Type')\n plt.tight_layout()\n plt.show()\n \n # 3D scatter plot of one experiment\n if len(experiments) > 0:\n sample_exp = experiments[0]\n sample_data = particle_df[particle_df['experiment'] == sample_exp]\n \n if len(sample_data) > 0:\n fig = plt.figure(figsize=(10, 8))\n ax = fig.add_subplot(111, projection='3d')\n \n for p_type, group in sample_data.groupby('particle_type'):\n ax.scatter(group['x'], group['y'], group['z'], \n label=p_type, \n color=particle_types[p_type]['color'],\n alpha=0.7)\n \n ax.set_xlabel('X coordinate')\n ax.set_ylabel('Y coordinate')\n ax.set_zlabel('Z coordinate')\n ax.set_title(f'3D Distribution of Particles in Experiment {sample_exp}')\n plt.legend()\n plt.tight_layout()\n plt.show()\n \n # Examine the coordinate ranges for each experiment\n coord_ranges = {}\n for exp in experiments:\n exp_data = particle_df[particle_df['experiment'] == exp]\n \n if not exp_data.empty:\n coord_ranges[exp] = {\n 'x': {'min': exp_data['x'].min(), 'max': exp_data['x'].max()},\n 'y': {'min': exp_data['y'].min(), 'max': exp_data['y'].max()},\n 'z': {'min': exp_data['z'].min(), 'max': exp_data['z'].max()}\n }\n \n print(\"\\nCoordinate ranges by experiment:\")\n for exp, ranges in coord_ranges.items():\n print(f\" {exp}:\")\n print(f\" X: {ranges['x']['min']:.1f} to {ranges['x']['max']:.1f}\")\n print(f\" Y: {ranges['y']['min']:.1f} to {ranges['y']['max']:.1f}\")\n print(f\" Z: {ranges['z']['min']:.1f} to {ranges['z']['max']:.1f}\")\nelse:\n print(\"No particle data found after extraction.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:53:20.185027Z","iopub.execute_input":"2025-04-03T07:53:20.185312Z","iopub.status.idle":"2025-04-03T07:53:21.287516Z","shell.execute_reply.started":"2025-04-03T07:53:20.185291Z","shell.execute_reply":"2025-04-03T07:53:21.286881Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Tomogram Visualization","metadata":{}},{"cell_type":"code","source":"import os\nimport zarr\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport json\nfrom tqdm.notebook import tqdm\n\n# Set paths\nbase_dir = '/kaggle/input/czii-cryo-et-object-identification'\ntrain_dir = os.path.join(base_dir, 'train')\ntest_dir = os.path.join(base_dir, 'test')\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:53:43.302227Z","iopub.execute_input":"2025-04-03T07:53:43.302519Z","iopub.status.idle":"2025-04-03T07:53:43.306991Z","shell.execute_reply.started":"2025-04-03T07:53:43.302498Z","shell.execute_reply":"2025-04-03T07:53:43.306109Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to load a tomogram using the correct zarr structure\ndef load_tomogram(zarr_path, resolution=0):\n \"\"\"\n Load a tomogram from a zarr file using the multiscale structure\n \n Parameters:\n zarr_path (str): Path to the zarr directory\n resolution (int): Resolution level (0 = highest, 1 = medium, 2 = lowest)\n \n Returns:\n numpy.ndarray: The loaded tomogram data\n \"\"\"\n try:\n # Open the zarr group\n z = zarr.open(zarr_path, mode='r')\n \n # Access the resolution level directly (based on the structure we observed)\n # Resolution level is a direct key in the group\n if str(resolution) in z:\n tomo_data = z[str(resolution)][:]\n print(f\"Loaded tomogram with shape {tomo_data.shape}\")\n return tomo_data\n else:\n print(f\"Resolution level {resolution} not found in zarr file\")\n return None\n except Exception as e:\n print(f\"Error loading tomogram: {str(e)}\")\n return None\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:53:49.488550Z","iopub.execute_input":"2025-04-03T07:53:49.488891Z","iopub.status.idle":"2025-04-03T07:53:49.493842Z","shell.execute_reply.started":"2025-04-03T07:53:49.488864Z","shell.execute_reply":"2025-04-03T07:53:49.492978Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to visualize tomogram slices\ndef visualize_tomogram(tomo_data, title=\"Tomogram Slices\", n_slices=3, figsize=(15, 5)):\n \"\"\"\n Visualize slices of a 3D tomogram\n \n Parameters:\n tomo_data (numpy.ndarray): 3D tomogram data\n title (str): Plot title\n n_slices (int): Number of slices to visualize\n figsize (tuple): Figure size\n \"\"\"\n if tomo_data is None:\n print(\"No tomogram data to visualize\")\n return\n \n # Get tomogram dimensions\n depth, height, width = tomo_data.shape\n print(f\"Tomogram dimensions: {depth} x {height} x {width}\")\n \n # Choose slice indices at different depths\n slice_indices = np.linspace(depth // 4, 3 * depth // 4, n_slices).astype(int)\n \n # Create figure\n fig, axes = plt.subplots(1, n_slices, figsize=figsize)\n if n_slices == 1:\n axes = [axes]\n \n # Plot each slice\n for i, slice_idx in enumerate(slice_indices):\n im = axes[i].imshow(tomo_data[slice_idx], cmap='gray')\n axes[i].set_title(f'Z-Slice {slice_idx}/{depth}')\n axes[i].axis('off')\n fig.colorbar(im, ax=axes[i], fraction=0.046, pad=0.04)\n \n plt.suptitle(title)\n plt.tight_layout()\n plt.show()\n \n # Show orthogonal views (XY, XZ, YZ planes)\n fig, axes = plt.subplots(1, 3, figsize=figsize)\n \n # XY plane (middle slice in Z)\n z_mid = depth // 2\n axes[0].imshow(tomo_data[z_mid], cmap='gray')\n axes[0].set_title(f'XY Plane (Z={z_mid})')\n axes[0].axis('off')\n \n # XZ plane (middle slice in Y)\n y_mid = height // 2\n axes[1].imshow(tomo_data[:, y_mid, :], cmap='gray')\n axes[1].set_title(f'XZ Plane (Y={y_mid})')\n axes[1].axis('off')\n \n # YZ plane (middle slice in X)\n x_mid = width // 2\n axes[2].imshow(tomo_data[:, :, x_mid], cmap='gray')\n axes[2].set_title(f'YZ Plane (X={x_mid})')\n axes[2].axis('off')\n \n plt.suptitle(f\"{title} - Orthogonal Views\")\n plt.tight_layout()\n plt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:53:54.590806Z","iopub.execute_input":"2025-04-03T07:53:54.591092Z","iopub.status.idle":"2025-04-03T07:53:54.598811Z","shell.execute_reply.started":"2025-04-03T07:53:54.591071Z","shell.execute_reply":"2025-04-03T07:53:54.597920Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to visualize tomogram with particle positions\ndef visualize_tomogram_with_particles(tomo_data, json_paths, title=\"Tomogram with Particles\", slice_idx=None):\n \"\"\"\n Visualize a tomogram slice with particle positions overlaid\n \n Parameters:\n tomo_data (numpy.ndarray): 3D tomogram data\n json_paths (list): List of paths to JSON files with particle coordinates\n title (str): Plot title\n slice_idx (int, optional): Specific slice to visualize. If None, middle slice is used.\n \"\"\"\n if tomo_data is None:\n print(\"No tomogram data to visualize\")\n return\n \n # Get tomogram dimensions\n depth, height, width = tomo_data.shape\n \n # Choose slice index if not specified\n if slice_idx is None:\n slice_idx = depth // 2\n \n # Load particle coordinates from all JSON files\n particles_by_type = {}\n \n for json_path in json_paths:\n particle_type = os.path.splitext(os.path.basename(json_path))[0]\n \n try:\n with open(json_path, 'r') as f:\n data = json.load(f)\n \n coords = []\n if 'points' in data:\n for point in data['points']:\n if 'location' in point:\n loc = point['location']\n coords.append((loc.get('x', 0), loc.get('y', 0), loc.get('z', 0)))\n \n particles_by_type[particle_type] = coords\n except Exception as e:\n print(f\"Error loading {json_path}: {str(e)}\")\n \n # Define colors for different particle types\n colors = {\n 'apo-ferritin': 'red',\n 'beta-amylase': 'gray',\n 'beta-galactosidase': 'blue',\n 'ribosome': 'green',\n 'thyroglobulin': 'purple',\n 'virus-like-particle': 'orange'\n }\n \n # Create figure\n plt.figure(figsize=(12, 10))\n \n # Show the tomogram slice\n plt.imshow(tomo_data[slice_idx], cmap='gray')\n \n # Overlay particle positions near the slice\n slice_range = 10 # Consider particles within ±10 slices\n \n for p_type, coords in particles_by_type.items():\n # Filter coordinates near the current slice\n slice_coords = []\n for x, y, z in coords:\n z_idx = int(z / 10) # Convert physical z to index (assuming 10 Å voxel spacing)\n if abs(z_idx - slice_idx) <= slice_range:\n slice_coords.append((x, y))\n \n if slice_coords:\n x_coords = [x for x, _ in slice_coords]\n y_coords = [y for _, y in slice_coords]\n \n plt.scatter(x_coords, y_coords, \n c=colors.get(p_type, 'white'), \n label=f'{p_type} ({len(slice_coords)})', \n alpha=0.7, s=30, edgecolors='white')\n \n plt.title(f'{title}\\nSlice {slice_idx}/{depth}')\n plt.legend(loc='upper right', bbox_to_anchor=(1.1, 1))\n plt.axis('off')\n plt.tight_layout()\n plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:54:04.584920Z","iopub.execute_input":"2025-04-03T07:54:04.585209Z","iopub.status.idle":"2025-04-03T07:54:04.594014Z","shell.execute_reply.started":"2025-04-03T07:54:04.585188Z","shell.execute_reply":"2025-04-03T07:54:04.593083Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# List all experiments in training data\ntrain_experiments = [os.path.basename(p) for p in glob.glob(os.path.join(train_dir, 'static/ExperimentRuns/*'))]\nprint(f\"Found {len(train_experiments)} experiments in training data: {train_experiments}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:54:12.817426Z","iopub.execute_input":"2025-04-03T07:54:12.817770Z","iopub.status.idle":"2025-04-03T07:54:12.825849Z","shell.execute_reply.started":"2025-04-03T07:54:12.817742Z","shell.execute_reply":"2025-04-03T07:54:12.825083Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Choose one experiment for visualization\nif train_experiments:\n sample_experiment = train_experiments[0]\n print(f\"\\nVisualizing data for experiment: {sample_experiment}\")\n \n # Path to the denoised tomogram\n denoised_path = os.path.join(train_dir, 'static/ExperimentRuns', sample_experiment, 'VoxelSpacing10.000/denoised.zarr')\n \n # Check if the path exists\n if os.path.exists(denoised_path):\n print(f\"Loading denoised tomogram from: {denoised_path}\")\n \n # Load the tomogram at different resolutions\n for resolution in range(3): # 0, 1, 2\n tomo_data = load_tomogram(denoised_path, resolution=resolution)\n \n if tomo_data is not None:\n # Visualize the tomogram\n visualize_tomogram(tomo_data, \n title=f\"Experiment {sample_experiment} - Denoised (Resolution {resolution})\",\n n_slices=3)\n \n # Show histogram of voxel values\n plt.figure(figsize=(10, 6))\n plt.hist(tomo_data.flatten(), bins=100, alpha=0.7, color='blue')\n plt.title(f\"Voxel Value Distribution - {sample_experiment} (Resolution {resolution})\")\n plt.xlabel(\"Voxel Value\")\n plt.ylabel(\"Frequency\")\n plt.grid(True, alpha=0.3)\n plt.show()\n \n # Get particle annotation files for this experiment\n particle_jsons = glob.glob(os.path.join(train_dir, 'overlay/ExperimentRuns', sample_experiment, 'Picks/*.json'))\n \n if particle_jsons:\n print(f\"Found {len(particle_jsons)} particle annotation files\")\n \n # Visualize tomogram with particles\n visualize_tomogram_with_particles(tomo_data, particle_jsons, \n title=f\"Experiment {sample_experiment} - With Particles\")\n else:\n print(\"No particle annotation files found for this experiment\")\n \n # Only visualize highest resolution (level 0)\n if resolution == 0:\n # Also compare with other tomogram types if available\n tomogram_types = ['ctfdeconvolved.zarr', 'isonetcorrected.zarr', 'wbp.zarr']\n \n for tomo_type in tomogram_types:\n tomo_path = os.path.join(train_dir, 'static/ExperimentRuns', sample_experiment, \n f'VoxelSpacing10.000/{tomo_type}')\n \n if os.path.exists(tomo_path):\n print(f\"\\nLoading {tomo_type} tomogram...\")\n other_tomo = load_tomogram(tomo_path, resolution=0)\n \n if other_tomo is not None:\n # Visualize middle slice\n mid_slice = other_tomo.shape[0] // 2\n plt.figure(figsize=(10, 8))\n plt.imshow(other_tomo[mid_slice], cmap='gray')\n plt.title(f\"{sample_experiment} - {tomo_type} (Slice {mid_slice})\")\n plt.axis('off')\n plt.colorbar(fraction=0.046, pad=0.04)\n plt.tight_layout()\n plt.show()\n \n # Break after visualizing level 0 (highest resolution)\n if resolution > 0:\n break\n else:\n print(f\"Tomogram path does not exist: {denoised_path}\")\nelse:\n print(\"No experiment directories found in training data\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:54:22.621875Z","iopub.execute_input":"2025-04-03T07:54:22.622167Z","iopub.status.idle":"2025-04-03T07:54:36.231453Z","shell.execute_reply.started":"2025-04-03T07:54:22.622147Z","shell.execute_reply":"2025-04-03T07:54:36.230562Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Also check a test experiment\ntest_experiments = [os.path.basename(p) for p in glob.glob(os.path.join(test_dir, 'static/ExperimentRuns/*'))]\nif test_experiments:\n test_experiment = test_experiments[0]\n print(f\"\\nVisualizing data for test experiment: {test_experiment}\")\n \n # Path to the denoised tomogram\n test_denoised_path = os.path.join(test_dir, 'static/ExperimentRuns', test_experiment, 'VoxelSpacing10.000/denoised.zarr')\n \n if os.path.exists(test_denoised_path):\n print(f\"Loading test tomogram from: {test_denoised_path}\")\n \n # Load and visualize test tomogram\n test_tomo = load_tomogram(test_denoised_path, resolution=0)\n \n if test_tomo is not None:\n visualize_tomogram(test_tomo, \n title=f\"Test Experiment {test_experiment} - Denoised\",\n n_slices=3)\n else:\n print(f\"Test tomogram path does not exist: {test_denoised_path}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:55:00.496173Z","iopub.execute_input":"2025-04-03T07:55:00.496468Z","iopub.status.idle":"2025-04-03T07:55:03.898885Z","shell.execute_reply.started":"2025-04-03T07:55:00.496447Z","shell.execute_reply":"2025-04-03T07:55:03.897825Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# F-beta Evaluation Metric","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom mpl_toolkits.mplot3d import Axes3D\nfrom scipy.spatial.distance import cdist\n\n# Particle types and their weights for scoring\nparticle_types = {\n 'apo-ferritin': {'difficulty': 'easy', 'weight': 1},\n 'beta-amylase': {'difficulty': 'impossible', 'weight': 0},\n 'beta-galactosidase': {'difficulty': 'hard', 'weight': 2},\n 'ribosome': {'difficulty': 'easy', 'weight': 1},\n 'thyroglobulin': {'difficulty': 'hard', 'weight': 2},\n 'virus-like-particle': {'difficulty': 'easy', 'weight': 1}\n}\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:55:14.143931Z","iopub.execute_input":"2025-04-03T07:55:14.144266Z","iopub.status.idle":"2025-04-03T07:55:14.149112Z","shell.execute_reply.started":"2025-04-03T07:55:14.144242Z","shell.execute_reply":"2025-04-03T07:55:14.148212Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Approximate particle radii in Angstroms\n# Values estimated from the literature\nparticle_radii = {\n 'apo-ferritin': 60,\n 'beta-amylase': 45,\n 'beta-galactosidase': 80,\n 'ribosome': 100,\n 'thyroglobulin': 85,\n 'virus-like-particle': 120\n}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:55:18.086074Z","iopub.execute_input":"2025-04-03T07:55:18.086362Z","iopub.status.idle":"2025-04-03T07:55:18.090223Z","shell.execute_reply.started":"2025-04-03T07:55:18.086340Z","shell.execute_reply":"2025-04-03T07:55:18.089329Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Define the F-beta metric calculation function\ndef calculate_fbeta(precision, recall, beta=4):\n \"\"\"\n Calculate F-beta score from precision and recall values.\n \n Parameters:\n precision (float): Precision value\n recall (float): Recall value\n beta (float): Beta value (defaults to 4 as per competition requirements)\n \n Returns:\n float: F-beta score\n \"\"\"\n if precision == 0 and recall == 0:\n return 0\n \n return (1 + beta**2) * (precision * recall) / ((beta**2 * precision) + recall)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:55:21.735572Z","iopub.execute_input":"2025-04-03T07:55:21.735992Z","iopub.status.idle":"2025-04-03T07:55:21.741795Z","shell.execute_reply.started":"2025-04-03T07:55:21.735957Z","shell.execute_reply":"2025-04-03T07:55:21.740998Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to determine if a predicted particle is a true positive\ndef is_true_positive(pred_coords, true_coords, particle_type, radius_factor=0.5):\n \"\"\"\n Determine if a predicted particle is a true positive.\n \n Parameters:\n pred_coords (tuple): Coordinates of the predicted particle (x, y, z)\n true_coords (list): List of coordinates of true particles\n particle_type (str): Type of particle\n radius_factor (float): Factor of particle radius for matching (0.5 means within half radius)\n \n Returns:\n bool: True if the predicted particle is a true positive, False otherwise\n int: Index of the matched true particle if found, -1 otherwise\n \"\"\"\n if not true_coords:\n return False, -1\n \n # Get particle radius\n particle_radius = particle_radii.get(particle_type, 60) # Default to 60 Angstroms if unknown\n \n # Calculate the distance threshold\n threshold = particle_radius * radius_factor\n \n # Convert to numpy arrays for vectorized operations\n pred_coords_array = np.array(pred_coords).reshape(1, 3)\n true_coords_array = np.array(true_coords)\n \n # Calculate distances to all true particles\n distances = cdist(pred_coords_array, true_coords_array)[0]\n \n # Find the minimum distance\n min_dist_idx = np.argmin(distances)\n min_dist = distances[min_dist_idx]\n \n # Check if the minimum distance is below the threshold\n if min_dist <= threshold:\n return True, min_dist_idx\n \n return False, -1\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:55:27.089214Z","iopub.execute_input":"2025-04-03T07:55:27.089504Z","iopub.status.idle":"2025-04-03T07:55:27.094604Z","shell.execute_reply.started":"2025-04-03T07:55:27.089484Z","shell.execute_reply":"2025-04-03T07:55:27.093778Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to evaluate predictions against ground truth\ndef evaluate_predictions(predictions, ground_truth, beta=4):\n \"\"\"\n Evaluate predictions against ground truth using the F-beta metric.\n \n Parameters:\n predictions (pd.DataFrame): DataFrame with columns 'experiment', 'particle_type', 'x', 'y', 'z'\n ground_truth (pd.DataFrame): DataFrame with columns 'experiment', 'particle_type', 'x', 'y', 'z'\n beta (float): Beta value for F-beta calculation (default: 4)\n \n Returns:\n dict: Dictionary with evaluation results\n \"\"\"\n results = {'overall': {}, 'by_type': {}, 'by_experiment': {}}\n \n # Initialize counters\n total_tp = 0\n total_fp = 0\n total_fn = 0\n weighted_tp = 0\n weighted_fp = 0\n weighted_fn = 0\n \n # Initialize counters for each particle type\n type_stats = {}\n for p_type in particle_types:\n if particle_types[p_type]['weight'] > 0: # Only consider scored particles\n type_stats[p_type] = {'tp': 0, 'fp': 0, 'fn': 0, 'weight': particle_types[p_type]['weight']}\n \n # Process each experiment\n experiments = ground_truth['experiment'].unique()\n for experiment in experiments:\n # Get predictions and ground truth for this experiment\n exp_pred = predictions[predictions['experiment'] == experiment]\n exp_true = ground_truth[ground_truth['experiment'] == experiment]\n \n # Process each particle type\n for p_type in particle_types:\n if particle_types[p_type]['weight'] == 0:\n continue # Skip particles with weight 0 (beta-amylase)\n \n # Get predictions and ground truth for this type\n type_pred = exp_pred[exp_pred['particle_type'] == p_type]\n type_true = exp_true[exp_true['particle_type'] == p_type]\n \n # Extract coordinates\n pred_coords = type_pred[['x', 'y', 'z']].values.tolist()\n true_coords = type_true[['x', 'y', 'z']].values.tolist()\n \n # Count true positives, false positives, and false negatives\n tp = 0\n fp = len(pred_coords) # Start assuming all predictions are false positives\n fn = len(true_coords) # Start assuming all true particles are false negatives\n \n # Track which true particles have been matched\n matched_true = [False] * len(true_coords)\n \n # Check each prediction\n for pred_coord in pred_coords:\n is_tp, match_idx = is_true_positive(pred_coord, true_coords, p_type)\n \n if is_tp and not matched_true[match_idx]:\n tp += 1\n fp -= 1 # One less false positive\n fn -= 1 # One less false negative\n matched_true[match_idx] = True\n \n # Update type statistics\n type_stats[p_type]['tp'] += tp\n type_stats[p_type]['fp'] += fp\n type_stats[p_type]['fn'] += fn\n \n # Update total counters\n total_tp += tp\n total_fp += fp\n total_fn += fn\n \n # Update weighted counters\n weight = particle_types[p_type]['weight']\n weighted_tp += tp * weight\n weighted_fp += fp * weight\n weighted_fn += fn * weight\n \n # Calculate overall micro-averaged precision, recall, and F-beta\n if weighted_tp + weighted_fp > 0:\n weighted_precision = weighted_tp / (weighted_tp + weighted_fp)\n else:\n weighted_precision = 0\n \n if weighted_tp + weighted_fn > 0:\n weighted_recall = weighted_tp / (weighted_tp + weighted_fn)\n else:\n weighted_recall = 0\n \n weighted_fbeta = calculate_fbeta(weighted_precision, weighted_recall, beta)\n \n # Store overall results\n results['overall'] = {\n 'true_positives': total_tp,\n 'false_positives': total_fp,\n 'false_negatives': total_fn,\n 'weighted_true_positives': weighted_tp,\n 'weighted_false_positives': weighted_fp,\n 'weighted_false_negatives': weighted_fn,\n 'precision': weighted_precision,\n 'recall': weighted_recall,\n 'f_beta': weighted_fbeta\n }\n \n # Calculate results for each particle type\n for p_type, stats in type_stats.items():\n tp = stats['tp']\n fp = stats['fp']\n fn = stats['fn']\n \n if tp + fp > 0:\n precision = tp / (tp + fp)\n else:\n precision = 0\n \n if tp + fn > 0:\n recall = tp / (tp + fn)\n else:\n recall = 0\n \n fbeta = calculate_fbeta(precision, recall, beta)\n \n results['by_type'][p_type] = {\n 'true_positives': tp,\n 'false_positives': fp,\n 'false_negatives': fn,\n 'precision': precision,\n 'recall': recall,\n 'f_beta': fbeta,\n 'weight': stats['weight']\n }\n \n return results","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:55:32.520452Z","iopub.execute_input":"2025-04-03T07:55:32.520806Z","iopub.status.idle":"2025-04-03T07:55:32.532897Z","shell.execute_reply.started":"2025-04-03T07:55:32.520776Z","shell.execute_reply":"2025-04-03T07:55:32.532055Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to visualize the F-beta metric\ndef visualize_fbeta():\n \"\"\"\n Create visualizations to understand the F-beta metric with beta=4.\n \"\"\"\n # Create a grid of precision and recall values\n precision_values = np.linspace(0.01, 1.0, 100)\n recall_values = np.linspace(0.01, 1.0, 100)\n P, R = np.meshgrid(precision_values, recall_values)\n \n # Calculate F-beta for each precision-recall pair\n F_beta = np.zeros_like(P)\n for i in range(P.shape[0]):\n for j in range(P.shape[1]):\n F_beta[i, j] = calculate_fbeta(P[i, j], R[i, j], beta=4)\n \n # 3D surface plot of F-beta\n fig = plt.figure(figsize=(12, 10))\n ax = fig.add_subplot(111, projection='3d')\n surf = ax.plot_surface(P, R, F_beta, cmap='viridis', alpha=0.8)\n ax.set_xlabel('Precision')\n ax.set_ylabel('Recall')\n ax.set_zlabel('F-beta (beta=4)')\n ax.set_title('F-beta Metric (beta=4) for Different Precision and Recall Values')\n fig.colorbar(surf, ax=ax, shrink=0.5, aspect=5)\n plt.tight_layout()\n plt.show()\n \n # Contour plot of F-beta\n plt.figure(figsize=(10, 8))\n contour = plt.contourf(P, R, F_beta, 20, cmap='viridis')\n plt.colorbar(contour, label='F-beta (beta=4)')\n plt.xlabel('Precision')\n plt.ylabel('Recall')\n plt.title('Contour Plot of F-beta Metric (beta=4)')\n plt.grid(True, alpha=0.3)\n \n # Add some contour lines with labels\n contour_lines = plt.contour(P, R, F_beta, levels=[0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9], \n colors='white', linestyles='dashed')\n plt.clabel(contour_lines, inline=True, fontsize=8, fmt='%.1f')\n plt.tight_layout()\n plt.show()\n \n # Plot F-beta for fixed precision or recall values\n plt.figure(figsize=(12, 6))\n \n # For fixed precision values\n plt.subplot(1, 2, 1)\n for precision in [0.2, 0.4, 0.6, 0.8, 1.0]:\n fbeta_values = [calculate_fbeta(precision, r, beta=4) for r in recall_values]\n plt.plot(recall_values, fbeta_values, label=f'Precision = {precision:.1f}')\n \n plt.xlabel('Recall')\n plt.ylabel('F-beta (beta=4)')\n plt.title('F-beta for Fixed Precision Values')\n plt.legend()\n plt.grid(True, alpha=0.3)\n \n # For fixed recall values\n plt.subplot(1, 2, 2)\n for recall in [0.2, 0.4, 0.6, 0.8, 1.0]:\n fbeta_values = [calculate_fbeta(p, recall, beta=4) for p in precision_values]\n plt.plot(precision_values, fbeta_values, label=f'Recall = {recall:.1f}')\n \n plt.xlabel('Precision')\n plt.ylabel('F-beta (beta=4)')\n plt.title('F-beta for Fixed Recall Values')\n plt.legend()\n plt.grid(True, alpha=0.3)\n \n plt.tight_layout()\n plt.show()\n \n # Compare different beta values\n plt.figure(figsize=(10, 6))\n \n # Fixed recall of 0.8\n recall = 0.8\n for beta in [0.5, 1, 2, 4, 8]:\n fbeta_values = [calculate_fbeta(p, recall, beta=beta) for p in precision_values]\n plt.plot(precision_values, fbeta_values, label=f'Beta = {beta}')\n \n plt.xlabel('Precision')\n plt.ylabel('F-beta')\n plt.title(f'F-beta Metrics for Different Beta Values (Recall = {recall})')\n plt.legend()\n plt.grid(True, alpha=0.3)\n plt.tight_layout()\n plt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:55:43.876133Z","iopub.execute_input":"2025-04-03T07:55:43.876417Z","iopub.status.idle":"2025-04-03T07:55:43.887570Z","shell.execute_reply.started":"2025-04-03T07:55:43.876396Z","shell.execute_reply":"2025-04-03T07:55:43.886724Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to create a sample prediction and ground truth for demonstration\ndef create_sample_data():\n \"\"\"\n Create sample prediction and ground truth data for demonstration.\n \"\"\"\n # Ground truth\n ground_truth_data = []\n \n # Experiment 1: 20 particles of each type (except beta-amylase)\n for _ in range(20):\n for p_type in ['apo-ferritin', 'beta-galactosidase', 'ribosome', 'thyroglobulin', 'virus-like-particle']:\n x = np.random.uniform(100, 1000)\n y = np.random.uniform(100, 1000)\n z = np.random.uniform(50, 150)\n ground_truth_data.append({'experiment': 'TS_5_4', 'particle_type': p_type, 'x': x, 'y': y, 'z': z})\n \n # Experiment 2: 15 particles of each type\n for _ in range(15):\n for p_type in ['apo-ferritin', 'beta-galactosidase', 'ribosome', 'thyroglobulin', 'virus-like-particle']:\n x = np.random.uniform(100, 1000)\n y = np.random.uniform(100, 1000)\n z = np.random.uniform(50, 150)\n ground_truth_data.append({'experiment': 'TS_6_4', 'particle_type': p_type, 'x': x, 'y': y, 'z': z})\n \n # Create ground truth DataFrame\n ground_truth_df = pd.DataFrame(ground_truth_data)\n \n # Create predictions with some noise and missing particles\n predictions_data = []\n \n # Copy 80% of ground truth with some noise\n for idx, row in ground_truth_df.iterrows():\n if np.random.random() < 0.8: # 80% chance of detecting the particle\n noise_x = np.random.normal(0, 10) # Add some noise\n noise_y = np.random.normal(0, 10)\n noise_z = np.random.normal(0, 5)\n predictions_data.append({\n 'experiment': row['experiment'],\n 'particle_type': row['particle_type'],\n 'x': row['x'] + noise_x,\n 'y': row['y'] + noise_y,\n 'z': row['z'] + noise_z\n })\n \n # Add some false positives (10% of the total)\n num_false_positives = int(0.1 * len(ground_truth_df))\n for _ in range(num_false_positives):\n experiment = np.random.choice(['TS_5_4', 'TS_6_4'])\n p_type = np.random.choice(['apo-ferritin', 'beta-galactosidase', 'ribosome', 'thyroglobulin', 'virus-like-particle'])\n x = np.random.uniform(100, 1000)\n y = np.random.uniform(100, 1000)\n z = np.random.uniform(50, 150)\n predictions_data.append({'experiment': experiment, 'particle_type': p_type, 'x': x, 'y': y, 'z': z})\n \n # Create predictions DataFrame\n predictions_df = pd.DataFrame(predictions_data)\n \n return predictions_df, ground_truth_df\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:55:52.897442Z","iopub.execute_input":"2025-04-03T07:55:52.897783Z","iopub.status.idle":"2025-04-03T07:55:52.905935Z","shell.execute_reply.started":"2025-04-03T07:55:52.897754Z","shell.execute_reply":"2025-04-03T07:55:52.905189Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Visualize the F-beta metric properties\nprint(\"Visualizing F-beta metric properties...\")\nvisualize_fbeta()\n\n# Demonstrate the evaluation with sample data\nprint(\"\\nDemonstrating the evaluation metric with sample data...\")\npredictions, ground_truth = create_sample_data()\n\nprint(f\"Ground truth shape: {ground_truth.shape}\")\nprint(f\"Predictions shape: {predictions.shape}\")\n\n# Evaluate the predictions\nresults = evaluate_predictions(predictions, ground_truth, beta=4)\n\n# Print overall results\nprint(\"\\nOverall Results:\")\nprint(f\"Precision: {results['overall']['precision']:.4f}\")\nprint(f\"Recall: {results['overall']['recall']:.4f}\")\nprint(f\"F-beta (beta=4): {results['overall']['f_beta']:.4f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:56:02.285476Z","iopub.execute_input":"2025-04-03T07:56:02.285841Z","iopub.status.idle":"2025-04-03T07:56:03.880977Z","shell.execute_reply.started":"2025-04-03T07:56:02.285811Z","shell.execute_reply":"2025-04-03T07:56:03.880067Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Print results by particle type\nprint(\"\\nResults by Particle Type:\")\nfor p_type, stats in results['by_type'].items():\n print(f\"{p_type}:\")\n print(f\" Precision: {stats['precision']:.4f}\")\n print(f\" Recall: {stats['recall']:.4f}\")\n print(f\" F-beta: {stats['f_beta']:.4f}\")\n print(f\" TP: {stats['true_positives']}, FP: {stats['false_positives']}, FN: {stats['false_negatives']}\")\n print(f\" Weight: {stats['weight']}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:56:16.016581Z","iopub.execute_input":"2025-04-03T07:56:16.016914Z","iopub.status.idle":"2025-04-03T07:56:16.027202Z","shell.execute_reply.started":"2025-04-03T07:56:16.016892Z","shell.execute_reply":"2025-04-03T07:56:16.026330Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Visualize results by particle type\nplt.figure(figsize=(12, 6))\nparticle_types_list = list(results['by_type'].keys())\nprecision_values = [results['by_type'][p]['precision'] for p in particle_types_list]\nrecall_values = [results['by_type'][p]['recall'] for p in particle_types_list]\nfbeta_values = [results['by_type'][p]['f_beta'] for p in particle_types_list]\n\nx = np.arange(len(particle_types_list))\nwidth = 0.25\n\nplt.bar(x - width, precision_values, width, label='Precision')\nplt.bar(x, recall_values, width, label='Recall')\nplt.bar(x + width, fbeta_values, width, label='F-beta')\n\nplt.xlabel('Particle Type')\nplt.ylabel('Score')\nplt.title('Precision, Recall, and F-beta by Particle Type')\nplt.xticks(x, particle_types_list, rotation=45, ha='right')\nplt.legend()\nplt.grid(True, alpha=0.3)\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:56:23.770383Z","iopub.execute_input":"2025-04-03T07:56:23.770915Z","iopub.status.idle":"2025-04-03T07:56:24.125993Z","shell.execute_reply.started":"2025-04-03T07:56:23.770876Z","shell.execute_reply":"2025-04-03T07:56:24.125191Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Data Preprocessing","metadata":{}},{"cell_type":"code","source":"import os\nimport glob\nimport zarr\nimport numpy as np\nimport pandas as pd\nimport json\nimport matplotlib.pyplot as plt\nfrom tqdm.notebook import tqdm\nimport gc\nfrom sklearn.model_selection import train_test_split\nfrom scipy.ndimage import zoom, gaussian_filter","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:56:49.315197Z","iopub.execute_input":"2025-04-03T07:56:49.315542Z","iopub.status.idle":"2025-04-03T07:56:49.500005Z","shell.execute_reply.started":"2025-04-03T07:56:49.315514Z","shell.execute_reply":"2025-04-03T07:56:49.499112Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Set paths\nbase_dir = '/kaggle/input/czii-cryo-et-object-identification'\ntrain_dir = os.path.join(base_dir, 'train')\ntest_dir = os.path.join(base_dir, 'test')\noutput_dir = '/kaggle/working/preprocessed'\n\n# Create output directory if it doesn't exist\nos.makedirs(output_dir, exist_ok=True)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:56:55.665494Z","iopub.execute_input":"2025-04-03T07:56:55.665848Z","iopub.status.idle":"2025-04-03T07:56:55.670161Z","shell.execute_reply.started":"2025-04-03T07:56:55.665812Z","shell.execute_reply":"2025-04-03T07:56:55.669296Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Particle types and their properties\nparticle_types = {\n 'apo-ferritin': {'difficulty': 'easy', 'weight': 1, 'radius': 60},\n 'beta-amylase': {'difficulty': 'impossible', 'weight': 0, 'radius': 45},\n 'beta-galactosidase': {'difficulty': 'hard', 'weight': 2, 'radius': 80},\n 'ribosome': {'difficulty': 'easy', 'weight': 1, 'radius': 100},\n 'thyroglobulin': {'difficulty': 'hard', 'weight': 2, 'radius': 85},\n 'virus-like-particle': {'difficulty': 'easy', 'weight': 1, 'radius': 120}\n}\n\n# Get scored particle types (weight > 0)\nscored_particle_types = [p for p, props in particle_types.items() if props['weight'] > 0]\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:57:01.009567Z","iopub.execute_input":"2025-04-03T07:57:01.009954Z","iopub.status.idle":"2025-04-03T07:57:01.014938Z","shell.execute_reply.started":"2025-04-03T07:57:01.009922Z","shell.execute_reply":"2025-04-03T07:57:01.013878Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to load a tomogram from a zarr file\ndef load_tomogram(zarr_path, resolution=0):\n \"\"\"\n Load a tomogram from a zarr file using the multiscale structure\n \n Parameters:\n zarr_path (str): Path to the zarr directory\n resolution (int): Resolution level (0 = highest, 1 = medium, 2 = lowest)\n \n Returns:\n numpy.ndarray: The loaded tomogram data\n \"\"\"\n try:\n # Open the zarr group\n z = zarr.open(zarr_path, mode='r')\n \n # Access the resolution level directly\n if str(resolution) in z:\n tomo_data = z[str(resolution)][:]\n return tomo_data\n else:\n print(f\"Resolution level {resolution} not found in zarr file\")\n return None\n except Exception as e:\n print(f\"Error loading tomogram: {str(e)}\")\n return None","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:57:04.977416Z","iopub.execute_input":"2025-04-03T07:57:04.977770Z","iopub.status.idle":"2025-04-03T07:57:04.982553Z","shell.execute_reply.started":"2025-04-03T07:57:04.977740Z","shell.execute_reply":"2025-04-03T07:57:04.981552Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to load particle coordinates from JSON\ndef load_particle_coords(json_path):\n \"\"\"\n Load particle coordinates from JSON file\n Uses 'points' key and 'location' field\n \n Parameters:\n json_path (str): Path to the JSON file\n \n Returns:\n list: List of (x, y, z) coordinate tuples\n \"\"\"\n try:\n with open(json_path, 'r') as f:\n data = json.load(f)\n \n # Extract coordinates from points\n coords = []\n if 'points' in data:\n for point in data['points']:\n if 'location' in point:\n loc = point['location']\n coords.append((loc.get('x', 0), loc.get('y', 0), loc.get('z', 0)))\n \n return coords\n except Exception as e:\n print(f\"Error loading {json_path}: {str(e)}\")\n return []","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:57:08.805750Z","iopub.execute_input":"2025-04-03T07:57:08.806036Z","iopub.status.idle":"2025-04-03T07:57:08.811210Z","shell.execute_reply.started":"2025-04-03T07:57:08.806015Z","shell.execute_reply":"2025-04-03T07:57:08.810106Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to preprocess a tomogram\ndef preprocess_tomogram(tomo_data):\n \"\"\"\n Preprocess a tomogram\n \n Parameters:\n tomo_data (numpy.ndarray): 3D tomogram data\n \n Returns:\n numpy.ndarray: Preprocessed tomogram data\n \"\"\"\n # Make a copy to avoid modifying the original\n processed = tomo_data.copy()\n \n # Apply Gaussian filtering to reduce noise\n processed = gaussian_filter(processed, sigma=1.0)\n \n # Normalize to [0, 1] range\n min_val = processed.min()\n max_val = processed.max()\n if max_val > min_val:\n processed = (processed - min_val) / (max_val - min_val)\n \n # Enhance contrast\n processed = np.clip((processed - 0.1) * 1.25, 0, 1)\n \n return processed","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:57:12.508324Z","iopub.execute_input":"2025-04-03T07:57:12.508660Z","iopub.status.idle":"2025-04-03T07:57:12.513416Z","shell.execute_reply.started":"2025-04-03T07:57:12.508633Z","shell.execute_reply":"2025-04-03T07:57:12.512535Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to create a density map for particles\ndef create_density_map(shape, coords, radius=10):\n \"\"\"\n Create a density map for particle locations\n \n Parameters:\n shape (tuple): Shape of the output density map (depth, height, width)\n coords (list): List of (x, y, z) coordinate tuples\n radius (float): Radius of particles in voxels\n \n Returns:\n numpy.ndarray: Density map with Gaussian-like peaks at particle locations\n \"\"\"\n # Initialize empty map\n density_map = np.zeros(shape, dtype=np.float32)\n \n # Convert coordinates to indices\n # Assuming 10 Angstrom voxel spacing\n voxel_spacing = 10.0\n \n depth, height, width = shape\n \n # Add Gaussian peaks for each particle\n for x, y, z in coords:\n # Convert physical coordinates to voxel indices\n z_idx, y_idx, x_idx = int(z / voxel_spacing), int(y / voxel_spacing), int(x / voxel_spacing)\n \n # Skip if outside the volume\n if not (0 <= z_idx < depth and 0 <= y_idx < height and 0 <= x_idx < width):\n continue\n \n # Create a spherical mask around the particle\n z_min = max(0, z_idx - radius)\n z_max = min(depth, z_idx + radius + 1)\n y_min = max(0, y_idx - radius)\n y_max = min(height, y_idx + radius + 1)\n x_min = max(0, x_idx - radius)\n x_max = min(width, x_idx + radius + 1)\n \n # Create coordinate grids\n z_grid, y_grid, x_grid = np.ogrid[z_min:z_max, y_min:y_max, x_min:x_max]\n \n # Calculate distance from center\n dist_from_center = np.sqrt(\n (z_grid - z_idx)**2 + \n (y_grid - y_idx)**2 + \n (x_grid - x_idx)**2\n )\n \n # Use a Gaussian-like function to create soft spheres\n mask = np.exp(-(dist_from_center**2) / (2 * (radius/2)**2))\n \n # Add to density map\n density_map[z_min:z_max, y_min:y_max, x_min:x_max] = np.maximum(\n density_map[z_min:z_max, y_min:y_max, x_min:x_max],\n mask\n )\n \n return density_map","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:57:17.337405Z","iopub.execute_input":"2025-04-03T07:57:17.337736Z","iopub.status.idle":"2025-04-03T07:57:17.344606Z","shell.execute_reply.started":"2025-04-03T07:57:17.337679Z","shell.execute_reply":"2025-04-03T07:57:17.343867Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to extract subvolumes (patches) from a tomogram\ndef extract_subvolumes(tomo_data, coords, patch_size=64, target_radius=6, max_patches_per_tomo=1000):\n \"\"\"\n Extract subvolumes (patches) from a tomogram\n \n Parameters:\n tomo_data (numpy.ndarray): 3D tomogram data\n coords (dict): Dictionary mapping particle types to coordinate lists\n patch_size (int): Size of cubic patches to extract\n target_radius (int): Radius of the target in the output density map (in voxels)\n max_patches_per_tomo (int): Maximum number of patches to extract per tomogram\n \n Returns:\n tuple: (patches, labels) where patches are subvolumes and labels are target density maps\n \"\"\"\n depth, height, width = tomo_data.shape\n half_size = patch_size // 2\n \n patches = []\n labels = []\n metadata = []\n \n # Combine all coordinates\n all_coords = []\n for p_type, coords_list in coords.items():\n for coord in coords_list:\n all_coords.append((coord[0], coord[1], coord[2], p_type))\n \n # Shuffle to ensure random selection if we hit max_patches_per_tomo\n np.random.shuffle(all_coords)\n \n # Extract patches around particle centers\n patch_count = 0\n for x, y, z, p_type in all_coords:\n if patch_count >= max_patches_per_tomo:\n break\n \n # Convert physical coordinates to voxel indices\n z_idx, y_idx, x_idx = int(z / 10.0), int(y / 10.0), int(x / 10.0)\n \n # Check if the patch is fully within the tomogram\n if (z_idx - half_size < 0 or z_idx + half_size >= depth or\n y_idx - half_size < 0 or y_idx + half_size >= height or\n x_idx - half_size < 0 or x_idx + half_size >= width):\n continue\n \n # Extract the patch\n patch = tomo_data[\n z_idx - half_size:z_idx + half_size,\n y_idx - half_size:y_idx + half_size,\n x_idx - half_size:x_idx + half_size\n ]\n \n # Skip if patch is invalid\n if patch.shape != (patch_size, patch_size, patch_size):\n continue\n \n # Create a label (target density map)\n if p_type in scored_particle_types:\n # Only create a target for scored particle types\n label = np.zeros((patch_size, patch_size, patch_size), dtype=np.float32)\n \n # Create a spherical mask\n z_grid, y_grid, x_grid = np.ogrid[\n :patch_size, \n :patch_size, \n :patch_size\n ]\n center = patch_size // 2\n \n dist_from_center = np.sqrt(\n (z_grid - center)**2 + \n (y_grid - center)**2 + \n (x_grid - center)**2\n )\n \n # Create Gaussian-like target\n label = np.exp(-(dist_from_center**2) / (2 * (target_radius/2)**2))\n else:\n # For non-scored particles, use empty labels\n label = np.zeros((patch_size, patch_size, patch_size), dtype=np.float32)\n \n patches.append(patch)\n labels.append(label)\n metadata.append({\n 'particle_type': p_type,\n 'x': x,\n 'y': y,\n 'z': z,\n 'weight': particle_types[p_type]['weight']\n })\n \n patch_count += 1\n \n # Also add some random negative patches (background)\n num_negative = min(len(patches) // 4, max_patches_per_tomo - patch_count)\n \n for _ in range(num_negative):\n # Generate random coordinates away from particles\n while True:\n z_idx = np.random.randint(half_size, depth - half_size)\n y_idx = np.random.randint(half_size, height - half_size)\n x_idx = np.random.randint(half_size, width - half_size)\n \n # Check if this point is far from all particles\n physical_x = x_idx * 10.0\n physical_y = y_idx * 10.0\n physical_z = z_idx * 10.0\n \n # Check distance from all particles\n min_dist = float('inf')\n for x, y, z, p_type in all_coords:\n dist = np.sqrt((x - physical_x)**2 + (y - physical_y)**2 + (z - physical_z)**2)\n min_dist = min(min_dist, dist)\n \n # If far enough from particles, use this location\n if min_dist > 100: # 100 Angstroms away from any particle\n break\n \n # Extract the patch\n patch = tomo_data[\n z_idx - half_size:z_idx + half_size,\n y_idx - half_size:y_idx + half_size,\n x_idx - half_size:x_idx + half_size\n ]\n \n # Skip if patch is invalid\n if patch.shape != (patch_size, patch_size, patch_size):\n continue\n \n # Use empty label for negative patches\n label = np.zeros((patch_size, patch_size, patch_size), dtype=np.float32)\n \n patches.append(patch)\n labels.append(label)\n metadata.append({\n 'particle_type': 'background',\n 'x': physical_x,\n 'y': physical_y,\n 'z': physical_z,\n 'weight': 0\n })\n \n return np.array(patches), np.array(labels), metadata","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:57:21.836724Z","iopub.execute_input":"2025-04-03T07:57:21.837038Z","iopub.status.idle":"2025-04-03T07:57:21.850052Z","shell.execute_reply.started":"2025-04-03T07:57:21.837015Z","shell.execute_reply":"2025-04-03T07:57:21.849042Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Process training data\ndef process_training_data(patch_size=64, split_ratio=0.2):\n \"\"\"\n Process all training data, extract patches, and save them\n \n Parameters:\n patch_size (int): Size of cubic patches to extract\n split_ratio (float): Ratio for validation split\n \n Returns:\n tuple: (train_patches, train_labels, val_patches, val_labels, metadata)\n \"\"\"\n all_patches = []\n all_labels = []\n all_metadata = []\n \n # Get list of training experiments\n train_experiments = [os.path.basename(p) for p in glob.glob(os.path.join(train_dir, 'static/ExperimentRuns/*'))]\n print(f\"Found {len(train_experiments)} training experiments: {train_experiments}\")\n \n # Process each experiment\n for experiment in tqdm(train_experiments, desc=\"Processing experiments\"):\n # Load tomogram\n zarr_path = os.path.join(train_dir, 'static/ExperimentRuns', experiment, 'VoxelSpacing10.000/denoised.zarr')\n \n if not os.path.exists(zarr_path):\n print(f\"Tomogram not found: {zarr_path}\")\n continue\n \n tomo_data = load_tomogram(zarr_path)\n \n if tomo_data is None:\n print(f\"Failed to load tomogram for experiment {experiment}\")\n continue\n \n # Preprocess tomogram\n tomo_data = preprocess_tomogram(tomo_data)\n \n # Load particle coordinates for each particle type\n coords = {}\n \n for p_type in particle_types.keys():\n json_path = os.path.join(train_dir, 'overlay/ExperimentRuns', experiment, 'Picks', f\"{p_type}.json\")\n \n if os.path.exists(json_path):\n coords_list = load_particle_coords(json_path)\n if coords_list:\n coords[p_type] = coords_list\n \n # Extract patches\n patches, labels, metadata = extract_subvolumes(tomo_data, coords, patch_size)\n \n # Add experiment to metadata\n for m in metadata:\n m['experiment'] = experiment\n \n # Collect patches, labels, and metadata\n all_patches.append(patches)\n all_labels.append(labels)\n all_metadata.extend(metadata)\n \n # Free memory\n del tomo_data, patches, labels\n gc.collect()\n \n # Combine all patches and labels\n all_patches = np.concatenate(all_patches, axis=0)\n all_labels = np.concatenate(all_labels, axis=0)\n \n # Create metadata DataFrame\n metadata_df = pd.DataFrame(all_metadata)\n \n # Save metadata\n os.makedirs(os.path.join(output_dir, 'metadata'), exist_ok=True)\n metadata_df.to_csv(os.path.join(output_dir, 'metadata', 'training_metadata.csv'), index=False)\n \n # Split into training and validation sets\n train_patches, val_patches, train_labels, val_labels = train_test_split(\n all_patches, all_labels, test_size=split_ratio, random_state=42\n )\n \n # Save patches and labels\n os.makedirs(os.path.join(output_dir, 'training'), exist_ok=True)\n np.save(os.path.join(output_dir, 'training', 'train_patches.npy'), train_patches)\n np.save(os.path.join(output_dir, 'training', 'train_labels.npy'), train_labels)\n np.save(os.path.join(output_dir, 'training', 'val_patches.npy'), val_patches)\n np.save(os.path.join(output_dir, 'training', 'val_labels.npy'), val_labels)\n \n # Print statistics\n print(f\"Processed {len(all_patches)} patches\")\n print(f\"Training set: {len(train_patches)} patches\")\n print(f\"Validation set: {len(val_patches)} patches\")\n \n # Count number of patches per particle type\n particle_counts = metadata_df['particle_type'].value_counts()\n print(\"\\nPatch distribution by particle type:\")\n for p_type, count in particle_counts.items():\n print(f\" - {p_type}: {count} patches\")\n \n # Plot a few random patches\n plot_random_patches(train_patches, train_labels, metadata_df, n_samples=5)\n \n return train_patches, train_labels, val_patches, val_labels, metadata_df\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:57:39.420942Z","iopub.execute_input":"2025-04-03T07:57:39.421252Z","iopub.status.idle":"2025-04-03T07:57:39.431160Z","shell.execute_reply.started":"2025-04-03T07:57:39.421228Z","shell.execute_reply":"2025-04-03T07:57:39.430408Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to process test data\ndef process_test_data(patch_size=64, stride=32):\n \"\"\"\n Process test data by extracting overlapping patches\n \n Parameters:\n patch_size (int): Size of cubic patches to extract\n stride (int): Stride between patches\n \n Returns:\n dict: Dictionary of test data by experiment\n \"\"\"\n test_data = {}\n \n # Get list of test experiments\n test_experiments = [os.path.basename(p) for p in glob.glob(os.path.join(test_dir, 'static/ExperimentRuns/*'))]\n print(f\"Found {len(test_experiments)} test experiments: {test_experiments}\")\n \n # Process each experiment\n for experiment in tqdm(test_experiments, desc=\"Processing test experiments\"):\n # Load tomogram\n zarr_path = os.path.join(test_dir, 'static/ExperimentRuns', experiment, 'VoxelSpacing10.000/denoised.zarr')\n \n if not os.path.exists(zarr_path):\n print(f\"Test tomogram not found: {zarr_path}\")\n continue\n \n tomo_data = load_tomogram(zarr_path)\n \n if tomo_data is None:\n print(f\"Failed to load test tomogram for experiment {experiment}\")\n continue\n \n # Preprocess tomogram\n tomo_data = preprocess_tomogram(tomo_data)\n \n # Extract overlapping patches\n depth, height, width = tomo_data.shape\n half_size = patch_size // 2\n \n # Initialize data structures\n patches = []\n coordinates = []\n \n # Extract patches with overlap (stride)\n for z in range(half_size, depth - half_size, stride):\n for y in range(half_size, height - half_size, stride):\n for x in range(half_size, width - half_size, stride):\n # Extract the patch\n patch = tomo_data[\n z - half_size:z + half_size,\n y - half_size:y + half_size,\n x - half_size:x + half_size\n ]\n \n # Skip if patch is invalid\n if patch.shape != (patch_size, patch_size, patch_size):\n continue\n \n patches.append(patch)\n coordinates.append((x, y, z))\n \n # Convert to numpy arrays\n patches = np.array(patches)\n \n # Store test data\n test_data[experiment] = {\n 'patches': patches,\n 'coordinates': coordinates,\n 'shape': tomo_data.shape\n }\n \n # Save the test data\n os.makedirs(os.path.join(output_dir, 'test', experiment), exist_ok=True)\n np.save(os.path.join(output_dir, 'test', experiment, 'patches.npy'), patches)\n np.save(os.path.join(output_dir, 'test', experiment, 'coordinates.npy'), coordinates)\n np.save(os.path.join(output_dir, 'test', experiment, 'shape.npy'), tomo_data.shape)\n \n # Print statistics\n print(f\"Processed {len(patches)} patches for experiment {experiment}\")\n \n # Free memory\n del tomo_data, patches\n gc.collect()\n \n return test_data\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:57:47.903636Z","iopub.execute_input":"2025-04-03T07:57:47.903979Z","iopub.status.idle":"2025-04-03T07:57:47.912453Z","shell.execute_reply.started":"2025-04-03T07:57:47.903952Z","shell.execute_reply":"2025-04-03T07:57:47.911512Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to plot random patches\ndef plot_random_patches(patches, labels, metadata_df, n_samples=5):\n \"\"\"\n Plot random patches and their labels\n \n Parameters:\n patches (numpy.ndarray): Array of patches\n labels (numpy.ndarray): Array of labels\n metadata_df (pandas.DataFrame): DataFrame with metadata\n n_samples (int): Number of samples to plot\n \"\"\"\n # Get random indices\n indices = np.random.choice(len(patches), size=n_samples, replace=False)\n \n # Plot each sample\n for i, idx in enumerate(indices):\n patch = patches[idx]\n label = labels[idx]\n \n # Get metadata for this patch\n if i < len(metadata_df):\n p_type = metadata_df.iloc[idx]['particle_type']\n p_weight = metadata_df.iloc[idx]['weight']\n else:\n p_type = \"Unknown\"\n p_weight = \"Unknown\"\n \n # Plot the middle slice of the patch and label\n middle_slice = patch.shape[0] // 2\n \n plt.figure(figsize=(12, 6))\n \n # Plot patch\n plt.subplot(1, 2, 1)\n plt.imshow(patch[middle_slice], cmap='gray')\n plt.title(f\"Patch {idx}: {p_type} (Weight: {p_weight})\")\n plt.axis('off')\n \n # Plot label\n plt.subplot(1, 2, 2)\n plt.imshow(label[middle_slice], cmap='hot')\n plt.title(f\"Label {idx}\")\n plt.axis('off')\n \n plt.tight_layout()\n plt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:57:56.129841Z","iopub.execute_input":"2025-04-03T07:57:56.130140Z","iopub.status.idle":"2025-04-03T07:57:56.136363Z","shell.execute_reply.started":"2025-04-03T07:57:56.130119Z","shell.execute_reply":"2025-04-03T07:57:56.135616Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Main execution\nprint(\"Starting data preprocessing...\")\n\n# Set parameters\npatch_size = 64 # Size of cubic patches\nsplit_ratio = 0.2 # Ratio for validation split\n\n# Process training data\ntrain_patches, train_labels, val_patches, val_labels, metadata_df = process_training_data(patch_size, split_ratio)\n\n# Process test data\ntest_data = process_test_data(patch_size, stride=32)\n\nprint(\"Data preprocessing complete.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:58:02.829789Z","iopub.execute_input":"2025-04-03T07:58:02.830084Z","iopub.status.idle":"2025-04-03T07:59:01.410815Z","shell.execute_reply.started":"2025-04-03T07:58:02.830063Z","shell.execute_reply":"2025-04-03T07:59:01.409831Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# 3D Particle Detection Model","metadata":{}},{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nfrom tqdm.notebook import tqdm\nimport gc\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom torch.utils.data import Dataset, DataLoader","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T07:59:55.938279Z","iopub.execute_input":"2025-04-03T07:59:55.938594Z","iopub.status.idle":"2025-04-03T07:59:58.900617Z","shell.execute_reply.started":"2025-04-03T07:59:55.938569Z","shell.execute_reply":"2025-04-03T07:59:58.899770Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Set paths\noutput_dir = '/kaggle/working/preprocessed'\nmodel_dir = '/kaggle/working/models'\n\n# Create models directory if it doesn't exist\nos.makedirs(model_dir, exist_ok=True)\n\n# Set device\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Using device: {device}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:00:01.912538Z","iopub.execute_input":"2025-04-03T08:00:01.913285Z","iopub.status.idle":"2025-04-03T08:00:01.963507Z","shell.execute_reply.started":"2025-04-03T08:00:01.913254Z","shell.execute_reply":"2025-04-03T08:00:01.962777Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Set random seeds for reproducibility\nnp.random.seed(42)\ntorch.manual_seed(42)\nif torch.cuda.is_available():\n torch.cuda.manual_seed(42)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:00:06.634563Z","iopub.execute_input":"2025-04-03T08:00:06.634881Z","iopub.status.idle":"2025-04-03T08:00:06.642799Z","shell.execute_reply.started":"2025-04-03T08:00:06.634857Z","shell.execute_reply":"2025-04-03T08:00:06.642068Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Particle types and their properties\nparticle_types = {\n 'apo-ferritin': {'difficulty': 'easy', 'weight': 1, 'radius': 60},\n 'beta-amylase': {'difficulty': 'impossible', 'weight': 0, 'radius': 45},\n 'beta-galactosidase': {'difficulty': 'hard', 'weight': 2, 'radius': 80},\n 'ribosome': {'difficulty': 'easy', 'weight': 1, 'radius': 100},\n 'thyroglobulin': {'difficulty': 'hard', 'weight': 2, 'radius': 85},\n 'virus-like-particle': {'difficulty': 'easy', 'weight': 1, 'radius': 120}\n}\n\n# Get scored particle types (weight > 0)\nscored_particle_types = [p for p, props in particle_types.items() if props['weight'] > 0]\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:00:09.882437Z","iopub.execute_input":"2025-04-03T08:00:09.882770Z","iopub.status.idle":"2025-04-03T08:00:09.887666Z","shell.execute_reply.started":"2025-04-03T08:00:09.882744Z","shell.execute_reply":"2025-04-03T08:00:09.886786Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 3D U-Net model\nclass UNet3D(nn.Module):\n def __init__(self, in_channels=1, out_channels=1, init_features=16):\n super(UNet3D, self).__init__()\n \n features = init_features\n self.encoder1 = self._block(in_channels, features, name=\"enc1\")\n self.pool1 = nn.MaxPool3d(kernel_size=2, stride=2)\n \n self.encoder2 = self._block(features, features * 2, name=\"enc2\")\n self.pool2 = nn.MaxPool3d(kernel_size=2, stride=2)\n \n self.encoder3 = self._block(features * 2, features * 4, name=\"enc3\")\n self.pool3 = nn.MaxPool3d(kernel_size=2, stride=2)\n \n self.bottleneck = self._block(features * 4, features * 8, name=\"bottleneck\")\n \n self.upconv3 = nn.ConvTranspose3d(features * 8, features * 4, kernel_size=2, stride=2)\n self.decoder3 = self._block((features * 4) * 2, features * 4, name=\"dec3\")\n \n self.upconv2 = nn.ConvTranspose3d(features * 4, features * 2, kernel_size=2, stride=2)\n self.decoder2 = self._block((features * 2) * 2, features * 2, name=\"dec2\")\n \n self.upconv1 = nn.ConvTranspose3d(features * 2, features, kernel_size=2, stride=2)\n self.decoder1 = self._block(features * 2, features, name=\"dec1\")\n \n self.conv = nn.Conv3d(in_channels=features, out_channels=out_channels, kernel_size=1)\n \n def forward(self, x):\n enc1 = self.encoder1(x)\n enc2 = self.encoder2(self.pool1(enc1))\n enc3 = self.encoder3(self.pool2(enc2))\n \n bottleneck = self.bottleneck(self.pool3(enc3))\n \n dec3 = self.upconv3(bottleneck)\n dec3 = torch.cat((dec3, enc3), dim=1)\n dec3 = self.decoder3(dec3)\n \n dec2 = self.upconv2(dec3)\n dec2 = torch.cat((dec2, enc2), dim=1)\n dec2 = self.decoder2(dec2)\n \n dec1 = self.upconv1(dec2)\n dec1 = torch.cat((dec1, enc1), dim=1)\n dec1 = self.decoder1(dec1)\n \n return torch.sigmoid(self.conv(dec1))\n \n @staticmethod\n def _block(in_channels, features, name):\n return nn.Sequential(\n nn.Conv3d(in_channels, features, kernel_size=3, padding=1, bias=False),\n nn.BatchNorm3d(features),\n nn.ReLU(inplace=True),\n nn.Conv3d(features, features, kernel_size=3, padding=1, bias=False),\n nn.BatchNorm3d(features),\n nn.ReLU(inplace=True)\n )\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:00:12.987191Z","iopub.execute_input":"2025-04-03T08:00:12.987510Z","iopub.status.idle":"2025-04-03T08:00:12.996967Z","shell.execute_reply.started":"2025-04-03T08:00:12.987484Z","shell.execute_reply":"2025-04-03T08:00:12.996186Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Dataset class for 3D patches\nclass PatchDataset(Dataset):\n def __init__(self, patches, labels):\n self.patches = patches\n self.labels = labels\n \n def __len__(self):\n return len(self.patches)\n \n def __getitem__(self, idx):\n # Add channel dimension and convert to torch tensors\n patch = torch.FloatTensor(self.patches[idx]).unsqueeze(0)\n label = torch.FloatTensor(self.labels[idx]).unsqueeze(0)\n \n return patch, label","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:00:17.962087Z","iopub.execute_input":"2025-04-03T08:00:17.962383Z","iopub.status.idle":"2025-04-03T08:00:17.967189Z","shell.execute_reply.started":"2025-04-03T08:00:17.962363Z","shell.execute_reply":"2025-04-03T08:00:17.966168Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Dice loss for 3D segmentation\nclass DiceLoss(nn.Module):\n def __init__(self, smooth=1.0):\n super(DiceLoss, self).__init__()\n self.smooth = smooth\n \n def forward(self, predictions, targets):\n # Flatten the predictions and targets\n predictions = predictions.view(-1)\n targets = targets.view(-1)\n \n intersection = (predictions * targets).sum()\n dice = (2. * intersection + self.smooth) / (predictions.sum() + targets.sum() + self.smooth)\n \n return 1 - dice","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:00:21.151398Z","iopub.execute_input":"2025-04-03T08:00:21.151784Z","iopub.status.idle":"2025-04-03T08:00:21.156857Z","shell.execute_reply.started":"2025-04-03T08:00:21.151752Z","shell.execute_reply":"2025-04-03T08:00:21.155882Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Focal loss for handling class imbalance\nclass FocalLoss(nn.Module):\n def __init__(self, alpha=0.25, gamma=2.0, reduction='mean'):\n super(FocalLoss, self).__init__()\n self.alpha = alpha\n self.gamma = gamma\n self.reduction = reduction\n \n def forward(self, inputs, targets):\n BCE_loss = F.binary_cross_entropy(inputs, targets, reduction='none')\n pt = torch.exp(-BCE_loss)\n F_loss = self.alpha * (1-pt)**self.gamma * BCE_loss\n \n if self.reduction == 'mean':\n return torch.mean(F_loss)\n elif self.reduction == 'sum':\n return torch.sum(F_loss)\n else:\n return F_loss","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:00:26.480619Z","iopub.execute_input":"2025-04-03T08:00:26.481046Z","iopub.status.idle":"2025-04-03T08:00:26.486309Z","shell.execute_reply.started":"2025-04-03T08:00:26.481005Z","shell.execute_reply":"2025-04-03T08:00:26.485467Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Combined loss function\nclass CombinedLoss(nn.Module):\n def __init__(self, dice_weight=0.5, focal_weight=0.5):\n super(CombinedLoss, self).__init__()\n self.dice_weight = dice_weight\n self.focal_weight = focal_weight\n self.dice_loss = DiceLoss()\n self.focal_loss = FocalLoss()\n \n def forward(self, predictions, targets):\n dice = self.dice_loss(predictions, targets)\n focal = self.focal_loss(predictions, targets)\n \n return self.dice_weight * dice + self.focal_weight * focal\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:00:31.880165Z","iopub.execute_input":"2025-04-03T08:00:31.880462Z","iopub.status.idle":"2025-04-03T08:00:31.885518Z","shell.execute_reply.started":"2025-04-03T08:00:31.880440Z","shell.execute_reply":"2025-04-03T08:00:31.884573Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Training function\ndef train_model(model, train_loader, val_loader, criterion, optimizer, scheduler, num_epochs=50, patience=10):\n \"\"\"\n Train the 3D detection model\n \n Parameters:\n model (nn.Module): Model to train\n train_loader (DataLoader): Training data loader\n val_loader (DataLoader): Validation data loader\n criterion (nn.Module): Loss function\n optimizer (optim.Optimizer): Optimizer\n scheduler (optim.lr_scheduler): Learning rate scheduler\n num_epochs (int): Maximum number of epochs\n patience (int): Early stopping patience\n \n Returns:\n model: Trained model\n history: Training history\n \"\"\"\n # Initialize variables\n best_val_loss = float('inf')\n best_model_state = None\n patience_counter = 0\n history = {'train_loss': [], 'val_loss': []}\n \n # Move model to device\n model = model.to(device)\n \n # Training loop\n for epoch in range(num_epochs):\n model.train()\n train_loss = 0.0\n \n # Training step\n progress_bar = tqdm(train_loader, desc=f\"Epoch {epoch+1}/{num_epochs}\")\n for patches, labels in progress_bar:\n # Move data to device\n patches = patches.to(device)\n labels = labels.to(device)\n \n # Zero the gradients\n optimizer.zero_grad()\n \n # Forward pass\n outputs = model(patches)\n \n # Calculate loss\n loss = criterion(outputs, labels)\n \n # Backward pass and optimize\n loss.backward()\n optimizer.step()\n \n # Update training loss\n train_loss += loss.item() * patches.size(0)\n \n # Update progress bar\n progress_bar.set_postfix(loss=loss.item())\n \n # Calculate average training loss\n train_loss /= len(train_loader.dataset)\n \n # Validation step\n model.eval()\n val_loss = 0.0\n \n with torch.no_grad():\n for patches, labels in val_loader:\n # Move data to device\n patches = patches.to(device)\n labels = labels.to(device)\n \n # Forward pass\n outputs = model(patches)\n \n # Calculate loss\n loss = criterion(outputs, labels)\n \n # Update validation loss\n val_loss += loss.item() * patches.size(0)\n \n # Calculate average validation loss\n val_loss /= len(val_loader.dataset)\n \n # Update learning rate\n scheduler.step(val_loss)\n \n # Print progress\n print(f\"Epoch {epoch+1}/{num_epochs}, Train Loss: {train_loss:.4f}, Val Loss: {val_loss:.4f}\")\n \n # Update history\n history['train_loss'].append(train_loss)\n history['val_loss'].append(val_loss)\n \n # Check for improvement\n if val_loss < best_val_loss:\n best_val_loss = val_loss\n best_model_state = model.state_dict()\n patience_counter = 0\n else:\n patience_counter += 1\n \n # Early stopping\n if patience_counter >= patience:\n print(f\"Early stopping at epoch {epoch+1}\")\n break\n \n # Free up memory\n gc.collect()\n if torch.cuda.is_available():\n torch.cuda.empty_cache()\n \n # Load best model\n if best_model_state is not None:\n model.load_state_dict(best_model_state)\n \n return model, history\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:00:36.927209Z","iopub.execute_input":"2025-04-03T08:00:36.927497Z","iopub.status.idle":"2025-04-03T08:00:36.936064Z","shell.execute_reply.started":"2025-04-03T08:00:36.927477Z","shell.execute_reply":"2025-04-03T08:00:36.934996Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to visualize training history\ndef plot_training_history(history):\n \"\"\"\n Plot training history\n \n Parameters:\n history (dict): Training history\n \"\"\"\n plt.figure(figsize=(10, 6))\n plt.plot(history['train_loss'], label='Train Loss')\n plt.plot(history['val_loss'], label='Val Loss')\n plt.xlabel('Epoch')\n plt.ylabel('Loss')\n plt.title('Training History')\n plt.legend()\n plt.grid(True)\n plt.tight_layout()\n plt.savefig(os.path.join(model_dir, 'training_history.png'))\n plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:00:43.286869Z","iopub.execute_input":"2025-04-03T08:00:43.287147Z","iopub.status.idle":"2025-04-03T08:00:43.291618Z","shell.execute_reply.started":"2025-04-03T08:00:43.287127Z","shell.execute_reply":"2025-04-03T08:00:43.290836Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to visualize model predictions\ndef visualize_predictions(model, val_loader, num_samples=5):\n \"\"\"\n Visualize model predictions\n \n Parameters:\n model (nn.Module): Trained model\n val_loader (DataLoader): Validation data loader\n num_samples (int): Number of samples to visualize\n \"\"\"\n model.eval()\n \n # Get random samples\n samples = []\n with torch.no_grad(): # This ensures no gradients are tracked\n for patches, labels in val_loader:\n samples.append((patches, labels))\n if len(samples) >= num_samples:\n break\n \n # Visualize each sample\n for i, (patches, labels) in enumerate(samples):\n # Move data to device\n patches = patches.to(device)\n \n # Forward pass\n with torch.no_grad(): # Add this to be extra safe\n outputs = model(patches)\n \n # Move back to CPU and convert to numpy\n patches = patches.cpu().numpy()\n labels = labels.cpu().numpy()\n outputs = outputs.detach().cpu().numpy() # Use detach() before converting to numpy\n \n # Plot middle slice of first batch item\n middle_slice = patches.shape[2] // 2\n \n plt.figure(figsize=(15, 5))\n \n # Plot patch\n plt.subplot(1, 3, 1)\n plt.imshow(patches[0, 0, middle_slice], cmap='gray')\n plt.title(f\"Input Patch {i+1}\")\n plt.axis('off')\n \n # Plot ground truth\n plt.subplot(1, 3, 2)\n plt.imshow(labels[0, 0, middle_slice], cmap='hot')\n plt.title(\"Ground Truth\")\n plt.axis('off')\n \n # Plot prediction\n plt.subplot(1, 3, 3)\n plt.imshow(outputs[0, 0, middle_slice], cmap='hot')\n plt.title(\"Prediction\")\n plt.axis('off')\n \n plt.savefig(os.path.join(model_dir, f'prediction_{i+1}.png'))\n plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:00:47.982467Z","iopub.execute_input":"2025-04-03T08:00:47.982818Z","iopub.status.idle":"2025-04-03T08:00:47.989497Z","shell.execute_reply.started":"2025-04-03T08:00:47.982788Z","shell.execute_reply":"2025-04-03T08:00:47.988798Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Main execution\nprint(\"Starting model training...\")\n\n# Load training and validation data\ntrain_patches = np.load(os.path.join(output_dir, 'training', 'train_patches.npy'))\ntrain_labels = np.load(os.path.join(output_dir, 'training', 'train_labels.npy'))\nval_patches = np.load(os.path.join(output_dir, 'training', 'val_patches.npy'))\nval_labels = np.load(os.path.join(output_dir, 'training', 'val_labels.npy'))\n\nprint(f\"Training data shape: {train_patches.shape}\")\nprint(f\"Validation data shape: {val_patches.shape}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:00:54.223198Z","iopub.execute_input":"2025-04-03T08:00:54.223481Z","iopub.status.idle":"2025-04-03T08:00:58.796856Z","shell.execute_reply.started":"2025-04-03T08:00:54.223459Z","shell.execute_reply":"2025-04-03T08:00:58.796026Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Create datasets and dataloaders\ntrain_dataset = PatchDataset(train_patches, train_labels)\nval_dataset = PatchDataset(val_patches, val_labels)\n\nbatch_size = 16 # Adjust based on available memory\ntrain_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)\nval_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False, num_workers=4)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:01:07.145742Z","iopub.execute_input":"2025-04-03T08:01:07.146069Z","iopub.status.idle":"2025-04-03T08:01:07.150819Z","shell.execute_reply.started":"2025-04-03T08:01:07.146046Z","shell.execute_reply":"2025-04-03T08:01:07.149996Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Initialize model\nmodel = UNet3D(in_channels=1, out_channels=1, init_features=16)\n\n# Initialize loss function and optimizer\ncriterion = CombinedLoss(dice_weight=0.5, focal_weight=0.5)\noptimizer = optim.Adam(model.parameters(), lr=1e-4, weight_decay=1e-5)\nscheduler = optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=5, verbose=True)\n\n# Train model\nmodel, history = train_model(\n model, \n train_loader, \n val_loader, \n criterion, \n optimizer, \n scheduler,\n num_epochs=50,\n patience=10\n)\n\n# Save model\ntorch.save(model.state_dict(), os.path.join(model_dir, 'model.pth'))\n\n# Visualize training history\nplot_training_history(history)\n\n# Visualize predictions\n#visualize_predictions(model, val_loader, num_samples=5)\n\nprint(\"Model training complete. Saved model to\", os.path.join(model_dir, 'model.pth'))\n\n# Free up memory\n#del train_patches, train_labels, val_patches, val_labels\n#del train_dataset, val_dataset, train_loader, val_loader\n#gc.collect()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:01:11.771504Z","iopub.execute_input":"2025-04-03T08:01:11.771993Z","iopub.status.idle":"2025-04-03T08:30:28.508683Z","shell.execute_reply.started":"2025-04-03T08:01:11.771956Z","shell.execute_reply":"2025-04-03T08:30:28.507573Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"visualize_predictions(model, val_loader, num_samples=5)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:30:36.405754Z","iopub.execute_input":"2025-04-03T08:30:36.406330Z","iopub.status.idle":"2025-04-03T08:30:39.370546Z","shell.execute_reply.started":"2025-04-03T08:30:36.406303Z","shell.execute_reply":"2025-04-03T08:30:39.369495Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Particle Overlay Visualization","metadata":{}},{"cell_type":"markdown","source":"# Validation Set","metadata":{}},{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport matplotlib.patches as patches\nfrom matplotlib.colors import LinearSegmentedColormap\nimport zarr\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nimport gc\nfrom tqdm.notebook import tqdm","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:30:52.815813Z","iopub.execute_input":"2025-04-03T08:30:52.816153Z","iopub.status.idle":"2025-04-03T08:30:52.820793Z","shell.execute_reply.started":"2025-04-03T08:30:52.816130Z","shell.execute_reply":"2025-04-03T08:30:52.819868Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Set paths\nbase_dir = '/kaggle/input/czii-cryo-et-object-identification'\ntrain_dir = os.path.join(base_dir, 'train')\noutput_dir = '/kaggle/working/preprocessed'\nmodel_dir = '/kaggle/working/models'\nvisualization_dir = '/kaggle/working/val_visualizations'\n\n# Create visualization directory if it doesn't exist\nos.makedirs(visualization_dir, exist_ok=True)\n\n# Set device\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Using device: {device}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:30:57.587972Z","iopub.execute_input":"2025-04-03T08:30:57.588276Z","iopub.status.idle":"2025-04-03T08:30:57.594241Z","shell.execute_reply.started":"2025-04-03T08:30:57.588255Z","shell.execute_reply":"2025-04-03T08:30:57.593294Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Particle types and their properties\nparticle_types = {\n 'apo-ferritin': {'difficulty': 'easy', 'weight': 1, 'radius': 60, 'color': 'red'},\n 'beta-amylase': {'difficulty': 'impossible', 'weight': 0, 'radius': 45, 'color': 'yellow'},\n 'beta-galactosidase': {'difficulty': 'hard', 'weight': 2, 'radius': 80, 'color': 'green'},\n 'ribosome': {'difficulty': 'easy', 'weight': 1, 'radius': 100, 'color': 'blue'},\n 'thyroglobulin': {'difficulty': 'hard', 'weight': 2, 'radius': 85, 'color': 'purple'},\n 'virus-like-particle': {'difficulty': 'easy', 'weight': 1, 'radius': 120, 'color': 'orange'}\n}\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:31:01.676089Z","iopub.execute_input":"2025-04-03T08:31:01.676396Z","iopub.status.idle":"2025-04-03T08:31:01.681287Z","shell.execute_reply.started":"2025-04-03T08:31:01.676373Z","shell.execute_reply":"2025-04-03T08:31:01.680260Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class UNet3D(nn.Module):\n def __init__(self, in_channels=1, out_channels=1, init_features=16):\n super(UNet3D, self).__init__()\n \n features = init_features\n self.encoder1 = self._block(in_channels, features, name=\"enc1\")\n self.pool1 = nn.MaxPool3d(kernel_size=2, stride=2)\n \n self.encoder2 = self._block(features, features * 2, name=\"enc2\")\n self.pool2 = nn.MaxPool3d(kernel_size=2, stride=2)\n \n self.encoder3 = self._block(features * 2, features * 4, name=\"enc3\")\n self.pool3 = nn.MaxPool3d(kernel_size=2, stride=2)\n \n self.bottleneck = self._block(features * 4, features * 8, name=\"bottleneck\")\n \n self.upconv3 = nn.ConvTranspose3d(features * 8, features * 4, kernel_size=2, stride=2)\n self.decoder3 = self._block((features * 4) * 2, features * 4, name=\"dec3\")\n \n self.upconv2 = nn.ConvTranspose3d(features * 4, features * 2, kernel_size=2, stride=2)\n self.decoder2 = self._block((features * 2) * 2, features * 2, name=\"dec2\")\n \n self.upconv1 = nn.ConvTranspose3d(features * 2, features, kernel_size=2, stride=2)\n self.decoder1 = self._block(features * 2, features, name=\"dec1\")\n \n self.conv = nn.Conv3d(in_channels=features, out_channels=out_channels, kernel_size=1)\n \n def forward(self, x):\n enc1 = self.encoder1(x)\n enc2 = self.encoder2(self.pool1(enc1))\n enc3 = self.encoder3(self.pool2(enc2))\n \n bottleneck = self.bottleneck(self.pool3(enc3))\n \n dec3 = self.upconv3(bottleneck)\n dec3 = torch.cat((dec3, enc3), dim=1)\n dec3 = self.decoder3(dec3)\n \n dec2 = self.upconv2(dec3)\n dec2 = torch.cat((dec2, enc2), dim=1)\n dec2 = self.decoder2(dec2)\n \n dec1 = self.upconv1(dec2)\n dec1 = torch.cat((dec1, enc1), dim=1)\n dec1 = self.decoder1(dec1)\n \n return torch.sigmoid(self.conv(dec1))\n \n @staticmethod\n def _block(in_channels, features, name):\n return nn.Sequential(\n nn.Conv3d(in_channels, features, kernel_size=3, padding=1, bias=False),\n nn.BatchNorm3d(features),\n nn.ReLU(inplace=True),\n nn.Conv3d(features, features, kernel_size=3, padding=1, bias=False),\n nn.BatchNorm3d(features),\n nn.ReLU(inplace=True)\n )\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:31:06.680470Z","iopub.execute_input":"2025-04-03T08:31:06.680807Z","iopub.status.idle":"2025-04-03T08:31:06.691611Z","shell.execute_reply.started":"2025-04-03T08:31:06.680782Z","shell.execute_reply":"2025-04-03T08:31:06.690788Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Dataset class for validation data\nclass ValDataset(Dataset):\n def __init__(self, patches, labels):\n self.patches = patches\n self.labels = labels\n \n def __len__(self):\n return len(self.patches)\n \n def __getitem__(self, idx):\n # Add channel dimension and convert to torch tensors\n patch = torch.FloatTensor(self.patches[idx]).unsqueeze(0)\n label = torch.FloatTensor(self.labels[idx]).unsqueeze(0)\n \n return patch, label","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:31:13.161622Z","iopub.execute_input":"2025-04-03T08:31:13.162091Z","iopub.status.idle":"2025-04-03T08:31:13.166993Z","shell.execute_reply.started":"2025-04-03T08:31:13.162051Z","shell.execute_reply":"2025-04-03T08:31:13.166083Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to load a tomogram\ndef load_tomogram(zarr_path, resolution=0):\n \"\"\"\n Load a tomogram from a zarr file\n \n Parameters:\n zarr_path (str): Path to the zarr directory\n resolution (int): Resolution level (0 = highest, 1 = medium, 2 = lowest)\n \n Returns:\n numpy.ndarray: The loaded tomogram data\n \"\"\"\n try:\n # Open the zarr group\n z = zarr.open(zarr_path, mode='r')\n \n # Access the resolution level directly (based on the structure we observed)\n if str(resolution) in z:\n tomo_data = z[str(resolution)][:]\n print(f\"Loaded tomogram with shape {tomo_data.shape}\")\n return tomo_data\n else:\n print(f\"Resolution level {resolution} not found in zarr file\")\n return None\n except Exception as e:\n print(f\"Error loading tomogram: {str(e)}\")\n return None","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:31:16.893547Z","iopub.execute_input":"2025-04-03T08:31:16.893893Z","iopub.status.idle":"2025-04-03T08:31:16.898752Z","shell.execute_reply.started":"2025-04-03T08:31:16.893865Z","shell.execute_reply":"2025-04-03T08:31:16.897921Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to preprocess a tomogram\ndef preprocess_tomogram(tomo_data):\n \"\"\"\n Preprocess a tomogram for better visualization\n \n Parameters:\n tomo_data (numpy.ndarray): 3D tomogram data\n \n Returns:\n numpy.ndarray: Preprocessed tomogram data\n \"\"\"\n # Make a copy to avoid modifying the original\n processed = tomo_data.copy()\n \n # Normalize to [0, 1] range\n min_val = processed.min()\n max_val = processed.max()\n if max_val > min_val:\n processed = (processed - min_val) / (max_val - min_val)\n \n # Enhance contrast\n processed = np.clip((processed - 0.1) * 1.25, 0, 1)\n \n return processed\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:31:20.997124Z","iopub.execute_input":"2025-04-03T08:31:20.997454Z","iopub.status.idle":"2025-04-03T08:31:21.001923Z","shell.execute_reply.started":"2025-04-03T08:31:20.997427Z","shell.execute_reply":"2025-04-03T08:31:21.000947Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to find local maxima in the density map\ndef find_local_maxima(density_map, min_distance=10, threshold_abs=0.3, threshold_rel=0.2):\n \"\"\"\n Find local maxima in the density map\n \n Parameters:\n density_map (numpy.ndarray): Density map\n min_distance (int): Minimum distance between peaks\n threshold_abs (float): Minimum absolute threshold for peak\n threshold_rel (float): Minimum relative threshold for peak\n \n Returns:\n numpy.ndarray: Array of peak coordinates [z, y, x]\n \"\"\"\n # Import here to keep dependencies clean\n from scipy.ndimage import gaussian_filter\n from skimage.feature import peak_local_max\n \n # Apply Gaussian smoothing to reduce noise\n smoothed_map = gaussian_filter(density_map, sigma=1.0)\n \n # Find local maxima\n coordinates = peak_local_max(\n smoothed_map,\n min_distance=min_distance,\n threshold_abs=threshold_abs,\n threshold_rel=threshold_rel,\n exclude_border=False\n )\n \n return coordinates","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:31:24.814944Z","iopub.execute_input":"2025-04-03T08:31:24.815276Z","iopub.status.idle":"2025-04-03T08:31:24.820164Z","shell.execute_reply.started":"2025-04-03T08:31:24.815248Z","shell.execute_reply":"2025-04-03T08:31:24.819126Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to prepare a sample of validation tomograms for visualization\ndef prepare_validation_tomograms():\n \"\"\"\n Prepare a few validation tomograms and their particle annotations for visualization\n \n This is a simplified approach that visualizes a few tomograms with their ground truth annotations\n \"\"\"\n # Get train experiments\n train_experiments = [os.path.basename(p) for p in glob.glob(os.path.join(train_dir, 'static/ExperimentRuns/*'))]\n \n if not train_experiments:\n print(\"No training experiments found.\")\n return\n \n # Use the first 3 experiments for visualization (or fewer if less are available)\n viz_experiments = train_experiments[:min(3, len(train_experiments))]\n print(f\"Using {len(viz_experiments)} experiments for visualization: {viz_experiments}\")\n \n # For each experiment, load the tomogram and particle annotations\n for experiment in viz_experiments:\n print(f\"\\nProcessing experiment: {experiment}\")\n \n # Load tomogram\n zarr_path = os.path.join(train_dir, 'static/ExperimentRuns', experiment, 'VoxelSpacing10.000/denoised.zarr')\n \n if not os.path.exists(zarr_path):\n print(f\"Tomogram not found for experiment {experiment}\")\n continue\n \n tomo_data = load_tomogram(zarr_path)\n if tomo_data is None:\n print(f\"Failed to load tomogram for experiment {experiment}\")\n continue\n \n tomo_data = preprocess_tomogram(tomo_data)\n \n # Load particle annotations\n ground_truth_coords = {}\n \n for p_type in particle_types.keys():\n json_path = os.path.join(train_dir, 'overlay/ExperimentRuns', experiment, 'Picks', f\"{p_type}.json\")\n \n if not os.path.exists(json_path):\n print(f\"No annotations found for {p_type} in experiment {experiment}\")\n continue\n \n # Load coordinates from JSON\n try:\n with open(json_path, 'r') as f:\n data = json.load(f)\n \n # Extract coordinates from points\n coords = []\n if 'points' in data:\n for point in data['points']:\n if 'location' in point:\n loc = point['location']\n coords.append((loc.get('x', 0), loc.get('y', 0), loc.get('z', 0)))\n \n if coords:\n ground_truth_coords[p_type] = coords\n print(f\"Loaded {len(coords)} {p_type} coordinates\")\n except Exception as e:\n print(f\"Error loading {json_path}: {str(e)}\")\n \n # Visualize ground truth\n visualize_ground_truth(tomo_data, ground_truth_coords, experiment)\n \n # Now visualize model predictions on the same tomogram\n visualize_model_predictions(tomo_data, experiment, ground_truth_coords)\n \n # Free memory\n #del tomo_data\n #gc.collect()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:31:30.397470Z","iopub.execute_input":"2025-04-03T08:31:30.397808Z","iopub.status.idle":"2025-04-03T08:31:30.405805Z","shell.execute_reply.started":"2025-04-03T08:31:30.397783Z","shell.execute_reply":"2025-04-03T08:31:30.404801Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to visualize ground truth\ndef visualize_ground_truth(tomo_data, ground_truth_coords, experiment):\n \"\"\"\n Visualize ground truth particles on the tomogram\n \n Parameters:\n tomo_data (numpy.ndarray): 3D tomogram data\n ground_truth_coords (dict): Dictionary mapping particle types to coordinates\n experiment (str): Experiment name\n \"\"\"\n # Get tomogram dimensions\n depth, height, width = tomo_data.shape\n \n # Choose slices for visualization\n slices = [depth // 4, depth // 2, 3 * depth // 4]\n \n # Create figure with subplots for different slices\n fig, axes = plt.subplots(1, len(slices), figsize=(6 * len(slices), 6))\n if len(slices) == 1:\n axes = [axes]\n \n # For each slice\n for i, slice_idx in enumerate(slices):\n # Show the tomogram slice\n axes[i].imshow(tomo_data[slice_idx], cmap='gray')\n axes[i].set_title(f'Ground Truth - Z-Slice {slice_idx}/{depth}')\n \n # Get slice range (particles near this slice)\n slice_range = 10 # Consider particles within ±10 slices\n z_min = (slice_idx - slice_range) * 10.0 # Convert to physical coordinates\n z_max = (slice_idx + slice_range) * 10.0\n \n # Add circles for each particle type\n for p_type, coords in ground_truth_coords.items():\n # Skip if particle type not in our dictionary\n if p_type not in particle_types:\n continue\n \n # Get color and radius for this particle type\n color = particle_types[p_type]['color']\n radius = particle_types[p_type]['radius'] * 0.1 # Scale down for visualization\n \n # Count particles in this slice\n slice_particles = [(x, y, z) for x, y, z in coords if z_min <= z <= z_max]\n n_particles = len(slice_particles)\n \n # Skip if no particles of this type in this slice\n if n_particles == 0:\n continue\n \n # Add to legend\n axes[i].plot([], [], 'o', color=color, label=f'{p_type} ({n_particles})')\n \n # Add circle for each particle\n for x, y, z in slice_particles:\n # Convert physical coordinates to pixel coordinates\n y_px = y / 10.0\n x_px = x / 10.0\n \n # Add circle\n circle = patches.Circle((x_px, y_px), radius, color=color, fill=False, linewidth=1.5, alpha=0.7)\n axes[i].add_patch(circle)\n \n # Add legend\n axes[i].legend(loc='upper right', bbox_to_anchor=(1.1, 1))\n axes[i].axis('off')\n \n plt.tight_layout()\n plt.savefig(os.path.join(visualization_dir, f'{experiment}_ground_truth.png'), dpi=200, bbox_inches='tight')\n plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:31:40.431889Z","iopub.execute_input":"2025-04-03T08:31:40.432181Z","iopub.status.idle":"2025-04-03T08:31:40.440779Z","shell.execute_reply.started":"2025-04-03T08:31:40.432160Z","shell.execute_reply":"2025-04-03T08:31:40.439752Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to visualize model predictions\ndef visualize_model_predictions(tomo_data, experiment, ground_truth_coords):\n \"\"\"\n Visualize model predictions on a tomogram\n \n Parameters:\n tomo_data (numpy.ndarray): 3D tomogram data\n experiment (str): Experiment name\n ground_truth_coords (dict): Dictionary mapping particle types to ground truth coordinates (for comparison)\n \"\"\"\n # Load the trained model\n model_path = os.path.join(model_dir, 'model.pth')\n \n if not os.path.exists(model_path):\n print(f\"Model not found at {model_path}. Please train the model first.\")\n return\n \n # Initialize model\n model = UNet3D(in_channels=1, out_channels=1, init_features=16)\n \n # Load model weights\n model.load_state_dict(torch.load(model_path))\n model = model.to(device)\n model.eval()\n \n # Get tomogram dimensions\n depth, height, width = tomo_data.shape\n \n # Create a simplified 3D density map\n print(\"Generating simplified density map for visualization...\")\n density_map = np.zeros_like(tomo_data)\n patch_size = 64\n half_size = patch_size // 2\n stride = 32 # Use a stride to reduce computation\n \n # Use a sliding window approach to generate density map\n with torch.no_grad():\n for z in tqdm(range(half_size, depth - half_size, stride)):\n for y in range(half_size, height - half_size, stride):\n for x in range(half_size, width - half_size, stride):\n # Extract patch\n patch = tomo_data[\n z - half_size:z + half_size,\n y - half_size:y + half_size,\n x - half_size:x + half_size\n ]\n \n # Skip if patch is invalid\n if patch.shape != (patch_size, patch_size, patch_size):\n continue\n \n # Convert to tensor\n patch_tensor = torch.FloatTensor(patch).unsqueeze(0).unsqueeze(0).to(device)\n \n # Forward pass\n output = model(patch_tensor)\n \n # Convert to numpy\n output = output.detach().cpu().numpy()[0, 0]\n \n # Add to density map\n density_map[\n z - half_size:z + half_size,\n y - half_size:y + half_size,\n x - half_size:x + half_size\n ] = np.maximum(\n density_map[\n z - half_size:z + half_size,\n y - half_size:y + half_size,\n x - half_size:x + half_size\n ],\n output\n )\n \n # Create a combined visualization with side-by-side ground truth and predictions\n visualize_comparison(tomo_data, density_map, ground_truth_coords, experiment)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:31:46.953338Z","iopub.execute_input":"2025-04-03T08:31:46.953677Z","iopub.status.idle":"2025-04-03T08:31:46.961528Z","shell.execute_reply.started":"2025-04-03T08:31:46.953639Z","shell.execute_reply":"2025-04-03T08:31:46.960535Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to visualize comparison between ground truth and predictions\ndef visualize_comparison(tomo_data, density_map, ground_truth_coords, experiment):\n \"\"\"\n Visualize comparison between ground truth and model predictions\n \n Parameters:\n tomo_data (numpy.ndarray): 3D tomogram data\n density_map (numpy.ndarray): 3D density map from model predictions\n ground_truth_coords (dict): Dictionary mapping particle types to ground truth coordinates\n experiment (str): Experiment name\n \"\"\"\n # Get tomogram dimensions\n depth, height, width = tomo_data.shape\n \n # Choose slices for visualization\n slices = [depth // 4, depth // 2, 3 * depth // 4]\n \n # Find local maxima in density map\n print(\"Finding local maxima in density map...\")\n coords = find_local_maxima(density_map, min_distance=10, threshold_abs=0.3, threshold_rel=0.2)\n print(f\"Found {len(coords)} local maxima\")\n \n # Simplified particle type assignment for visualization\n # This doesn't replicate the full clustering logic but is sufficient for visualization\n predicted_coords = {}\n \n # Basic heuristic: stronger peaks are more likely to be larger/easier particles\n peak_values = np.array([density_map[z, y, x] for z, y, x in coords])\n peak_order = np.argsort(peak_values)[::-1] # Sort in descending order\n \n # Assign top 20% to easier particles, next 30% to harder particles\n n_peaks = len(coords)\n easy_threshold = int(n_peaks * 0.2)\n hard_threshold = int(n_peaks * 0.5)\n \n easy_particles = ['apo-ferritin', 'ribosome', 'virus-like-particle']\n hard_particles = ['beta-galactosidase', 'thyroglobulin']\n \n # Randomly assign among the categories\n np.random.seed(42) # For reproducibility\n \n for i, idx in enumerate(peak_order):\n z, y, x = coords[idx]\n \n # Convert voxel coordinates to physical coordinates\n physical_x = x * 10.0\n physical_y = y * 10.0\n physical_z = z * 10.0\n \n if i < easy_threshold:\n # Assign to an easy particle type\n p_type = np.random.choice(easy_particles)\n elif i < hard_threshold:\n # Assign to a hard particle type\n p_type = np.random.choice(hard_particles)\n else:\n # Skip the rest\n continue\n \n if p_type not in predicted_coords:\n predicted_coords[p_type] = []\n \n predicted_coords[p_type].append((physical_x, physical_y, physical_z))\n \n # For each slice, create a side-by-side comparison\n for slice_idx in slices:\n fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(20, 8))\n \n # Ground truth visualization\n ax1.imshow(tomo_data[slice_idx], cmap='gray')\n ax1.set_title(f'Ground Truth - Z-Slice {slice_idx}/{depth}')\n \n # Add a semi-transparent overlay of the density map\n ax2.imshow(tomo_data[slice_idx], cmap='gray')\n density_overlay = ax2.imshow(density_map[slice_idx], cmap='hot', alpha=0.5)\n fig.colorbar(density_overlay, ax=ax2, fraction=0.046, pad=0.04)\n ax2.set_title(f'Model Predictions - Z-Slice {slice_idx}/{depth}')\n \n # Get slice range (particles near this slice)\n slice_range = 10 # Consider particles within ±10 slices\n z_min = (slice_idx - slice_range) * 10.0 # Convert to physical coordinates\n z_max = (slice_idx + slice_range) * 10.0\n \n # Add ground truth particles\n for p_type, coords in ground_truth_coords.items():\n # Skip if particle type not in our dictionary\n if p_type not in particle_types:\n continue\n \n # Get color and radius for this particle type\n color = particle_types[p_type]['color']\n radius = particle_types[p_type]['radius'] * 0.1 # Scale down for visualization\n \n # Count particles in this slice\n slice_particles = [(x, y, z) for x, y, z in coords if z_min <= z <= z_max]\n n_particles = len(slice_particles)\n \n # Skip if no particles of this type in this slice\n if n_particles == 0:\n continue\n \n # Add to legend\n ax1.plot([], [], 'o', color=color, label=f'{p_type} ({n_particles})')\n \n # Add circle for each particle\n for x, y, z in slice_particles:\n # Convert physical coordinates to pixel coordinates\n y_px = y / 10.0\n x_px = x / 10.0\n \n # Add circle\n circle = patches.Circle((x_px, y_px), radius, color=color, fill=False, linewidth=1.5, alpha=0.7)\n ax1.add_patch(circle)\n \n # Add predicted particles\n for p_type, coords in predicted_coords.items():\n # Skip if particle type not in our dictionary\n if p_type not in particle_types:\n continue\n \n # Get color and radius for this particle type\n color = particle_types[p_type]['color']\n radius = particle_types[p_type]['radius'] * 0.1 # Scale down for visualization\n \n # Count particles in this slice\n slice_particles = [(x, y, z) for x, y, z in coords if z_min <= z <= z_max]\n n_particles = len(slice_particles)\n \n # Skip if no particles of this type in this slice\n if n_particles == 0:\n continue\n \n # Add to legend\n ax2.plot([], [], 'o', color=color, label=f'{p_type} ({n_particles})')\n \n # Add circle for each particle\n for x, y, z in slice_particles:\n # Convert physical coordinates to pixel coordinates\n y_px = y / 10.0\n x_px = x / 10.0\n \n # Add circle\n circle = patches.Circle((x_px, y_px), radius, color=color, fill=False, linewidth=1.5, alpha=0.7)\n ax2.add_patch(circle)\n \n # Add legends\n ax1.legend(loc='upper right', bbox_to_anchor=(1.1, 1))\n ax2.legend(loc='upper right', bbox_to_anchor=(1.1, 1))\n \n # Remove axes\n ax1.axis('off')\n ax2.axis('off')\n \n plt.tight_layout()\n plt.savefig(os.path.join(visualization_dir, f'{experiment}_slice_{slice_idx}_comparison.png'), dpi=200, bbox_inches='tight')\n plt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:32:00.584741Z","iopub.execute_input":"2025-04-03T08:32:00.585069Z","iopub.status.idle":"2025-04-03T08:32:00.599134Z","shell.execute_reply.started":"2025-04-03T08:32:00.585042Z","shell.execute_reply":"2025-04-03T08:32:00.598164Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Main execution\nprint(\"Starting validation tomogram visualization...\")\nprepare_validation_tomograms()\nprint(\"Validation tomogram visualization complete.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:32:26.015631Z","iopub.execute_input":"2025-04-03T08:32:26.016007Z","iopub.status.idle":"2025-04-03T08:34:29.271877Z","shell.execute_reply.started":"2025-04-03T08:32:26.015979Z","shell.execute_reply":"2025-04-03T08:34:29.271023Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Test Prediction and Visualization","metadata":{}},{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport matplotlib.pyplot as plt\nimport matplotlib.patches as patches\nfrom matplotlib.colors import LinearSegmentedColormap\nimport zarr\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nimport gc\nimport glob\nfrom tqdm.notebook import tqdm\nfrom scipy.ndimage import gaussian_filter\nfrom skimage.feature import peak_local_max\nimport json","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:35:06.585567Z","iopub.execute_input":"2025-04-03T08:35:06.585912Z","iopub.status.idle":"2025-04-03T08:35:06.590684Z","shell.execute_reply.started":"2025-04-03T08:35:06.585887Z","shell.execute_reply":"2025-04-03T08:35:06.589792Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Set paths\nbase_dir = '/kaggle/input/czii-cryo-et-object-identification'\ntrain_dir = os.path.join(base_dir, 'train')\ntest_dir = os.path.join(base_dir, 'test')\noutput_dir = '/kaggle/working/preprocessed'\nmodel_dir = '/kaggle/working/models'\nsubmission_dir = '/kaggle/working/final_submission'\nvisualization_dir = '/kaggle/working/final_visualizations'\n\n# Create directories if they don't exist\nos.makedirs(submission_dir, exist_ok=True)\nos.makedirs(visualization_dir, exist_ok=True)\nos.makedirs(output_dir, exist_ok=True)\n\n# Set device\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"Using device: {device}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:35:13.112044Z","iopub.execute_input":"2025-04-03T08:35:13.112339Z","iopub.status.idle":"2025-04-03T08:35:13.118498Z","shell.execute_reply.started":"2025-04-03T08:35:13.112318Z","shell.execute_reply":"2025-04-03T08:35:13.117813Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Set random seeds for reproducibility\nnp.random.seed(42)\ntorch.manual_seed(42)\nif torch.cuda.is_available():\n torch.cuda.manual_seed(42)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:35:17.306250Z","iopub.execute_input":"2025-04-03T08:35:17.306540Z","iopub.status.idle":"2025-04-03T08:35:17.312733Z","shell.execute_reply.started":"2025-04-03T08:35:17.306519Z","shell.execute_reply":"2025-04-03T08:35:17.311951Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Particle types and their properties\nparticle_types = {\n 'apo-ferritin': {'difficulty': 'easy', 'weight': 1, 'radius': 60, 'color': 'red'},\n 'beta-amylase': {'difficulty': 'impossible', 'weight': 0, 'radius': 45, 'color': 'yellow'},\n 'beta-galactosidase': {'difficulty': 'hard', 'weight': 2, 'radius': 80, 'color': 'green'},\n 'ribosome': {'difficulty': 'easy', 'weight': 1, 'radius': 100, 'color': 'blue'},\n 'thyroglobulin': {'difficulty': 'hard', 'weight': 2, 'radius': 85, 'color': 'purple'},\n 'virus-like-particle': {'difficulty': 'easy', 'weight': 1, 'radius': 120, 'color': 'orange'}\n}\n\n# Get scored particle types (weight > 0)\nscored_particle_types = [p for p, props in particle_types.items() if props['weight'] > 0]\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:35:21.061931Z","iopub.execute_input":"2025-04-03T08:35:21.062249Z","iopub.status.idle":"2025-04-03T08:35:21.067803Z","shell.execute_reply.started":"2025-04-03T08:35:21.062225Z","shell.execute_reply":"2025-04-03T08:35:21.066779Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 3D U-Net model (same as in training)\nclass UNet3D(nn.Module):\n def __init__(self, in_channels=1, out_channels=1, init_features=16):\n super(UNet3D, self).__init__()\n \n features = init_features\n self.encoder1 = self._block(in_channels, features, name=\"enc1\")\n self.pool1 = nn.MaxPool3d(kernel_size=2, stride=2)\n \n self.encoder2 = self._block(features, features * 2, name=\"enc2\")\n self.pool2 = nn.MaxPool3d(kernel_size=2, stride=2)\n \n self.encoder3 = self._block(features * 2, features * 4, name=\"enc3\")\n self.pool3 = nn.MaxPool3d(kernel_size=2, stride=2)\n \n self.bottleneck = self._block(features * 4, features * 8, name=\"bottleneck\")\n \n self.upconv3 = nn.ConvTranspose3d(features * 8, features * 4, kernel_size=2, stride=2)\n self.decoder3 = self._block((features * 4) * 2, features * 4, name=\"dec3\")\n \n self.upconv2 = nn.ConvTranspose3d(features * 4, features * 2, kernel_size=2, stride=2)\n self.decoder2 = self._block((features * 2) * 2, features * 2, name=\"dec2\")\n \n self.upconv1 = nn.ConvTranspose3d(features * 2, features, kernel_size=2, stride=2)\n self.decoder1 = self._block(features * 2, features, name=\"dec1\")\n \n self.conv = nn.Conv3d(in_channels=features, out_channels=out_channels, kernel_size=1)\n \n def forward(self, x):\n enc1 = self.encoder1(x)\n enc2 = self.encoder2(self.pool1(enc1))\n enc3 = self.encoder3(self.pool2(enc2))\n \n bottleneck = self.bottleneck(self.pool3(enc3))\n \n dec3 = self.upconv3(bottleneck)\n dec3 = torch.cat((dec3, enc3), dim=1)\n dec3 = self.decoder3(dec3)\n \n dec2 = self.upconv2(dec3)\n dec2 = torch.cat((dec2, enc2), dim=1)\n dec2 = self.decoder2(dec2)\n \n dec1 = self.upconv1(dec2)\n dec1 = torch.cat((dec1, enc1), dim=1)\n dec1 = self.decoder1(dec1)\n \n return torch.sigmoid(self.conv(dec1))\n \n @staticmethod\n def _block(in_channels, features, name):\n return nn.Sequential(\n nn.Conv3d(in_channels, features, kernel_size=3, padding=1, bias=False),\n nn.BatchNorm3d(features),\n nn.ReLU(inplace=True),\n nn.Conv3d(features, features, kernel_size=3, padding=1, bias=False),\n nn.BatchNorm3d(features),\n nn.ReLU(inplace=True)\n )\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:35:26.019976Z","iopub.execute_input":"2025-04-03T08:35:26.020278Z","iopub.status.idle":"2025-04-03T08:35:26.032031Z","shell.execute_reply.started":"2025-04-03T08:35:26.020257Z","shell.execute_reply":"2025-04-03T08:35:26.031140Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to load a tomogram\ndef load_tomogram(zarr_path, resolution=0):\n \"\"\"\n Load a tomogram from a zarr file\n \n Parameters:\n zarr_path (str): Path to the zarr directory\n resolution (int): Resolution level (0 = highest, 1 = medium, 2 = lowest)\n \n Returns:\n numpy.ndarray: The loaded tomogram data\n \"\"\"\n try:\n # Open the zarr group\n z = zarr.open(zarr_path, mode='r')\n \n # Access the resolution level directly\n if str(resolution) in z:\n tomo_data = z[str(resolution)][:]\n print(f\"Loaded tomogram with shape {tomo_data.shape}\")\n return tomo_data\n else:\n print(f\"Resolution level {resolution} not found in zarr file\")\n return None\n except Exception as e:\n print(f\"Error loading tomogram: {str(e)}\")\n return None\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:35:33.285824Z","iopub.execute_input":"2025-04-03T08:35:33.286129Z","iopub.status.idle":"2025-04-03T08:35:33.291402Z","shell.execute_reply.started":"2025-04-03T08:35:33.286106Z","shell.execute_reply":"2025-04-03T08:35:33.290616Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to preprocess a tomogram\ndef preprocess_tomogram(tomo_data):\n \"\"\"\n Preprocess a tomogram for better prediction\n \n Parameters:\n tomo_data (numpy.ndarray): 3D tomogram data\n \n Returns:\n numpy.ndarray: Preprocessed tomogram data\n \"\"\"\n # Make a copy to avoid modifying the original\n processed = tomo_data.copy()\n \n # Normalize to [0, 1] range\n min_val = processed.min()\n max_val = processed.max()\n if max_val > min_val:\n processed = (processed - min_val) / (max_val - min_val)\n \n # Enhance contrast\n processed = np.clip((processed - 0.1) * 1.25, 0, 1)\n \n return processed","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:35:38.635614Z","iopub.execute_input":"2025-04-03T08:35:38.635969Z","iopub.status.idle":"2025-04-03T08:35:38.640560Z","shell.execute_reply.started":"2025-04-03T08:35:38.635943Z","shell.execute_reply":"2025-04-03T08:35:38.639425Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to load particle coordinates from JSON\ndef load_particle_coords(json_path):\n \"\"\"\n Load particle coordinates from JSON file\n \n Parameters:\n json_path (str): Path to the JSON file\n \n Returns:\n list: List of (x, y, z) coordinate tuples\n \"\"\"\n try:\n with open(json_path, 'r') as f:\n data = json.load(f)\n \n # Extract coordinates from points\n coords = []\n if 'points' in data:\n for point in data['points']:\n if 'location' in point:\n loc = point['location']\n coords.append((loc.get('x', 0), loc.get('y', 0), loc.get('z', 0)))\n \n return coords\n except Exception as e:\n print(f\"Error loading {json_path}: {str(e)}\")\n return []","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:35:45.340562Z","iopub.execute_input":"2025-04-03T08:35:45.340934Z","iopub.status.idle":"2025-04-03T08:35:45.346229Z","shell.execute_reply.started":"2025-04-03T08:35:45.340893Z","shell.execute_reply":"2025-04-03T08:35:45.345263Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to generate a density map for a full tomogram\ndef generate_density_map(model, tomo_data, patch_size=64, stride=32, batch_size=8):\n \"\"\"\n Generate a density map for a full tomogram\n \n Parameters:\n model (nn.Module): Trained model\n tomo_data (numpy.ndarray): 3D tomogram data\n patch_size (int): Size of cubic patches\n stride (int): Stride between patches\n batch_size (int): Batch size for prediction\n \n Returns:\n numpy.ndarray: Density map\n \"\"\"\n model.eval()\n \n # Get dimensions\n depth, height, width = tomo_data.shape\n \n # Initialize density map and count array\n density_map = np.zeros_like(tomo_data, dtype=np.float32)\n count = np.zeros_like(tomo_data, dtype=np.float32)\n \n # Half size for patch extraction\n half_size = patch_size // 2\n \n # Extract patches\n patches = []\n coordinates = []\n \n print(\"Extracting patches...\")\n for z in tqdm(range(half_size, depth - half_size, stride)):\n for y in range(half_size, height - half_size, stride):\n for x in range(half_size, width - half_size, stride):\n # Extract the patch\n patch = tomo_data[\n z - half_size:z + half_size,\n y - half_size:y + half_size,\n x - half_size:x + half_size\n ]\n \n # Skip if patch is invalid\n if patch.shape != (patch_size, patch_size, patch_size):\n continue\n \n patches.append(patch)\n coordinates.append((z, y, x))\n \n # Process in batches to avoid memory issues\n if len(patches) >= batch_size:\n # Convert to tensor\n batch_patches = torch.FloatTensor(np.array(patches)).unsqueeze(1).to(device)\n \n # Predict\n with torch.no_grad():\n batch_outputs = model(batch_patches)\n \n # Move to CPU and convert to numpy\n batch_outputs = batch_outputs.detach().cpu().numpy()\n \n # Add to density map\n for i, (z, y, x) in enumerate(coordinates):\n output = batch_outputs[i, 0]\n \n # Add to density map\n density_map[\n z - half_size:z + half_size,\n y - half_size:y + half_size,\n x - half_size:x + half_size\n ] += output\n \n # Increment count\n count[\n z - half_size:z + half_size,\n y - half_size:y + half_size,\n x - half_size:x + half_size\n ] += 1\n \n # Clear lists\n patches = []\n coordinates = []\n \n # Process remaining patches\n if patches:\n # Convert to tensor\n batch_patches = torch.FloatTensor(np.array(patches)).unsqueeze(1).to(device)\n \n # Predict\n with torch.no_grad():\n batch_outputs = model(batch_patches)\n \n # Move to CPU and convert to numpy\n batch_outputs = batch_outputs.detach().cpu().numpy()\n \n # Add to density map\n for i, (z, y, x) in enumerate(coordinates):\n output = batch_outputs[i, 0]\n \n # Add to density map\n density_map[\n z - half_size:z + half_size,\n y - half_size:y + half_size,\n x - half_size:x + half_size\n ] += output\n \n # Increment count\n count[\n z - half_size:z + half_size,\n y - half_size:y + half_size,\n x - half_size:x + half_size\n ] += 1\n \n # Average overlapping regions\n density_map = np.divide(density_map, count, out=np.zeros_like(density_map), where=count > 0)\n \n return density_map","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:35:53.840834Z","iopub.execute_input":"2025-04-03T08:35:53.841165Z","iopub.status.idle":"2025-04-03T08:35:53.853374Z","shell.execute_reply.started":"2025-04-03T08:35:53.841140Z","shell.execute_reply":"2025-04-03T08:35:53.852386Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to analyze validation tomograms for threshold calibration\ndef analyze_validation_tomograms():\n \"\"\"\n Analyze validation tomograms to calibrate thresholds\n \n Returns:\n dict: Dictionary of calibrated thresholds\n \"\"\"\n # Load the trained model\n model_path = os.path.join(model_dir, 'model.pth')\n \n if not os.path.exists(model_path):\n print(f\"Model not found at {model_path}. Please train the model first.\")\n return None\n \n # Initialize model\n model = UNet3D(in_channels=1, out_channels=1, init_features=16)\n \n # Load model weights\n model.load_state_dict(torch.load(model_path))\n model = model.to(device)\n model.eval()\n \n # Get list of training experiments\n train_experiments = [os.path.basename(p) for p in glob.glob(os.path.join(train_dir, 'static/ExperimentRuns/*'))]\n \n if not train_experiments:\n print(\"No training experiments found.\")\n return None\n \n # Sample a few experiments for analysis\n sample_size = min(3, len(train_experiments))\n validation_experiments = train_experiments[:sample_size]\n print(f\"Using {len(validation_experiments)} experiments for threshold calibration: {validation_experiments}\")\n \n # Dictionary to store signal intensities for each particle type\n particle_intensities = {p_type: [] for p_type in particle_types}\n \n # Process each validation experiment\n for experiment in validation_experiments:\n print(f\"\\nProcessing validation experiment: {experiment}\")\n \n # Load tomogram\n zarr_path = os.path.join(train_dir, 'static/ExperimentRuns', experiment, 'VoxelSpacing10.000/denoised.zarr')\n \n if not os.path.exists(zarr_path):\n print(f\"Tomogram not found for experiment {experiment}\")\n continue\n \n tomo_data = load_tomogram(zarr_path)\n if tomo_data is None:\n print(f\"Failed to load tomogram for experiment {experiment}\")\n continue\n \n tomo_data = preprocess_tomogram(tomo_data)\n \n # Generate density map\n print(\"Generating density map...\")\n density_map = generate_density_map(model, tomo_data, patch_size=64, stride=32, batch_size=8)\n \n # Load ground truth particle positions\n ground_truth = {}\n \n for p_type in particle_types:\n json_path = os.path.join(train_dir, 'overlay/ExperimentRuns', experiment, 'Picks', f\"{p_type}.json\")\n \n if os.path.exists(json_path):\n coords = load_particle_coords(json_path)\n if coords:\n ground_truth[p_type] = coords\n print(f\"Loaded {len(coords)} {p_type} coordinates\")\n \n # Sample intensity values at ground truth positions\n for p_type, coords in ground_truth.items():\n for x, y, z in coords:\n # Convert physical coordinates to voxel indices\n z_idx, y_idx, x_idx = int(z / 10), int(y / 10), int(x / 10)\n \n # Check if the position is within the density map\n if (0 <= z_idx < density_map.shape[0] and \n 0 <= y_idx < density_map.shape[1] and \n 0 <= x_idx < density_map.shape[2]):\n \n # Get the intensity value\n intensity = density_map[z_idx, y_idx, x_idx]\n particle_intensities[p_type].append(intensity)\n \n # Free memory\n del tomo_data, density_map\n gc.collect()\n if torch.cuda.is_available():\n torch.cuda.empty_cache()\n \n # Calculate intensity statistics for each particle type\n intensity_stats = {}\n \n print(\"\\nParticle intensity statistics:\")\n for p_type, intensities in particle_intensities.items():\n if intensities:\n stats = {\n 'mean': np.mean(intensities),\n 'median': np.median(intensities),\n 'min': np.min(intensities),\n 'max': np.max(intensities),\n 'p25': np.percentile(intensities, 25),\n 'p10': np.percentile(intensities, 10),\n 'p5': np.percentile(intensities, 5),\n 'count': len(intensities)\n }\n intensity_stats[p_type] = stats\n \n print(f\" - {p_type} ({len(intensities)} particles):\")\n print(f\" Mean: {stats['mean']:.4f}, Median: {stats['median']:.4f}\")\n print(f\" Min: {stats['min']:.4f}, Max: {stats['max']:.4f}\")\n print(f\" P5: {stats['p5']:.4f}, P10: {stats['p10']:.4f}, P25: {stats['p25']:.4f}\")\n \n # Visualize intensity distributions\n plt.figure(figsize=(12, 8))\n \n for p_type, intensities in particle_intensities.items():\n if len(intensities) > 10: # Only plot if we have enough samples\n plt.hist(intensities, bins=20, alpha=0.6, label=f\"{p_type} (n={len(intensities)})\")\n \n plt.xlabel('Density Value')\n plt.ylabel('Frequency')\n plt.title('Distribution of Density Values at Ground Truth Particle Positions')\n plt.legend()\n plt.grid(True, alpha=0.3)\n plt.savefig(os.path.join(visualization_dir, 'intensity_distributions.png'), dpi=200)\n plt.close()\n \n # Calculate calibrated thresholds based on validation analysis\n # Use much lower percentiles to ensure we find all types\n calibrated_thresholds = {}\n \n for p_type in scored_particle_types:\n if p_type in intensity_stats:\n stats = intensity_stats[p_type]\n \n # Use very low percentiles - prioritize recall over precision\n if p_type == 'apo-ferritin' or p_type == 'ribosome':\n # For easier particle types\n calibrated_thresholds[p_type] = stats['p5'] * 0.7 # 70% of 5th percentile\n elif p_type == 'beta-galactosidase':\n # For harder particle types\n calibrated_thresholds[p_type] = stats['p5'] * 0.6 # 60% of 5th percentile\n elif p_type == 'thyroglobulin':\n # More difficult to detect\n calibrated_thresholds[p_type] = stats['p5'] * 0.5 # 50% of 5th percentile\n elif p_type == 'virus-like-particle':\n # Most difficult to detect\n calibrated_thresholds[p_type] = stats['p5'] * 0.4 # 40% of 5th percentile\n else:\n # Default\n calibrated_thresholds[p_type] = stats['p5'] * 0.6\n \n # If we don't have statistics for some particle types, use defaults\n default_threshold = 0.002 # Very low default threshold\n for p_type in scored_particle_types:\n if p_type not in calibrated_thresholds:\n if p_type == 'thyroglobulin':\n calibrated_thresholds[p_type] = 0.0015 # Very low threshold\n elif p_type == 'virus-like-particle':\n calibrated_thresholds[p_type] = 0.0010 # Extremely low threshold\n elif p_type == 'beta-galactosidase':\n calibrated_thresholds[p_type] = 0.0018 # Low threshold\n else:\n calibrated_thresholds[p_type] = default_threshold\n \n print(\"\\nCalibrated thresholds:\")\n for p_type, threshold in calibrated_thresholds.items():\n print(f\" - {p_type}: {threshold:.4f}\")\n \n return calibrated_thresholds\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:36:04.854498Z","iopub.execute_input":"2025-04-03T08:36:04.854861Z","iopub.status.idle":"2025-04-03T08:36:04.871332Z","shell.execute_reply.started":"2025-04-03T08:36:04.854831Z","shell.execute_reply":"2025-04-03T08:36:04.870450Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to find local maxima in the density map\ndef find_local_maxima(density_map, min_distance=6, threshold_abs=0.05, threshold_rel=0.03):\n \"\"\"\n Find local maxima in the density map\n \n Parameters:\n density_map (numpy.ndarray): Density map\n min_distance (int): Minimum distance between peaks\n threshold_abs (float): Minimum absolute threshold for peak\n threshold_rel (float): Minimum relative threshold for peak\n \n Returns:\n numpy.ndarray: Array of peak coordinates [z, y, x]\n \"\"\"\n # Apply Gaussian smoothing to reduce noise\n smoothed_map = gaussian_filter(density_map, sigma=1.0)\n \n # Find local maxima\n coordinates = peak_local_max(\n smoothed_map,\n min_distance=min_distance,\n threshold_abs=threshold_abs,\n threshold_rel=threshold_rel,\n exclude_border=False\n )\n \n return coordinates","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:36:15.255511Z","iopub.execute_input":"2025-04-03T08:36:15.255873Z","iopub.status.idle":"2025-04-03T08:36:15.260319Z","shell.execute_reply.started":"2025-04-03T08:36:15.255844Z","shell.execute_reply":"2025-04-03T08:36:15.259450Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Enhanced clustering function to detect all particle types\ndef calibrated_cluster_particles(coordinates, volume, particle_types, calibrated_thresholds):\n \"\"\"\n Cluster particles using calibrated thresholds with a multi-pass approach to ensure all types are detected\n \n Parameters:\n coordinates (numpy.ndarray): Array of peak coordinates [z, y, x]\n volume (numpy.ndarray): Density map\n particle_types (list): List of particle types to consider\n calibrated_thresholds (dict): Calibrated thresholds for each particle type\n \n Returns:\n dict: Dictionary mapping particle type to list of coordinates\n \"\"\"\n if len(coordinates) == 0:\n print(\"No coordinates found for clustering\")\n return {p_type: [] for p_type in particle_types}\n \n # Sort peaks by intensity\n peak_values = np.array([volume[z, y, x] for z, y, x in coordinates])\n sorted_indices = np.argsort(peak_values)[::-1] # Sort in descending order\n \n sorted_coordinates = coordinates[sorted_indices]\n sorted_values = peak_values[sorted_indices]\n \n # Dictionary to store particles by type\n particles_by_type = {p_type: [] for p_type in particle_types}\n \n # Assign peaks to particle types based on calibrated thresholds\n n_particles = {p_type: 0 for p_type in particle_types}\n \n # Set target and maximum number of particles per tomogram\n # Increased target numbers to ensure we find all types\n target_particles = {\n 'apo-ferritin': 80,\n 'beta-galactosidase': 60,\n 'ribosome': 80,\n 'thyroglobulin': 60,\n 'virus-like-particle': 40\n }\n \n max_particles = {\n 'apo-ferritin': 120,\n 'beta-galactosidase': 100,\n 'ribosome': 120,\n 'thyroglobulin': 100,\n 'virus-like-particle': 80\n }\n \n # Define size categories by radius for more balanced distribution\n size_categories = {\n 'small': ['apo-ferritin'],\n 'medium': ['beta-galactosidase', 'thyroglobulin'],\n 'large': ['ribosome', 'virus-like-particle']\n }\n \n # Create reverse mapping of particle type to category\n particle_to_category = {}\n for category, p_types in size_categories.items():\n for p_type in p_types:\n particle_to_category[p_type] = category\n \n # Track assigned peaks\n already_assigned = set()\n \n # First pass: Assign top peaks to different size categories as evenly as possible\n size_counts = {cat: 0 for cat in size_categories}\n category_quotas = {\n 'small': 0.30, # 30% for small particles\n 'medium': 0.40, # 40% for medium particles\n 'large': 0.30 # 30% for large particles\n }\n \n total_expected = min(len(sorted_coordinates), sum(target_particles.values()))\n category_limits = {\n cat: int(total_expected * quota) for cat, quota in category_quotas.items()\n }\n \n print(\"\\nCategory distribution targets:\")\n for cat, limit in category_limits.items():\n print(f\" - {cat}: {limit} particles\")\n \n # First pass: Assign peaks to categories based on thresholds and type quotas\n for i, (z, y, x) in enumerate(sorted_coordinates):\n peak_value = sorted_values[i]\n \n # Skip if already assigned\n if i in already_assigned:\n continue\n \n # Skip if this peak is too close to an already assigned peak\n too_close = False\n for idx in already_assigned:\n z2, y2, x2 = sorted_coordinates[idx]\n dist = np.sqrt((z - z2)**2 + (y - y2)**2 + (x - x2)**2)\n if dist < 5: # 5 voxels minimum distance (reduced from original 6)\n too_close = True\n break\n \n if too_close:\n continue\n \n # Try to assign to particle types that are under their limits\n assigned = False\n \n # Prioritize underrepresented categories\n categories_sorted = sorted(size_categories.keys(), key=lambda cat: size_counts[cat] / max(1, category_limits[cat]))\n \n for category in categories_sorted:\n if size_counts[category] >= category_limits[category] * 1.5: # Allow going over by 50%\n continue # Skip if category is well beyond limit\n \n # Try each particle type in this category\n p_types = size_categories[category]\n \n # Sort types by how far they are from their target\n p_types_sorted = sorted(p_types, \n key=lambda p: (n_particles[p] / max(1, target_particles[p])))\n \n for p_type in p_types_sorted:\n if p_type not in calibrated_thresholds:\n continue # Skip if we don't have a threshold\n \n if n_particles[p_type] >= max_particles[p_type]:\n continue # Skip if at limit\n \n # Check threshold - use a lower threshold as we go deeper into the sorted list\n threshold_factor = 1.0 - (i / len(sorted_coordinates)) * 0.2 # Gradually decrease threshold (reduced from 0.3)\n effective_threshold = calibrated_thresholds[p_type] * threshold_factor\n \n if peak_value >= effective_threshold:\n # Convert voxel coordinates to physical coordinates\n physical_coords = (x * 10.0, y * 10.0, z * 10.0)\n particles_by_type[p_type].append(physical_coords)\n n_particles[p_type] += 1\n already_assigned.add(i)\n size_counts[category] += 1\n assigned = True\n break\n \n if assigned:\n break\n \n # Print distribution after first pass\n print(\"\\nParticle distribution after first pass:\")\n for p_type in particle_types:\n print(f\" - {p_type}: {n_particles[p_type]} particles\")\n \n # Check for missing or underrepresented types\n underrepresented = []\n for p_type in particle_types:\n # If we have less than 20% of target for this type, consider it underrepresented\n if n_particles[p_type] < target_particles.get(p_type, 50) * 0.2:\n underrepresented.append(p_type)\n \n print(f\"Underrepresented types: {underrepresented}\")\n \n # Second pass: Focus on underrepresented types with much lower thresholds\n if underrepresented:\n print(\"Running second pass for underrepresented types...\")\n \n # Sort by intensity again - we'll use the remaining strong signals\n remaining_indices = [i for i in range(len(sorted_coordinates)) if i not in already_assigned]\n \n for i in remaining_indices:\n z, y, x = sorted_coordinates[i]\n peak_value = sorted_values[i]\n \n # Skip if too close to an already assigned peak\n too_close = False\n for idx in already_assigned:\n z2, y2, x2 = sorted_coordinates[idx]\n dist = np.sqrt((z - z2)**2 + (y - y2)**2 + (x - x2)**2)\n if dist < 5: # 5 voxels minimum distance\n too_close = True\n break\n \n if too_close:\n continue\n \n # Try to assign to underrepresented types with very low thresholds\n for p_type in underrepresented:\n # Target at least 20% of the expected count for each type\n min_target = target_particles.get(p_type, 50) * 0.2\n \n if n_particles[p_type] >= min_target:\n continue\n \n # Use a very low threshold - just to get some representation\n very_low_threshold = calibrated_thresholds[p_type] * 0.4 # 40% of calibrated threshold\n \n if peak_value >= very_low_threshold:\n # Convert voxel coordinates to physical coordinates\n physical_coords = (x * 10.0, y * 10.0, z * 10.0)\n particles_by_type[p_type].append(physical_coords)\n n_particles[p_type] += 1\n already_assigned.add(i)\n \n # Update category count\n category = particle_to_category.get(p_type)\n if category:\n size_counts[category] += 1\n \n break # Assign at most one particle per coordinate\n \n # Print distribution after second pass\n print(\"\\nParticle distribution after second pass:\")\n for p_type in particle_types:\n print(f\" - {p_type}: {n_particles[p_type]} particles\")\n \n # Check if we still have missing particle types\n missing_types = [p_type for p_type in particle_types if n_particles[p_type] == 0]\n \n # Third pass: Desperate measures for still-missing types\n if missing_types:\n print(f\"Still missing types: {missing_types}. Using desperate measures.\")\n \n # Use remaining peaks with extremely low thresholds\n remaining_indices = [i for i in range(len(sorted_coordinates)) if i not in already_assigned]\n \n # For each missing type, try to find at least a few candidates\n for p_type in missing_types:\n # Count needed for minimum representation\n min_count = 5 # At least 5 particles of each type\n \n # Find peaks for this type, sorted by intensity\n candidate_indices = []\n for i in remaining_indices:\n z, y, x = sorted_coordinates[i]\n # Use almost no threshold - we just need some representation\n candidate_indices.append((i, sorted_values[i], z, y, x))\n \n # Sort candidates by intensity\n candidate_indices.sort(key=lambda x: x[1], reverse=True)\n \n # Take top candidates that aren't too close to existing particles\n for i, value, z, y, x in candidate_indices:\n if n_particles[p_type] >= min_count:\n break\n \n # Check if too close to already assigned particles\n too_close = False\n for idx in already_assigned:\n z2, y2, x2 = sorted_coordinates[idx]\n dist = np.sqrt((z - z2)**2 + (y - y2)**2 + (x - x2)**2)\n if dist < 5: # 5 voxels minimum distance\n too_close = True\n break\n \n if too_close:\n continue\n \n # Convert voxel coordinates to physical coordinates\n physical_coords = (x * 10.0, y * 10.0, z * 10.0)\n particles_by_type[p_type].append(physical_coords)\n n_particles[p_type] += 1\n already_assigned.add(i)\n \n # Update category count\n category = particle_to_category.get(p_type)\n if category:\n size_counts[category] += 1\n \n print(f\"Added {n_particles[p_type]} emergency particles for {p_type}\")\n \n # Fourth pass: Fill in with remaining high-value peaks to reach targets\n remaining_indices = [i for i in range(len(sorted_coordinates)) if i not in already_assigned]\n \n # For each particle type that's below target\n below_target_types = [(p, target_particles.get(p, 50) - n_particles[p]) \n for p in particle_types \n if n_particles[p] < target_particles.get(p, 50)]\n \n # Sort by how far below target they are\n below_target_types.sort(key=lambda x: x[1], reverse=True)\n \n if below_target_types:\n print(\"\\nRunning fourth pass to reach targets for types below target...\")\n \n for p_type, deficit in below_target_types:\n if deficit <= 0:\n continue\n \n # Use lower thresholds for final pass\n final_threshold = calibrated_thresholds[p_type] * 0.3 # 30% of calibrated threshold\n \n # Count particles added for this type\n added = 0\n \n for i in remaining_indices[:]: # Create a copy to modify during iteration\n if added >= deficit:\n break\n \n z, y, x = sorted_coordinates[i]\n peak_value = sorted_values[i]\n \n # Skip if already assigned (could happen during this pass)\n if i in already_assigned:\n continue\n \n # Skip if too close to an already assigned peak\n too_close = False\n for idx in already_assigned:\n z2, y2, x2 = sorted_coordinates[idx]\n dist = np.sqrt((z - z2)**2 + (y - y2)**2 + (x - x2)**2)\n if dist < 5: # 5 voxels minimum distance\n too_close = True\n break\n \n if too_close:\n continue\n \n # Check if this peak meets our threshold\n if peak_value >= final_threshold:\n # Convert voxel coordinates to physical coordinates\n physical_coords = (x * 10.0, y * 10.0, z * 10.0)\n particles_by_type[p_type].append(physical_coords)\n n_particles[p_type] += 1\n already_assigned.add(i)\n added += 1\n \n # Update category count\n category = particle_to_category.get(p_type)\n if category:\n size_counts[category] += 1\n \n print(f\"Added {added} additional particles for {p_type}\")\n \n # Print final statistics\n print(\"\\nFinal particle distribution:\")\n for p_type in particle_types:\n print(f\" - {p_type}: {n_particles[p_type]} particles\")\n \n print(f\"Total particles found: {sum(n_particles.values())}\")\n \n return particles_by_type","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:36:23.594925Z","iopub.execute_input":"2025-04-03T08:36:23.595246Z","iopub.status.idle":"2025-04-03T08:36:23.619137Z","shell.execute_reply.started":"2025-04-03T08:36:23.595220Z","shell.execute_reply":"2025-04-03T08:36:23.618199Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to write predictions to JSON\ndef write_predictions_to_json(particles_by_type, output_path):\n \"\"\"\n Write particle predictions to JSON\n \n Parameters:\n particles_by_type (dict): Dictionary mapping particle type to list of coordinates\n output_path (str): Path to save the JSON file\n \"\"\"\n # Format the predictions according to the submission format\n prediction = {\"points\": []}\n \n for p_type, coords in particles_by_type.items():\n for x, y, z in coords:\n prediction[\"points\"].append({\n \"location\": {\"x\": x, \"y\": y, \"z\": z},\n \"type\": p_type\n })\n \n # Write to file\n with open(output_path, 'w') as f:\n json.dump(prediction, f)\n \n print(f\"Wrote {len(prediction['points'])} particles to {output_path}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:36:38.037837Z","iopub.execute_input":"2025-04-03T08:36:38.038163Z","iopub.status.idle":"2025-04-03T08:36:38.043109Z","shell.execute_reply.started":"2025-04-03T08:36:38.038141Z","shell.execute_reply":"2025-04-03T08:36:38.042197Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to visualize predictions with different colors for different particle types\ndef visualize_tomogram_with_particles(tomo_data, particles_by_type, output_path, slices=None):\n \"\"\"\n Visualize tomogram slices with colored particle markers\n \n Parameters:\n tomo_data (numpy.ndarray): 3D tomogram data\n particles_by_type (dict): Dictionary mapping particle type to list of coordinates\n output_path (str): Path to save the visualization\n slices (list): List of slice indices to visualize (default is middle slice)\n \"\"\"\n # Get tomogram dimensions\n depth, height, width = tomo_data.shape\n \n # Choose slices if not provided\n if slices is None:\n slices = [depth // 4, depth // 2, 3 * depth // 4]\n \n # Create figure\n fig, axes = plt.subplots(1, len(slices), figsize=(6 * len(slices), 6))\n if len(slices) == 1:\n axes = [axes]\n \n # For each slice\n for i, slice_idx in enumerate(slices):\n # Show the tomogram slice\n axes[i].imshow(tomo_data[slice_idx], cmap='gray')\n axes[i].set_title(f'Z-Slice {slice_idx}/{depth}')\n \n # Get slice range (particles near this slice)\n slice_range = 10 # Consider particles within ±10 slices\n z_min = (slice_idx - slice_range) * 10.0 # Convert to physical coordinates\n z_max = (slice_idx + slice_range) * 10.0\n \n # Add circles for each particle type\n for p_type, coords in particles_by_type.items():\n # Skip if particle type not in our dictionary or no coordinates\n if p_type not in particle_types or not coords:\n continue\n \n # Get color and radius for this particle type\n color = particle_types[p_type]['color']\n radius = particle_types[p_type]['radius'] / 10.0 # Convert to voxel units\n \n # Count particles in this slice\n slice_particles = [(x, y, z) for x, y, z in coords if z_min <= z <= z_max]\n n_particles = len(slice_particles)\n \n # Skip if no particles of this type in this slice\n if n_particles == 0:\n continue\n \n # Add to legend\n axes[i].plot([], [], 'o', color=color, label=f'{p_type} ({n_particles})')\n \n # Add circle for each particle\n for x, y, z in slice_particles:\n # Convert physical coordinates to pixel coordinates\n y_px = y / 10.0\n x_px = x / 10.0\n \n # Calculate alpha based on distance from the slice\n z_px = z / 10.0\n alpha = 1.0 - abs(z_px - slice_idx) / slice_range\n \n # Add circle\n circle = plt.Circle((x_px, y_px), radius, color=color, fill=False, alpha=alpha, linewidth=1.5)\n axes[i].add_patch(circle)\n \n # Add legend\n axes[i].legend(loc='upper right', bbox_to_anchor=(1.1, 1))\n axes[i].axis('off')\n \n plt.tight_layout()\n plt.savefig(output_path, dpi=200, bbox_inches='tight')\n plt.close()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:36:43.647173Z","iopub.execute_input":"2025-04-03T08:36:43.647471Z","iopub.status.idle":"2025-04-03T08:36:43.656195Z","shell.execute_reply.started":"2025-04-03T08:36:43.647449Z","shell.execute_reply":"2025-04-03T08:36:43.655243Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to visualize density map with particle locations\ndef visualize_density_map_with_particles(density_map, particles_by_type, output_path, slice_idx=None):\n \"\"\"\n Visualize density map with particle locations\n \n Parameters:\n density_map (numpy.ndarray): 3D density map\n particles_by_type (dict): Dictionary mapping particle type to list of coordinates\n output_path (str): Path to save the visualization\n slice_idx (int): Index of slice to visualize (default is middle slice)\n \"\"\"\n # Get dimensions\n depth, height, width = density_map.shape\n \n # Choose slice if not provided\n if slice_idx is None:\n slice_idx = depth // 2\n \n # Create a custom colormap for density\n cmap_name = 'hot_alpha'\n colors = [(0, 0, 0, 0)] # Start with transparent\n for i in range(1, 256):\n # Red-yellow colormap with increasing alpha\n alpha = i / 255.0\n if i < 128:\n # From transparent to red\n colors.append((i / 127.0, 0, 0, alpha * 0.7))\n else:\n # From red to yellow\n colors.append((1, (i - 128) / 127.0, 0, alpha * 0.7))\n \n custom_cmap = LinearSegmentedColormap.from_list(cmap_name, colors, N=256)\n \n # Create figure\n plt.figure(figsize=(12, 10))\n \n # Show the density map slice\n plt.imshow(density_map[slice_idx], cmap=custom_cmap)\n \n # Get slice range (particles near this slice)\n slice_range = 10 # Consider particles within ±10 slices\n z_min = (slice_idx - slice_range) * 10.0 # Convert to physical coordinates\n z_max = (slice_idx + slice_range) * 10.0\n \n # Add markers for each particle type\n for p_type, coords in particles_by_type.items():\n # Skip if particle type not in our dictionary or no coordinates\n if p_type not in particle_types or not coords:\n continue\n \n # Get color for this particle type\n color = particle_types[p_type]['color']\n \n # Count particles in this slice\n slice_particles = [(x, y, z) for x, y, z in coords if z_min <= z <= z_max]\n n_particles = len(slice_particles)\n \n # Skip if no particles of this type in this slice\n if n_particles == 0:\n continue\n \n # Extract x, y coordinates for this slice\n x_coords = [x / 10.0 for x, y, z in slice_particles]\n y_coords = [y / 10.0 for x, y, z in slice_particles]\n \n # Plot particles\n plt.scatter(x_coords, y_coords, color=color, marker='o', s=50, facecolors='none', \n label=f'{p_type} ({n_particles})', linewidth=1.5)\n \n plt.title(f'Density Map with Particles (Z-Slice {slice_idx}/{depth})', fontsize=14)\n plt.colorbar(label='Density Value')\n plt.legend(loc='upper right', bbox_to_anchor=(1.25, 1))\n plt.axis('off')\n plt.tight_layout()\n plt.savefig(output_path, dpi=200, bbox_inches='tight')\n plt.close()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:36:50.561960Z","iopub.execute_input":"2025-04-03T08:36:50.562347Z","iopub.status.idle":"2025-04-03T08:36:50.575123Z","shell.execute_reply.started":"2025-04-03T08:36:50.562317Z","shell.execute_reply":"2025-04-03T08:36:50.574134Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to evaluate the prediction quality\ndef evaluate_prediction(particles_by_type):\n \"\"\"\n Evaluate prediction quality using heuristic metrics\n \n Parameters:\n particles_by_type (dict): Dictionary mapping particle type to list of coordinates\n \n Returns:\n float: Heuristic quality score (0-1)\n \"\"\"\n # Calculate score based on particle distribution\n scores = []\n \n # Expected particle counts based on observations\n expected_counts = {\n 'apo-ferritin': 80,\n 'beta-galactosidase': 60,\n 'ribosome': 80,\n 'thyroglobulin': 60,\n 'virus-like-particle': 40\n }\n \n for p_type, coords in particles_by_type.items():\n if p_type not in expected_counts:\n continue\n \n # Count\n count = len(coords)\n expected = expected_counts[p_type]\n \n # Calculate normalized score (1.0 if count matches expected, less otherwise)\n # Use min-max normalization\n ratio = min(count / expected, 1.0) if expected > 0 else 0.0\n \n # Weight by particle importance\n weight = particle_types[p_type].get('weight', 1.0)\n \n # Add to scores\n scores.append(ratio * weight)\n \n # Calculate final score\n total_weight = sum(particle_types[p_type].get('weight', 1.0) for p_type in expected_counts if p_type in particle_types)\n \n if total_weight > 0:\n return sum(scores) / total_weight\n else:\n return 0.0","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:36:58.066904Z","iopub.execute_input":"2025-04-03T08:36:58.067236Z","iopub.status.idle":"2025-04-03T08:36:58.072927Z","shell.execute_reply.started":"2025-04-03T08:36:58.067212Z","shell.execute_reply":"2025-04-03T08:36:58.072089Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Main function to process test tomograms\ndef process_test_tomograms():\n \"\"\"\n Process test tomograms and generate predictions\n \"\"\"\n # Load the trained model\n model_path = os.path.join(model_dir, 'model.pth')\n \n if not os.path.exists(model_path):\n print(f\"Model not found at {model_path}. Please train the model first.\")\n return\n \n # Initialize model\n model = UNet3D(in_channels=1, out_channels=1, init_features=16)\n \n # Load model weights\n model.load_state_dict(torch.load(model_path, map_location=device))\n model = model.to(device)\n model.eval()\n \n # Get calibrated thresholds\n calibrated_thresholds = analyze_validation_tomograms()\n \n if calibrated_thresholds is None:\n print(\"Failed to calibrate thresholds. Using default values.\")\n calibrated_thresholds = {\n 'apo-ferritin': 0.0020,\n 'beta-galactosidase': 0.0018,\n 'ribosome': 0.0020,\n 'thyroglobulin': 0.0015,\n 'virus-like-particle': 0.0010\n }\n \n # Modify thresholds to ensure better detection of all types\n # Further lower thresholds for all types to ensure better detection\n print(\"\\nAdjusting thresholds for better particle type detection:\")\n \n # Hardcode lower thresholds to ensure all types are detected\n for p_type in calibrated_thresholds:\n # Reduce all thresholds by 20%\n calibrated_thresholds[p_type] *= 0.8\n \n # Specifically lower thresholds for problematic types\n if 'thyroglobulin' in calibrated_thresholds:\n calibrated_thresholds['thyroglobulin'] = min(calibrated_thresholds['thyroglobulin'], 0.0015)\n \n if 'virus-like-particle' in calibrated_thresholds:\n calibrated_thresholds['virus-like-particle'] = min(calibrated_thresholds['virus-like-particle'], 0.0008)\n \n if 'beta-galactosidase' in calibrated_thresholds:\n calibrated_thresholds['beta-galactosidase'] = min(calibrated_thresholds['beta-galactosidase'], 0.0015)\n \n for p_type, threshold in calibrated_thresholds.items():\n print(f\" - {p_type}: {threshold:.4f}\")\n \n # Get list of test experiments\n test_experiments = [os.path.basename(p) for p in glob.glob(os.path.join(test_dir, 'static/ExperimentRuns/*'))]\n \n if not test_experiments:\n print(\"No test experiments found.\")\n return\n \n print(f\"\\nFound {len(test_experiments)} test experiments: {test_experiments}\")\n \n # List to store all particle predictions for final submission\n all_particles = []\n \n # Process each test experiment\n for experiment in test_experiments:\n print(f\"\\nProcessing test experiment: {experiment}\")\n \n # Create experiment output directory\n experiment_output_dir = os.path.join(output_dir, 'test', experiment)\n os.makedirs(experiment_output_dir, exist_ok=True)\n \n # Load tomogram\n zarr_path = os.path.join(test_dir, 'static/ExperimentRuns', experiment, 'VoxelSpacing10.000/denoised.zarr')\n \n if not os.path.exists(zarr_path):\n print(f\"Tomogram not found for experiment {experiment}\")\n continue\n \n tomo_data = load_tomogram(zarr_path)\n if tomo_data is None:\n print(f\"Failed to load tomogram for experiment {experiment}\")\n continue\n \n tomo_data = preprocess_tomogram(tomo_data)\n \n # Generate density map\n print(\"Generating density map...\")\n density_map = generate_density_map(model, tomo_data, patch_size=64, stride=32, batch_size=8)\n \n # Save density map\n np.save(os.path.join(experiment_output_dir, 'density_map.npy'), density_map)\n \n # Find local maxima\n print(\"Finding local maxima...\")\n # Use very low thresholds to ensure we detect harder particles\n coordinates = find_local_maxima(\n density_map, \n min_distance=6, # Decreased from 8 to detect smaller particles\n threshold_abs=0.05, # Decreased to be more sensitive\n threshold_rel=0.03 # Decreased to be more sensitive\n )\n \n print(f\"Found {len(coordinates)} potential particle locations\")\n \n # Cluster particles\n print(\"Clustering particles...\")\n particles_by_type = calibrated_cluster_particles(\n coordinates, \n density_map, \n scored_particle_types,\n calibrated_thresholds\n )\n \n # Evaluate prediction quality\n quality_score = evaluate_prediction(particles_by_type)\n print(f\"Prediction quality score: {quality_score:.4f}\")\n \n # Check for missing particle types\n missing_types = [p_type for p_type in scored_particle_types if len(particles_by_type[p_type]) == 0]\n \n # If quality score is too low or we're missing particle types, try again with even lower thresholds\n if quality_score < 0.5 or missing_types:\n print(f\"Low quality prediction (score: {quality_score:.4f}) or missing types: {missing_types}. Retrying with lower thresholds...\")\n \n # Try with even lower thresholds\n lower_thresholds = {p_type: t * 0.6 for p_type, t in calibrated_thresholds.items()}\n \n # Especially lower thresholds for missing types\n for p_type in missing_types:\n lower_thresholds[p_type] = calibrated_thresholds[p_type] * 0.4\n \n # Find local maxima again with even lower thresholds\n coordinates = find_local_maxima(\n density_map, \n min_distance=5, # Even smaller minimum distance\n threshold_abs=0.03, # Much lower threshold\n threshold_rel=0.02 # Much lower relative threshold\n )\n \n print(f\"Found {len(coordinates)} potential particle locations (retry)\")\n \n # Cluster particles again\n particles_by_type = calibrated_cluster_particles(\n coordinates, \n density_map, \n scored_particle_types,\n lower_thresholds\n )\n \n # Re-evaluate\n quality_score = evaluate_prediction(particles_by_type)\n print(f\"New prediction quality score: {quality_score:.4f}\")\n \n # If still missing particle types, use desperate measures\n missing_types = [p_type for p_type in scored_particle_types if len(particles_by_type[p_type]) == 0]\n if missing_types:\n print(f\"Still missing types: {missing_types}. Trying desperate measures...\")\n \n # For each missing type, find at least a few candidates\n for p_type in missing_types:\n # Use peak_local_max with very low thresholds specific to this type\n desperate_coords = peak_local_max(\n density_map,\n min_distance=4,\n threshold_abs=0.01,\n threshold_rel=0.01,\n exclude_border=False,\n num_peaks=10 # Limit to top 10 peaks\n )\n \n # Add top 5 coordinates to this particle type\n for z, y, x in desperate_coords[:5]:\n physical_coords = (x * 10.0, y * 10.0, z * 10.0)\n particles_by_type[p_type].append(physical_coords)\n \n print(f\"Added {min(5, len(desperate_coords))} emergency particles for {p_type}\")\n \n # Write predictions to JSON\n output_path = os.path.join(submission_dir, f\"{experiment}.json\")\n write_predictions_to_json(particles_by_type, output_path)\n \n # Visualize tomogram with particles\n vis_path = os.path.join(visualization_dir, f\"{experiment}_tomogram_with_particles.png\")\n visualize_tomogram_with_particles(tomo_data, particles_by_type, vis_path)\n \n # Visualize density map with particles\n density_vis_path = os.path.join(visualization_dir, f\"{experiment}_density_map_with_particles.png\")\n visualize_density_map_with_particles(density_map, particles_by_type, density_vis_path)\n \n # Add to list of all particles for CSV submission\n for p_type, coords in particles_by_type.items():\n for x, y, z in coords:\n all_particles.append({\n 'experiment': experiment,\n 'particle_type': p_type,\n 'x': x,\n 'y': y,\n 'z': z\n })\n \n # Free memory\n del tomo_data, density_map\n gc.collect()\n if torch.cuda.is_available():\n torch.cuda.empty_cache()\n \n # Create CSV submission file\n if all_particles:\n # Create DataFrame\n submission_df = pd.DataFrame(all_particles)\n \n # Add id column\n submission_df['id'] = range(len(submission_df))\n \n # Reorder columns\n submission_df = submission_df[['id', 'experiment', 'particle_type', 'x', 'y', 'z']]\n \n # Save submission\n submission_path = os.path.join(submission_dir, 'submission.csv')\n submission_df.to_csv(submission_path, index=False)\n \n print(f\"\\nSaved submission to {submission_path}\")\n print(f\"Total predictions: {len(submission_df)}\")\n \n # Print submission statistics\n print(\"\\nSubmission statistics:\")\n print(submission_df.groupby(['experiment', 'particle_type']).size().unstack(fill_value=0))\n \n print(\"\\nAll test tomograms processed.\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:37:07.229679Z","iopub.execute_input":"2025-04-03T08:37:07.230043Z","iopub.status.idle":"2025-04-03T08:37:07.249123Z","shell.execute_reply.started":"2025-04-03T08:37:07.230017Z","shell.execute_reply":"2025-04-03T08:37:07.248111Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Run the main function\nif __name__ == \"__main__\":\n print(\"Starting test tomogram processing...\")\n process_test_tomograms()\n print(\"Test tomogram processing completed.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:37:21.247613Z","iopub.execute_input":"2025-04-03T08:37:21.247982Z","iopub.status.idle":"2025-04-03T08:39:33.095956Z","shell.execute_reply.started":"2025-04-03T08:37:21.247955Z","shell.execute_reply":"2025-04-03T08:39:33.094986Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Function to display test images\ndef display_test_images():\n \"\"\"\n Display all test images generated in the visualization directory\n \"\"\"\n import glob\n import matplotlib.pyplot as plt\n from IPython.display import display\n \n # Get all visualization files\n vis_files = glob.glob(os.path.join(visualization_dir, '*.png'))\n \n if not vis_files:\n print(\"No visualization files found in\", visualization_dir)\n return\n \n print(f\"Found {len(vis_files)} visualization files.\")\n \n # Display each visualization file\n for vis_file in sorted(vis_files):\n filename = os.path.basename(vis_file)\n print(f\"\\nDisplaying: {filename}\")\n \n # Load and display the image\n img = plt.imread(vis_file)\n plt.figure(figsize=(15, 10))\n plt.imshow(img)\n plt.axis('off')\n plt.title(filename, fontsize=14)\n plt.tight_layout()\n plt.show()\n\n# Call the function to display the images\ndisplay_test_images()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-03T08:42:03.142946Z","iopub.execute_input":"2025-04-03T08:42:03.143327Z","iopub.status.idle":"2025-04-03T08:42:10.205899Z","shell.execute_reply.started":"2025-04-03T08:42:03.143299Z","shell.execute_reply":"2025-04-03T08:42:10.204853Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}