diff --git a/server_api/chatbot/agent_eval.py b/server_api/chatbot/agent_eval.py new file mode 100644 index 00000000..9b58b957 --- /dev/null +++ b/server_api/chatbot/agent_eval.py @@ -0,0 +1,101 @@ +""" +End-to-end agent evaluation for training command generation. + +Sends 10 realistic user requests through the full agent pipeline and +prints tool call logs + full responses for manual review. +""" + +import sys +import re +from io import StringIO + +TRAINING_REQUESTS = [ + # Basic config selection + "Give me the training command for mitochondria segmentation on the Lucchi dataset", + "Give me the command to train synapse detection on CREMI", + "Write the training command for neuron segmentation on SNEMI3D", + "Generate a command to train on the MitoEM-R dataset", + "Give me the command to train nucleus segmentation on zebrafish data", + "Write the command for vesicle segmentation training", + "Give me a quick demo training command to test if everything works", + # Single override + "Give me the command to train CREMI synapse detection with learning rate 0.0001", + "Write the training command for Lucchi mitochondria with 300 epochs", + "Generate the command to train SNEMI3D neuron segmentation with batch size 4", + "Give me the command to train MitoEM-H with AdamW optimizer", + "Write the command for Lucchi mito training with bf16 mixed precision", + # Multiple overrides + "Generate the training command for vesicle segmentation with learning rate 0.001 and gradient clipping of 1.0", + "Give me the command to train Lucchi mito with 200 epochs and learning rate 0.0005", + "Write the command to train CREMI synapse detection, saving checkpoints every 5 epochs with early stopping patience 20", + # Advanced / tricky + "Give me the command for neuron instance segmentation training at 40nm resolution with gradient clipping 1.0", + "Generate the training command for MitoEM-R with cosine learning rate schedule and 500 epochs", + "Write the command to train on the MitoEM human+rat combined dataset with 4 GPUs", + "Give me the training command for Lucchi mitochondria, but I want to save only the top 1 checkpoint", + "Generate a command to train SNEMI3D neurons with learning rate 0.0003, batch size 2, and gradient clipping 0.5", +] + +INFERENCE_REQUESTS = [ + # Basic inference + "Give me the inference command for my trained Lucchi mito model. The checkpoint is at outputs/mito_lucchi/checkpoints/epoch=100.ckpt", + "Write the command to run inference on CREMI with my checkpoint at outputs/syn_cremi/checkpoints/last.ckpt", + "Generate the inference command for SNEMI3D using checkpoint outputs/neuron_snemi/checkpoints/epoch=050.ckpt", + # With overrides + "Give me the inference command for Lucchi mito with checkpoint outputs/mito_lucchi/checkpoints/epoch=100.ckpt and Gaussian blending", + "Write the command to run inference on CREMI with checkpoint outputs/syn_cremi/checkpoints/last.ckpt and test-time augmentation enabled", + "Generate the inference command for MitoEM-R with checkpoint outputs/mitoem_r/checkpoints/epoch=200.ckpt and output saved to results/mitoem_predictions/", + # Advanced + "Give me the command to run inference on vesicle data with checkpoint outputs/vesicle_xm/checkpoints/epoch=150.ckpt, sliding window overlap 0.75, and Gaussian blending", + "Write the inference command for SNEMI3D neuron segmentation using outputs/neuron_snemi/checkpoints/epoch=100.ckpt with TTA enabled and ensemble mode set to mean", + "Generate the command to run inference on the Lucchi dataset with checkpoint outputs/mito_lucchi/checkpoints/best.ckpt and batch size 2", + "Give me the inference command for nucleus segmentation on zebrafish with checkpoint outputs/nuc_nucmm/checkpoints/epoch=080.ckpt and save output to results/nuc_output/", +] + + +def main(): + from server_api.chatbot.chatbot import build_chain + + print("Building agent chain...") + supervisor, reset_search = build_chain() + print("Agent chain ready.\n") + + all_requests = [("TRAIN", r) for r in TRAINING_REQUESTS] + + for i, (category, request) in enumerate(all_requests): + print(f"\n{'=' * 80}") + print(f"[{category}] TEST {i + 1}/{len(all_requests)}") + print(f"USER: {request}") + print("=" * 80) + + # Capture stdout to get tool call logs + old_stdout = sys.stdout + captured = StringIO() + sys.stdout = captured + + try: + reset_search() + result = supervisor.invoke( + {"messages": [{"role": "user", "content": request}]} + ) + messages = result.get("messages", []) + response = messages[-1].content if messages else "(no response)" + except Exception as e: + response = f"ERROR: {e}" + finally: + sys.stdout = old_stdout + tool_log = captured.getvalue() + + # Print tool calls + tool_calls = re.findall(r"\[TOOL\] .*", tool_log) + if tool_calls: + print("\nTOOL CALLS:") + for tc in tool_calls: + print(f" {tc}") + + print(f"\nRESPONSE:\n{response}") + print() + + +if __name__ == "__main__": + main() diff --git a/server_api/chatbot/chatbot.py b/server_api/chatbot/chatbot.py index 6469a748..c2800cce 100644 --- a/server_api/chatbot/chatbot.py +++ b/server_api/chatbot/chatbot.py @@ -9,6 +9,15 @@ import os from pathlib import Path + +def _format_admin_llm_error(error): + """Format an LLM connection error for display to the user.""" + return ( + "The AI assistant could not connect to its configured language model. " + "Please contact your system administrator with this error: " + f"{str(error).strip() or error.__class__.__name__}" + ) + from langchain_ollama import OllamaEmbeddings, ChatOllama from langchain_community.vectorstores import FAISS from langchain_core.tools import tool @@ -30,24 +39,16 @@ 4. Be concise. State the facts, generate the command, stop. WORKFLOW: The available configs are provided in your task message. Pick the best match, then: -1. Call read_config on the chosen config path to see its YAML overrides. -2. For common parameters (learning rate, batch size, iterations, optimizer, checkpoint interval), ALWAYS use the keys listed below. DO NOT search for these. -3. For specialized parameters (augmentation settings, loss functions, architecture details), call search_documentation. -4. Build the command with overrides using the SECTION.KEY=value format. +1. ALWAYS call read_config on the chosen config path BEFORE generating any command. Never skip this. +2. Use ONLY the exact key paths you see in the returned YAML to build overrides (Hydra dot-path format: key=value). +3. For parameters not visible in the config, call search_documentation. +4. Build the command. -IMPORTANT: YAML configs only show overrides — many valid keys exist in the defaults but are not shown in read_config output. +CRITICAL: PyTC has its own config key names — they are NOT PyTorch Lightning CLI names. +Do NOT guess key names. You MUST read the config first to see the actual keys. -Common override keys (ALWAYS use these exact keys, never search for alternatives): -- SOLVER.BASE_LR, SOLVER.SAMPLES_PER_BATCH, SOLVER.ITERATION_TOTAL -- SOLVER.ITERATION_SAVE (checkpoint save interval), SOLVER.ITERATION_STEP (LR decay steps) -- SOLVER.NAME (values: SGD, Adam, AdamW) -- SOLVER.LR_SCHEDULER_NAME (values: WarmupMultiStepLR, WarmupCosineLR) -- SOLVER.CLIP_GRADIENTS.ENABLED (True/False), SOLVER.CLIP_GRADIENTS.CLIP_VALUE -- MODEL.ARCHITECTURE, MODEL.BLOCK_TYPE, MODEL.FILTERS - -NEVER invent keys like TRAIN.MAX_ITER, TRAINING.BATCH_SIZE, or CLI flags like --batch-size, --checkpoint-interval — these do not exist. - -Command format: `python scripts/main.py --config-file [SECTION.KEY=value ...]` +Command format: `python scripts/main.py --config [key=value ...]` +Example: `python scripts/main.py --config tutorials/syn_cremi.yaml optimization.optimizer.lr=0.0001 optimization.max_epochs=300` Always generate commands for the user to run — never execute directly.""" @@ -58,26 +59,13 @@ 2. Be concise. State the facts, generate the command, stop. WORKFLOW: -1. If the user did NOT provide a checkpoint path, call list_checkpoints first to see available checkpoints. -2. If the user DID provide a checkpoint path (e.g., outputs/model/checkpoint.pth.tar), skip list_checkpoints. -3. Call read_config to see the INFERENCE section keys. +1. If the user did NOT provide a checkpoint path, call list_checkpoints first. +2. If the user DID provide one, skip list_checkpoints. +3. Call read_config to see the config’s keys, then use those for overrides (Hydra dot-path format). 4. For specialized inference parameters, call search_documentation if needed. -Here is the correct override key mapping (use these exact keys): -- Output path → INFERENCE.OUTPUT_PATH -- TTA augmentation count → INFERENCE.AUG_NUM -- TTA mode → INFERENCE.AUG_MODE (values: mean, max) -- Blending → INFERENCE.BLENDING (values: gaussian, bump) -- Stride → INFERENCE.STRIDE -- Process volumes one at a time → INFERENCE.DO_SINGLY -- Batch size → INFERENCE.SAMPLES_PER_BATCH - -Command format: `python scripts/main.py --config-file --inference --checkpoint [SECTION.KEY=value ...]` - -IMPORTANT: Overrides use SECTION.KEY=value format (NO -- prefix). Example: - ✅ CORRECT: INFERENCE.AUG_NUM=8 - ❌ WRONG: --inference.AUG_NUM=8 - +Command format: `python scripts/main.py --config --mode test --checkpoint [key=value ...]` +Example: `python scripts/main.py --config tutorials/syn_cremi.yaml --mode test --checkpoint outputs/model/checkpoints/epoch=100.ckpt test.output_path=results/` Always generate commands for the user to run — never execute directly.""" @@ -224,12 +212,53 @@ def delegate_to_training_agent(task: str) -> str: # Auto-inject available configs so the agent doesn't need to call list_training_configs configs = list_training_configs.invoke({}) config_summary = "\n".join( - f"- {c['name']} ({c['model']}) → {c['path']}" for c in configs if isinstance(c, dict) and 'name' in c + f"- {c['name']} ({c['model']}): {c.get('description', '')} → {c['path']}" + for c in configs if isinstance(c, dict) and 'name' in c ) + + # Auto-read a representative config and flatten grouping keys (default/train/test) + # so the agent sees the actual Hydra override paths, not YAML grouping structure + ref_config = next( + (c for c in configs if isinstance(c, dict) and 'lucchi' in c.get('name', '').lower()), + configs[0] if configs and isinstance(configs[0], dict) else None, + ) + key_structure_section = "" + if ref_config: + ref_content = read_config.invoke({"config_path": ref_config['path']}) + if isinstance(ref_content, dict) and 'error' not in ref_content: + import yaml as _yaml + + def _deep_merge(base, override): + """Merge override into base dict recursively.""" + for k, v in override.items(): + if k in base and isinstance(base[k], dict) and isinstance(v, dict): + _deep_merge(base[k], v) + else: + base[k] = v + + # Flatten: merge default, train, test into one level + flat = {} + for group_key in ("default", "train", "test"): + group = ref_content.get(group_key, {}) + if isinstance(group, dict): + _deep_merge(flat, group) + + # Also include top-level non-grouping keys + for k, v in ref_content.items(): + if k not in ("default", "train", "test", "_base_", "description", "experiment_name"): + flat[k] = v + + key_structure_section = ( + f"\n\nOVERRIDE KEY PATHS (from {ref_config['name']} — all configs share this structure):\n" + f"```yaml\n{_yaml.dump(flat, default_flow_style=False)[:3000]}```\n" + f"These are the exact Hydra override paths. Use them as-is (e.g. optimization.max_epochs=200)." + ) + enriched_task = ( f"{task}\n\n" - f"AVAILABLE CONFIGS (already fetched for you):\n{config_summary}\n\n" - f"Pick the best match and call read_config on its path to see the exact YAML keys before generating the command." + f"AVAILABLE CONFIGS (already fetched for you):\n{config_summary}" + f"{key_structure_section}\n\n" + f"Pick the best match. Use the key paths shown above for overrides — do NOT guess key names." ) result = training_agent.invoke( {"messages": [{"role": "user", "content": enriched_task}]} diff --git a/server_api/chatbot/file_summaries/PyTC-Augmentation.md b/server_api/chatbot/file_summaries/PyTC-Augmentation.md index 3745fc35..0bbe0513 100644 --- a/server_api/chatbot/file_summaries/PyTC-Augmentation.md +++ b/server_api/chatbot/file_summaries/PyTC-Augmentation.md @@ -2,152 +2,164 @@ This document describes all data augmentation options available in PyTorch Connectomics. These augmentations are specifically designed for electron microscopy (EM) and biomedical volumetric data. -All augmentations are controlled under the `AUGMENTOR` section of the YAML config. Set `AUGMENTOR.ENABLED: False` to disable all augmentation. +Augmentations are controlled under the `data.augmentation` section and use a **profile-based system**. Set `data.augmentation.profile` to choose a preset, then override individual fields as needed. -## Available Augmentations +## Augmentation Profiles -### Rotation (`AUGMENTOR.ROTATE`) -Applies random 90-degree rotations. +Choose ONE profile via `data.augmentation.profile`: -| Key | Default | Description | -|-----|---------|-------------| -| `ENABLED` | `True` | Enable rotation | -| `ROT90` | `True` | Restrict to 90° increments | -| `P` | `1.0` | Probability of applying | +| Profile | Description | Best For | +|---------|-------------|----------| +| `aug_light` | Flip + rotate + mild intensity | Quick experiments, clean data, fine-tuning | +| `aug_standard` | Light + EM artifact simulation (misalignment, missing sections) | Most 3D EM tasks (RECOMMENDED) | +| `aug_strong` | Maximum regularization (affine, elastic, motion blur, missing parts, cut noise) | Small datasets, overfitting prevention | +| `aug_em_neuron` | DeepEM-matched augmentation with defect mutex | Neuron affinity learning on anisotropic EM (e.g., SNEMI3D) | +| `aug_instance` | Standard + copy-paste + mixup | Instance segmentation | +| `aug_superres` | CutBlur-centric for multi-scale learning | Super-resolution, denoising | + +## Available Augmentations -### Rescale (`AUGMENTOR.RESCALE`) -Randomly rescales the volume. +All augmentation keys are under `data.augmentation.`: + +### Flip (`data.augmentation.flip`) +Randomly flips along spatial axes. | Key | Default | Description | |-----|---------|-------------| -| `ENABLED` | `True` | Enable rescaling | -| `FIX_ASPECT` | `False` | Keep aspect ratio fixed | -| `P` | `0.5` | Probability | +| `enabled` | `true` | Enable flipping | +| `prob` | `0.5` | Probability of applying | +| `spatial_axis` | `[1, 2]` | Axes to flip. Add `0` for z-axis (isotropic data) | -### Flip (`AUGMENTOR.FLIP`) -Randomly flips along axes. For isotropic data, enable z-axis flips. +### Rotate (`data.augmentation.rotate`) +Applies random 90-degree rotations. | Key | Default | Description | |-----|---------|-------------| -| `ENABLED` | `True` | Enable flipping | -| `DO_ZTRANS` | `0` | Set to `1` to enable x-z and y-z flips (for isotropic cubic data) | -| `P` | `1.0` | Probability | +| `enabled` | `true` | Enable rotation | +| `prob` | `0.5` | Probability | +| `spatial_axes` | `[1, 2]` | Spatial axes for rotation | -### Elastic Deformation (`AUGMENTOR.ELASTIC`) -Applies smooth elastic deformation to simulate tissue warping. +### Affine (`data.augmentation.affine`) +Random affine transforms (rotation, scaling, shearing). | Key | Default | Description | |-----|---------|-------------| -| `ENABLED` | `True` | Enable elastic deformation | -| `ALPHA` | `16.0` | Maximum pixel displacement | -| `SIGMA` | `4.0` | Gaussian filter standard deviation | -| `P` | `0.75` | Probability | +| `enabled` | `false` | Enable affine (only in `aug_strong`) | +| `prob` | `0.3` | Probability | +| `rotate_range` | `[0.1, 0.1, 0.1]` | Rotation range per axis | +| `scale_range` | `[0.05, 0.05, 0.05]` | Scale range per axis | +| `shear_range` | `[0.05, 0.05, 0.05]` | Shear range per axis | -### Grayscale Augmentation (`AUGMENTOR.GRAYSCALE`) -Randomly adjusts brightness and contrast. +### Elastic Deformation (`data.augmentation.elastic`) +Smooth elastic deformation to simulate tissue warping. | Key | Default | Description | |-----|---------|-------------| -| `ENABLED` | `True` | Enable grayscale augmentation | -| `P` | `0.75` | Probability | +| `enabled` | varies | Enable elastic deformation | +| `prob` | `0.3` | Probability | +| `sigma_range` | `[5.0, 8.0]` | Gaussian filter sigma range | +| `magnitude_range` | `[50.0, 150.0]` | Displacement magnitude range | -### Missing Parts (`AUGMENTOR.MISSINGPARTS`) -Simulates missing tissue regions (common artifact in EM data). +### Intensity (`data.augmentation.intensity`) +Adjusts brightness, contrast, and adds Gaussian noise. | Key | Default | Description | |-----|---------|-------------| -| `ENABLED` | `True` | Enable missing parts simulation | -| `ITER` | `64` | Number of missing region iterations | -| `P` | `0.9` | Probability | +| `enabled` | `true` | Enable intensity augmentation | +| `gaussian_noise_prob` | `0.3–0.5` | Probability of Gaussian noise | +| `gaussian_noise_std` | `0.05–0.1` | Noise standard deviation | +| `shift_intensity_prob` | `0.3–0.7` | Probability of brightness shift | +| `shift_intensity_offset` | `0.1–0.2` | Brightness shift magnitude | +| `contrast_prob` | `0.3–0.7` | Probability of contrast change | +| `contrast_range` | `[0.9, 1.1]` or `[0.7, 1.4]` | Contrast scaling range | -### Missing Section (`AUGMENTOR.MISSINGSECTION`) -Simulates entirely missing z-slices (another common EM artifact). +### Misalignment (`data.augmentation.misalignment`) +Simulates section-to-section misalignment (common EM artifact). | Key | Default | Description | |-----|---------|-------------| -| `ENABLED` | `True` | Enable missing section simulation | -| `NUM_SECTION` | `2` | Number of sections to remove | -| `P` | `0.5` | Probability | +| `enabled` | `true` | Enable misalignment | +| `prob` | `0.4–1.0` | Probability | +| `displacement` | `10–17` | Maximum pixel displacement | +| `rotate_ratio` | `0.0` | Fraction using rotation vs. translation | -### Misalignment (`AUGMENTOR.MISALIGNMENT`) -Simulates section-to-section misalignment (shift or rotation between z-slices). +### Missing Section (`data.augmentation.missing_section`) +Simulates entirely missing z-slices. | Key | Default | Description | |-----|---------|-------------| -| `ENABLED` | `True` | Enable misalignment simulation | -| `DISPLACEMENT` | `16` | Maximum pixel displacement | -| `ROTATE_RATIO` | `0.5` | Fraction of misalignments that use rotation instead of translation | -| `P` | `0.5` | Probability | +| `enabled` | `true` | Enable missing section | +| `prob` | `0.3–1.0` | Probability | +| `num_sections` | `2` or `[0, 5]` | Number of sections to remove | -### Motion Blur (`AUGMENTOR.MOTIONBLUR`) +### Motion Blur (`data.augmentation.motion_blur`) Simulates motion blur artifacts in EM sections. | Key | Default | Description | |-----|---------|-------------| -| `ENABLED` | `True` | Enable motion blur | -| `SECTIONS` | `2` | Number of sections to blur | -| `KERNEL_SIZE` | `11` | Blur kernel size | -| `P` | `0.5` | Probability | +| `enabled` | varies | Enable motion blur | +| `prob` | `0.3–1.0` | Probability | +| `sections` | `[1, 3]` | Number of sections to blur | +| `kernel_size` | `11` | Blur kernel size | -### CutBlur (`AUGMENTOR.CUTBLUR`) -Replaces a rectangular region with a downsampled-then-upsampled version, simulating resolution variation. +### Missing Parts (`data.augmentation.missing_parts`) +Simulates missing tissue regions. | Key | Default | Description | |-----|---------|-------------| -| `ENABLED` | `True` | Enable CutBlur | -| `LENGTH_RATIO` | `0.4` | Ratio of region size to volume size | -| `DOWN_RATIO_MIN` | `2.0` | Minimum downsampling factor | -| `DOWN_RATIO_MAX` | `8.0` | Maximum downsampling factor | -| `DOWNSAMPLE_Z` | `False` | Also downsample along z-axis | -| `P` | `0.5` | Probability | - -### CutNoise (`AUGMENTOR.CUTNOISE`) +| `enabled` | varies | Enable missing parts | +| `prob` | `0.2` | Probability | +| `hole_range` | `[0.1, 0.25]` | Hole size range relative to volume | + +### Cut Noise (`data.augmentation.cut_noise`) Adds noise to a rectangular region. | Key | Default | Description | |-----|---------|-------------| -| `ENABLED` | `True` | Enable CutNoise | -| `LENGTH_RATIO` | `0.4` | Ratio of region size | -| `SCALE` | `0.3` | Noise scale | -| `P` | `0.75` | Probability | +| `enabled` | varies | Enable cut noise | +| `prob` | `0.3` | Probability | +| `length_ratio` | `[0.1, 0.25]` | Region size ratio | +| `noise_scale` | `[0.05, 0.15]` | Noise scale range | -### CopyPaste (`AUGMENTOR.COPYPASTE`) -Copy-pastes object instances for data augmentation (disabled by default). +### Cut Blur (`data.augmentation.cut_blur`) +Replaces a region with downsampled-then-upsampled version. | Key | Default | Description | |-----|---------|-------------| -| `ENABLED` | `False` | Enable CopyPaste (disabled by default) | -| `AUG_THRES` | `0.7` | Augmentation threshold | -| `P` | `0.8` | Probability | +| `enabled` | varies | Enable cut blur | +| `prob` | `0.8` | Probability | +| `length_ratio` | `[0.3, 0.5]` | Region size ratio | +| `down_ratio_range` | `[2.0, 8.0]` | Downsampling factor range | +| `downsample_z` | `false` | Also downsample along z-axis | -## Global Augmentation Options +### Copy-Paste (`data.augmentation.copy_paste`) +Copy-pastes object instances for data augmentation. | Key | Default | Description | |-----|---------|-------------| -| `AUGMENTOR.ENABLED` | `True` | Master switch for all augmentations | -| `AUGMENTOR.SMOOTH` | `False` | Apply Gaussian smoothing to label masks after augmentation. WARNING: can erase thin structures. | -| `AUGMENTOR.ADDITIONAL_TARGETS_NAME` | `['label']` | Names of additional targets to augment alongside the image | -| `AUGMENTOR.ADDITIONAL_TARGETS_TYPE` | `['mask']` | Type of each additional target (`'mask'` uses nearest-neighbor interpolation) | +| `enabled` | `false` | Enable copy-paste (only in `aug_instance`) | +| `prob` | `0.6` | Probability | +| `max_obj_ratio` | `0.7` | Maximum object area ratio | -## Recommended Settings by Data Type +### Mixup (`data.augmentation.mixup`) +Blends two training samples together. -### Anisotropic EM data (e.g., SNEMI, CREMI) -- `AUGMENTOR.FLIP.DO_ZTRANS: 0` (do not flip across z) -- All other augmentations enabled by default -- Consider disabling `CUTNOISE` for clean datasets - -### Isotropic EM data (e.g., Lucchi) -- `AUGMENTOR.FLIP.DO_ZTRANS: 1` (enable z-axis flips) -- `AUGMENTOR.CUTBLUR.DOWNSAMPLE_Z: True` +| Key | Default | Description | +|-----|---------|-------------| +| `enabled` | `false` | Enable mixup (only in `aug_instance`) | +| `prob` | `0.3` | Probability | +| `alpha_range` | `[0.7, 0.9]` | Blending alpha range | -### 2D datasets (e.g., Cellpose) -- Disable EM-specific augmentations: `ELASTIC.ENABLED: False`, `RESCALE.ENABLED: False`, `MISSINGPARTS.ENABLED: False` -- Keep flip, rotation, and grayscale augmentations +## Special Features -### Sparse labels (e.g., Scutoid, NucMM) -- Enable reject sampling: `DATASET.REJECT_SAMPLING.SIZE_THRES: 1000` -- Consider disabling `CUTNOISE` to avoid corrupting sparse regions +### Defect Mutex (`data.augmentation.defect_mutex`) +When `true`, only one defect augmentation (misalignment, missing_section, or motion_blur) is applied per sample. Used in the `aug_em_neuron` profile to match the DeepEM augmentation strategy. -## Per-Augmentation Skipping +## Recommended Profile by Data Type -Each augmentation has a `SKIP` parameter (list of sample keys to skip). This allows skipping certain augmentations for specific data channels. Default is an empty list (no skipping). +- **Anisotropic EM (SNEMI, CREMI)**: `aug_standard` or `aug_em_neuron` +- **Isotropic EM (Lucchi)**: `aug_standard` with `flip.spatial_axis: [0, 1, 2]` +- **Small datasets / overfitting**: `aug_strong` +- **Instance segmentation**: `aug_instance` +- **Super-resolution**: `aug_superres` +- **Quick experiments / fine-tuning**: `aug_light` diff --git a/server_api/chatbot/file_summaries/PyTC-Configs.md b/server_api/chatbot/file_summaries/PyTC-Configs.md index 59e65fff..dd7b7db9 100644 --- a/server_api/chatbot/file_summaries/PyTC-Configs.md +++ b/server_api/chatbot/file_summaries/PyTC-Configs.md @@ -1,137 +1,103 @@ # PyTC Bundled Configuration Files -This document describes every YAML configuration file bundled with PyTorch Connectomics. Users should start from the closest matching config and modify it rather than writing one from scratch. +This document describes the YAML configuration files bundled with PyTorch Connectomics. All configs are in the `tutorials/` directory and use Hydra/OmegaConf format with a profile-based system. Users should start from the closest matching config and modify it rather than writing one from scratch. -## Lucchi (Mitochondria Segmentation) +## Config Architecture -**Dataset**: Lucchi — isotropic EM volume of mitochondria in hippocampus. -**Task**: Binary semantic segmentation (mitochondria vs. background). +Configs use a profile-based inheritance system. Each tutorial config references base profiles from `tutorials/bases/`: +- `arch_profiles.yaml` — Model architecture presets +- `optimizer_profiles.yaml` — Optimizer and scheduler presets +- `augmentation_profiles.yaml` — Data augmentation presets +- `pipeline_profiles.yaml` — End-to-end pipeline presets (loss, labels, decoding) +- `dataloader_profiles.yaml` — Dataloader presets +- `system_profiles.yaml` — GPU/CPU resource presets -| Config | Architecture | Key Settings | -|--------|-------------|-------------| -| `configs/Lucchi-Mitochondria.yaml` | `unet_3d` / `residual_se` | Input: `[112,112,112]`, LR: `0.04`, Loss: BCE+Dice, Isotropic data, 100K iterations | - -## SNEMI (Neuron Segmentation) +## Mitochondria Segmentation -**Dataset**: SNEMI3D — anisotropic EM serial section images of mouse cortex neurons. -**Task**: Instance segmentation via affinity maps. +### Lucchi++ (Isotropic EM) | Config | Architecture | Key Settings | |--------|-------------|-------------| -| `configs/SNEMI/SNEMI-Base.yaml` | `unet_3d` / `residual` | Affinity target (`"2"`), 3-channel output, GroupNorm, 150K iters | -| `configs/SNEMI/SNEMI-Affinity-UNet.yaml` | `unet_3d` / `residual_se` | Extends SNEMI-Base with SE attention blocks | -| `configs/SNEMI/SNEMI-Affinity-UNet-2x.yaml` | `unet_3d` | Larger model variant | -| `configs/SNEMI/SNEMI-Affinity-UNet-LR.yaml` | `unet_3d` | Learning rate experiment variant | -| `configs/SNEMI/SNEMI-Affinity-UNet-MER.yaml` | `unet_3d` | Model ensemble with regularization | -| `configs/SNEMI/SNEMI-Affinity-Contour-UNet-2x.yaml` | `unet_3d` | Multi-task: affinity + contour prediction | -| `configs/SNEMI/SNEMI-Affinity-ResNet.yaml` | `fpn_3d` / ResNet backbone | FPN architecture with ResNet backbone | -| `configs/SNEMI/SNEMI-Affinity-EffNet.yaml` | `fpn_3d` / EfficientNet backbone | FPN with EfficientNet backbone | -| `configs/SNEMI/SNEMI-Affinity-SwinUNETR.yaml` | `swinunetr` | Swin Transformer-based segmentation | -| `configs/SNEMI/SNEMI-Affinity-UNETR.yaml` | `unetr` | Vision Transformer-based segmentation | -| `configs/SNEMI/SNEMI-Base_multiGPU.yaml` | `unet_3d` | Multi-GPU (4 GPUs) distributed training variant | - -## CREMI (Synapse Detection) - -**Dataset**: CREMI — anisotropic EM volumes (A, B, C) with synaptic cleft annotations. -**Task**: Synapse detection / binary segmentation with reject sampling for sparse labels. +| `tutorials/mito_lucchi++.yaml` | `monai_unet` | Input: `[112,112,112]`, binary pipeline, WarmupCosineLR, batch_size=8, 100 epochs, isotropic data | + +### MitoEM (Large-Scale) | Config | Architecture | Key Settings | |--------|-------------|-------------| -| `configs/CREMI/CREMI-Base.yaml` | `unet_3d` / `residual` | SyncBN, Reject sampling (threshold 1000), 150K iters | -| `configs/CREMI/CREMI-Base_multiGPU.yaml` | `unet_3d` | Multi-GPU distributed variant | -| `configs/CREMI/CREMI-Foreground-UNet.yaml` | `unet_3d` | Foreground prediction with UNet | -| `configs/CREMI/CREMI-Foreground-DT-UNet.yaml` | `unet_3d` | Multi-task: foreground + distance transform | -| `configs/CREMI/CREMI-Foreground-DT-Regu-UNet.yaml` | `unet_3d` | Foreground + DT with regularization | -| `configs/CREMI/CREMI-Foreground-FPN-ResNet.yaml` | `fpn_3d` / ResNet | FPN with ResNet for foreground detection | -| `configs/CREMI/CREMI-Foreground-FPN-RepVGG.yaml` | `fpn_3d` / RepVGG | FPN with RepVGG backbone | - -## MitoEM (Large-Scale Mitochondria) +| `tutorials/mito_mitoEM_common.yaml` | shared base | Common settings for MitoEM configs | +| `tutorials/mito_mitoEM_H.yaml` | variant | MitoEM-H dataset variant | +| `tutorials/mito_mitoEM_R.yaml` | variant | MitoEM-R dataset variant | +| `tutorials/mito_mitoEM_HR.yaml` | variant | Combined H+R dataset variant | -**Dataset**: MitoEM — large-scale EM dataset for mitochondria instance segmentation. -**Task**: Binary/instance segmentation with tile-based training for large volumes. +### BetaSeg (Mitochondria) | Config | Architecture | Key Settings | |--------|-------------|-------------| -| `configs/MitoEM/MitoEM-R-Base.yaml` | `unet_plus_3d` / `residual_se` | Chunked training (4×8×8), 4 GPUs, SyncBN, 150K iters | -| `configs/MitoEM/MitoEM-R-A.yaml` | variant | Architecture A experiment | -| `configs/MitoEM/MitoEM-R-AC.yaml` | variant | Architecture AC experiment | -| `configs/MitoEM/MitoEM-R-BC.yaml` | variant | Boundary + Contour multi-task | -| `configs/MitoEM/MitoEM-R-BCD.yaml` | variant | Boundary + Contour + Distance Transform multi-task | +| `tutorials/mito_betaseg.yaml` | varies | BetaSeg mitochondria segmentation | -## NucMM (Nucleus Segmentation) - -**Dataset**: NucMM — 3D nuclei segmentation in mouse brain and zebrafish. -**Task**: Instance segmentation of cell nuclei. +### MitoLab | Config | Architecture | Key Settings | |--------|-------------|-------------| -| `configs/NucMM/NucMM-Mouse-Base.yaml` | `unet_3d` / `residual_se` | 4 GPUs, GroupNorm, Reject sampling, 100K iters | -| `configs/NucMM/NucMM-Mouse-UNet-BC.yaml` | `unet_3d` | Boundary + Contour multi-task | -| `configs/NucMM/NucMM-Mouse-UNet-BCD.yaml` | `unet_3d` | Boundary + Contour + Distance Transform | -| `configs/NucMM/NucMM-Zebrafish-Base.yaml` | `unet_3d` / `residual_se` | Zebrafish data, similar setup | -| `configs/NucMM/NucMM-Zebrafish-UNet-BC.yaml` | `unet_3d` | Zebrafish BC variant | -| `configs/NucMM/NucMM-Zebrafish-UNet-BCD.yaml` | `unet_3d` | Zebrafish BCD variant | +| `tutorials/mito_mitolab.yaml` | varies | MitoLab mitochondria segmentation | -## JWR15 (Neuron and Synapse Segmentation) +## Neuron Segmentation -**Dataset**: JWR15 — EM dataset with both neuron and synapse annotations. +### SNEMI3D (Anisotropic EM) -### Neuron configs | Config | Architecture | Key Settings | |--------|-------------|-------------| -| `configs/JWR15/neuron/JWR15-Neuron-Base.yaml` | `unet_3d` | Neuron segmentation baseline | -| `configs/JWR15/neuron/JWR15-Neuron-Affinity-UNet-MER.yaml` | `unet_3d` | Affinity-based with model ensemble | +| `tutorials/neuron_snemi.yaml` | `rsunet` | 12-channel affinity maps, anisotropic, WarmupCosineLR, 100 epochs, ABISS decoding with parameter tuning | + +### NISB (Neuron Instance Segmentation Benchmark) -### Synapse configs | Config | Architecture | Key Settings | |--------|-------------|-------------| -| `configs/JWR15/synapse/JWR15-Synapse-Base.yaml` | `unet_3d` | Synapse detection baseline | -| `configs/JWR15/synapse/JWR15-Synapse-BCE.yaml` | `unet_3d` | BCE loss variant | -| `configs/JWR15/synapse/JWR15-Synapse-BCE-DICE.yaml` | `unet_3d` | BCE + Dice loss | -| `configs/JWR15/synapse/JWR15-Synapse-BCE-DICE-Regu.yaml` | `unet_3d` | BCE + Dice + Regularization | -| `configs/JWR15/synapse/JWR15-Synapse-Semantic-CE.yaml` | `unet_3d` | Multi-class semantic variant | +| `tutorials/neuron_nisb_common.yaml` | shared base | Common NISB settings | +| `tutorials/neuron_nisb_40nm_common.yaml` | shared base | 40nm resolution common settings | +| `tutorials/neuron_nisb_40nm_base.yaml` | varies | 40nm baseline | +| `tutorials/neuron_nisb_40nm_liconn.yaml` | varies | 40nm LiConn variant | +| `tutorials/neuron_nisb_9nm_common.yaml` | shared base | 9nm resolution common settings | +| `tutorials/neuron_nisb_9nm_base.yaml` | varies | 9nm baseline | +| `tutorials/neuron_nisb_9nm_liconn.yaml` | varies | 9nm LiConn variant | -## Cellpose (2D Cell Segmentation) +## Synapse Detection -**Dataset**: Cellpose — 2D microscopy images of cells. -**Task**: Instance segmentation using flow-field prediction. +### CREMI (Synaptic Cleft) | Config | Architecture | Key Settings | |--------|-------------|-------------| -| `configs/Cellpose/Cellpose-Base.yaml` | `unet_2d` / `residual_se` | 2D mode, Multi-task (mask + flow), LeakyReLU, Batch=32, 50K iters, Elastic/Rescale/MissingParts disabled | +| `tutorials/syn_cremi.yaml` | `rsunet` | Binary synapse cleft detection, BCE+Dice loss, reject sampling, 150K steps, CREMI evaluation metrics | -## Scutoid (Epithelial Cell Segmentation) - -**Dataset**: Scutoid — 3D EM volumes of epithelial cells. -**Task**: Instance segmentation with scaled data. +## Nucleus Segmentation | Config | Architecture | Key Settings | |--------|-------------|-------------| -| `configs/Scutoid/Scutoid-Base.yaml` | `unet_3d` / `residual_se` | 4 GPUs, Data scale `[1.0, 0.5, 0.5]`, Reject sampling, CutNoise disabled | -| `configs/Scutoid/Scutoid-UNet-BCD.yaml` | `unet_3d` | Boundary + Contour + Distance Transform | - -## Zebrafinch (Neuron Segmentation) +| `tutorials/nuc_nucmm-z.yaml` | varies | NucMM zebrafish nucleus segmentation | -**Dataset**: Zebrafinch — anisotropic EM of songbird brain neurons. -**Task**: Instance segmentation via affinity maps. +## Other Datasets | Config | Architecture | Key Settings | |--------|-------------|-------------| -| `configs/Zebrafinch/Zebrafinch-Base.yaml` | `unet_3d` | Neuron segmentation baseline | -| `configs/Zebrafinch/Zebrafinch-Affinity-UNet.yaml` | `unet_3d` / `residual_se` | Affinity-based with SE blocks | -| `configs/Zebrafinch/Zebrafinch-Affinity-UNet-MER.yaml` | `unet_3d` | Affinity with model ensemble regularization | +| `tutorials/vesicle_xm.yaml` | varies | Vesicle segmentation | +| `tutorials/fiber_linghu26.yaml` | varies | Fiber segmentation | +| `tutorials/minimal.yaml` | varies | Minimal config for demo/testing | -## Generic / Template Configs +### Misc Configs -| Config | Task | Description | -|--------|------|-------------| -| `configs/Multiclass-Semantic-Seg.yaml` | Multi-class segmentation | Template for N-class semantic segmentation. Uses `TARGET_OPT: ["9-12"]` and `WeightedCE` loss. 12 output channels. | -| `configs/Distance-Transform-Quantized.yaml` | Distance transform | Template for quantized distance transform prediction. Uses `TARGET_OPT: ["5"]` and `WeightedCE` loss. | +| Config | Description | +|--------|-------------| +| `tutorials/misc/mito_2dsem_seg.yaml` | 2D SEM mitochondria segmentation | +| `tutorials/misc/tsai_axon.yaml` | Axon segmentation | +| `tutorials/misc/worm2d.yaml` | 2D worm segmentation | +| `tutorials/misc/zebrafish_neurons.yaml` | Zebrafish neuron segmentation | +| `tutorials/misc/hydra-lv.yaml` | Hydra large-volume config | +| `tutorials/misc/hydra-lv-finetune.yaml` | Hydra large-volume fine-tuning | ## How to Choose a Config -1. **Identify your task**: binary segmentation, instance segmentation, multi-class, or cell segmentation. -2. **Find the closest dataset**: Pick the bundled config whose dataset most resembles yours (isotropic vs. anisotropic, resolution, structure type). -3. **Start from Base**: Use the `-Base.yaml` config as your starting point. -4. **Override paths**: Change `DATASET.INPUT_PATH`, `DATASET.IMAGE_NAME`, `DATASET.LABEL_NAME`, and `DATASET.OUTPUT_PATH` to point to your data. -5. **Adjust as needed**: Modify learning rate, batch size, iteration count, and model architecture based on your GPU resources and dataset size. +1. **Identify your task**: binary segmentation (mitochondria, synapse), instance segmentation (neurons, nuclei), or other. +2. **Find the closest dataset**: Pick the tutorial config whose dataset most resembles yours (isotropic vs. anisotropic, resolution, structure type). +3. **Override data paths**: Set `data.train.image`, `data.train.label`, and `test.data.test.image` to point to your data. +4. **Choose profiles**: Select appropriate `model.arch.profile`, `optimization.profile`, and `data.augmentation.profile`. +5. **Adjust as needed**: Override learning rate, batch size, epochs, etc. using Hydra command-line overrides. diff --git a/server_api/chatbot/file_summaries/PyTC-Evaluation.md b/server_api/chatbot/file_summaries/PyTC-Evaluation.md index f94a7e1a..82c52e8a 100644 --- a/server_api/chatbot/file_summaries/PyTC-Evaluation.md +++ b/server_api/chatbot/file_summaries/PyTC-Evaluation.md @@ -4,12 +4,12 @@ This document explains how to evaluate segmentation results produced by PyTorch ## Running Inference First -Before evaluation, you must run inference to produce predictions. Use the `--inference` flag: +Before evaluation, you must run inference to produce predictions. Use `--mode test`: ``` -python scripts/main.py --config-file --inference --checkpoint +python scripts/main.py --config --mode test --checkpoint ``` -The predictions are saved to `INFERENCE.OUTPUT_PATH` (usually as HDF5 files). +The predictions are saved to the output directory (usually as HDF5 files). Evaluation can be run automatically during inference by setting `inference.evaluation.enabled=true` and specifying `inference.evaluation.metrics` in the config. ## Available Evaluation Metrics @@ -114,6 +114,7 @@ Computes mean distance-based false positive and false negative statistics. ## Notes -- PyTC does NOT automatically compute evaluation metrics after inference. You must write a separate script. -- The `--inference` flag only generates predictions; it does not compare them to ground truth. +- PyTC can automatically compute evaluation metrics during inference when `inference.evaluation.enabled=true` is set in the config. +- Available built-in metrics include `jaccard`, `adapted_rand`, `cremi_distance`, and more. +- For standalone evaluation, you can also write a separate script using `connectomics.utils.evaluate`. - For large volumes, load data in chunks to avoid memory issues. diff --git a/server_api/chatbot/file_summaries/PyTC-Inference.md b/server_api/chatbot/file_summaries/PyTC-Inference.md index c89f3d1a..7012a8b9 100644 --- a/server_api/chatbot/file_summaries/PyTC-Inference.md +++ b/server_api/chatbot/file_summaries/PyTC-Inference.md @@ -5,85 +5,115 @@ This document explains how to run inference (prediction) with a trained PyTorch ## Inference Command ``` -python scripts/main.py --config-file --inference --checkpoint [OVERRIDES] +python scripts/main.py --config --mode test --checkpoint [key=value overrides] ``` Example: ``` -python scripts/main.py --config-file configs/Lucchi-Mitochondria.yaml --inference --checkpoint outputs/Lucchi_UNet/checkpoint_100000.pth.tar +python scripts/main.py --config tutorials/mito_lucchi++.yaml --mode test --checkpoint outputs/mito_lucchi++/20241124_203930/checkpoints/last.ckpt ``` -The `--inference` flag switches the trainer to test mode. A `--checkpoint` path is required to load the trained model weights. +The `--mode test` flag switches the trainer to test/inference mode. A `--checkpoint` path is required to load the trained model weights. ## How Inference Works -1. The config file is loaded and inference-specific settings override the training defaults (input path, image name, output path, pad size, input/output size). +1. The config file is loaded and test-stage settings override the training defaults (test data paths, output paths). 2. The model is built from the config and checkpoint weights are loaded. -3. The test volume is loaded and divided into overlapping patches based on `INFERENCE.INPUT_SIZE` and `INFERENCE.STRIDE`. +3. The test volume is loaded and divided into overlapping patches via sliding window inference. 4. Each patch is fed through the model. If test-time augmentation is enabled, multiple augmented copies are predicted and aggregated. 5. Overlapping predictions are blended together (Gaussian blending by default). -6. The final prediction is saved to the output directory. +6. Postprocessing and decoding are applied if configured. +7. The final prediction is saved to the output directory. -## INFERENCE Configuration Options +## Inference Configuration Options + +### Sliding Window Settings + +| Key | Description | Default | +|-----|-------------|---------| +| `inference.batch_size` | Batch size for inference | `1` | +| `inference.sliding_window.window_size` | Patch size `[z, y, x]` for sliding window inference | matches model input_size | +| `inference.sliding_window.sw_batch_size` | Number of patches per sliding window batch | `4` | +| `inference.sliding_window.overlap` | Overlap ratio between adjacent patches (0.0–1.0) | `0.5` | +| `inference.sliding_window.blending` | Blending mode: `'gaussian'` or `'constant'` | `'gaussian'` | +| `inference.sliding_window.sigma_scale` | Sigma scale for Gaussian blending | `0.25` | +| `inference.sliding_window.padding_mode` | Padding mode: `'reflect'`, `'constant'`, etc. | `'reflect'` | + +### Test-Time Augmentation (TTA) + +| Key | Description | Default | +|-----|-------------|---------| +| `inference.test_time_augmentation.enabled` | Enable TTA | `true` or `false` (config-dependent) | +| `inference.test_time_augmentation.flip_axes` | Flip axes for TTA: `'all'` or list of axis combinations | `'all'` | +| `inference.test_time_augmentation.rotation90_axes` | Rotation axes for TTA: `'all'` or list | None | +| `inference.test_time_augmentation.ensemble_mode` | Aggregation mode: `'mean'` or `'min'` | `'mean'` | +| `inference.test_time_augmentation.channel_activations` | Per-channel activation functions | varies | +| `inference.test_time_augmentation.distributed_sharding` | Distribute TTA across GPUs | `false` | + +### Postprocessing and Decoding | Key | Description | Default | |-----|-------------|---------| -| `INFERENCE.INPUT_SIZE` | Patch size `[z, y, x]` for inference (overrides MODEL.INPUT_SIZE) | None (uses model size) | -| `INFERENCE.OUTPUT_SIZE` | Output patch size `[z, y, x]` (overrides MODEL.OUTPUT_SIZE) | None (uses model size) | -| `INFERENCE.IMAGE_NAME` | Test image filename(s) relative to input path | None | -| `INFERENCE.INPUT_PATH` | Override DATASET.INPUT_PATH for inference | None | -| `INFERENCE.OUTPUT_PATH` | Directory where predictions are saved | `""` | -| `INFERENCE.OUTPUT_NAME` | Output filename (e.g., `result.h5`, `pred`) | `'result.h5'` | -| `INFERENCE.PAD_SIZE` | Padding `[z, y, x]` for inference volumes | None (uses dataset pad) | -| `INFERENCE.STRIDE` | Stride `[z, y, x]` between patches. Smaller stride = more overlap = smoother results but slower | `[4, 128, 128]` | -| `INFERENCE.BLENDING` | Blending mode for overlapping patches: `'gaussian'` or `'bump'` | `'gaussian'` | -| `INFERENCE.OUTPUT_ACT` | Activation function(s) for output: `'sigmoid'`, `'softmax'`, `'tanh'` | `['sigmoid']` | -| `INFERENCE.AUG_MODE` | Test-time augmentation aggregation: `'mean'` or `'min'` | `'mean'` | -| `INFERENCE.AUG_NUM` | Number of augmented copies for TTA. Set to None to disable TTA | None | -| `INFERENCE.SAMPLES_PER_BATCH` | Batch size for inference (patches per GPU) | `4` | -| `INFERENCE.DO_EVAL` | Run with `model.eval()` (affects batchnorm/dropout) | `True` | -| `INFERENCE.DO_SINGLY` | Process volumes one at a time (for multi-volume datasets) | `False` | -| `INFERENCE.DO_SINGLY_START_INDEX` | Starting index when DO_SINGLY is True | `0` | -| `INFERENCE.DO_SINGLY_STEP` | Step between volumes in DO_SINGLY mode | `1` | -| `INFERENCE.DATA_SCALE` | Override DATASET.DATA_SCALE for inference | None | -| `INFERENCE.OUTPUT_SCALE` | Rescale predictions after model output `[z, y, x]` | `[1.0, 1.0, 1.0]` | - -## Inference Output - -- By default, predictions are saved as HDF5 files (`.h5`) in the `INFERENCE.OUTPUT_PATH` directory. -- For chunked or DO_SINGLY inference, the output filename is a stem (e.g., `result`) and files are saved per volume. -- For multi-class segmentation (`TARGET_OPT: ["9-N"]`), the output activation is automatically set to softmax if not specified. +| `inference.postprocessing.enabled` | Enable postprocessing | `false` | +| `inference.postprocessing.crop_pad` | Crop padding from predictions | None | +| `inference.decoding` | List of decoding profiles (e.g., ABISS watershed) | None | + +### Save Prediction + +| Key | Description | Default | +|-----|-------------|---------| +| `inference.save_prediction.enabled` | Save prediction to disk | `true` | +| `inference.save_prediction.intensity_scale` | Scale intensity values (-1 = disabled) | `-1` | +| `inference.save_prediction.intensity_dtype` | Output data type | `float32` | +| `inference.save_prediction.output_formats` | Output file format(s) | `[h5]` | + +### Evaluation + +| Key | Description | Default | +|-----|-------------|---------| +| `inference.evaluation.enabled` | Run evaluation metrics after inference | `true` or `false` | +| `inference.evaluation.metrics` | List of metrics: `jaccard`, `adapted_rand`, `cremi_distance`, etc. | varies | + +### Test Data + +| Key | Description | Example | +|-----|-------------|---------| +| `test.data.test.image` | Test image file path(s) | `datasets/lucchi++/test_im.h5` | +| `test.data.test.label` | Test label file path(s) (for evaluation) | `datasets/lucchi++/test_mito.h5` | +| `test.data.test.resolution` | Test data voxel resolution | `[5, 5, 5]` | +| `test.output_path` | Directory where predictions are saved | `outputs//results/` | ## Test-Time Augmentation (TTA) -TTA applies geometric augmentations (flips, rotations) to the input, runs the model on each augmented copy, and aggregates the predictions. +TTA applies geometric augmentations (flips, 90° rotations) to the input, runs the model on each augmented copy, and aggregates predictions. -- `INFERENCE.AUG_NUM`: Number of augmented copies. Common values: `4`, `8`, `16`. Set to `None` to disable. -- `INFERENCE.AUG_MODE`: - - `'mean'`: Average all predictions (recommended for most tasks). - - `'min'`: Take the minimum prediction (useful for conservative boundary detection). +- **`flip_axes: all`**: All combinations of axis flips (2^D variants for D spatial dimensions). +- **`ensemble_mode: mean`**: Average all predictions (recommended for most tasks). +- **`ensemble_mode: min`**: Take the minimum prediction (useful for conservative boundary detection in affinity maps). -## Inference Stride and Blending +## Sliding Window and Blending -The inference stride controls how much overlap there is between adjacent patches: -- **Smaller stride** = more overlap = smoother results, but slower. -- **Larger stride** = less overlap = faster, but may show tile boundary artifacts. +The sliding window controls how the volume is split into overlapping patches: +- **Higher overlap** (e.g., `0.5`) = smoother results, but slower. +- **Lower overlap** (e.g., `0.25`) = faster, but may show tile boundary artifacts. - `'gaussian'` blending weights the center of each patch more heavily, producing smooth transitions. -## Large-Scale (Chunked) Inference +## Multi-Volume and Sharded Inference -For datasets too large to fit in memory, use tiled inference: -```yaml -INFERENCE: - DO_CHUNK_TITLE: 1 - DATA_CHUNK_NUM: [4, 8, 8] +For multi-volume datasets, test data paths can be lists. For distributed inference across multiple machines: +``` +python scripts/main.py --config --mode test --checkpoint --shard-id 0 --num-shards 4 +python scripts/main.py --config --mode test --checkpoint --shard-id 1 --num-shards 4 ``` -The volume is split into chunks that are processed sequentially. + +## Cache-Aware Inference + +PyTC automatically detects cached prediction files and skips redundant inference. This is useful when re-running postprocessing or evaluation without re-computing predictions. ## Tips -- Use the same config file for inference that was used for training. Only override inference-specific settings. -- Increase `INFERENCE.SAMPLES_PER_BATCH` to speed up inference if GPU memory allows. -- Increase `INFERENCE.AUG_NUM` (e.g., 8 or 16) for better accuracy at the cost of speed. -- For multi-volume datasets (like CREMI with volumes A, B, C), set `INFERENCE.DO_SINGLY: True`. -- The output directory must exist or the script will create it automatically. +- Use the same config file for inference that was used for training. Only override test-specific settings. +- Increase `inference.sliding_window.sw_batch_size` to speed up inference if GPU memory allows. +- Enable TTA (`inference.test_time_augmentation.enabled=true`) for better accuracy at the cost of speed. +- Checkpoint files use `.ckpt` extension (PyTorch Lightning format). +- The output directory is automatically created based on the checkpoint path. diff --git a/server_api/chatbot/file_summaries/PyTC-Models.md b/server_api/chatbot/file_summaries/PyTC-Models.md index 2f101760..ac026627 100644 --- a/server_api/chatbot/file_summaries/PyTC-Models.md +++ b/server_api/chatbot/file_summaries/PyTC-Models.md @@ -1,153 +1,132 @@ # PyTC Model Zoo -This document describes all model architectures, block types, backbones, loss functions, and target options available in PyTorch Connectomics. +This document describes all model architectures, loss functions, label transforms, and pipeline profiles available in PyTorch Connectomics. ## Model Architectures -Set `MODEL.ARCHITECTURE` to one of: +Set `model.arch.profile` to use a predefined architecture, or set `model.arch.type` directly: -| Architecture | Key | Description | Best For | -|-------------|-----|-------------|----------| -| 3D U-Net | `unet_3d` | Standard 3D encoder-decoder with skip connections. Most commonly used. | General 3D segmentation | -| 2D U-Net | `unet_2d` | 2D variant for slice-by-slice processing. Requires `DATASET.DO_2D: True`. | 2D cell segmentation (e.g., Cellpose) | -| 3D U-Net++ | `unet_plus_3d` | Nested U-Net with dense skip connections. | Large-scale datasets (e.g., MitoEM) | -| 2D U-Net++ | `unet_plus_2d` | 2D variant of U-Net++. | 2D tasks with dense skip connections | -| 3D FPN | `fpn_3d` | Feature Pyramid Network. Requires a backbone (`MODEL.BACKBONE`). | Multi-scale feature extraction | -| DeepLabV3 | `deeplabv3a` / `deeplabv3b` / `deeplabv3c` | 2D DeepLab with atrous convolutions. Three variants with different dilation rates. | 2D semantic segmentation | -| UNETR | `unetr` | Vision Transformer encoder with CNN decoder. | Transformer-based 3D segmentation | -| Swin UNETR | `swinunetr` | Shifted-window Transformer encoder (MONAI-based). Supports v2. | State-of-the-art transformer segmentation | +| Architecture | Profile / Type | Description | Best For | +|-------------|----------------|-------------|----------| +| MONAI UNet | `monai_unet` | MONAI-based 3D UNet with configurable filters and dropout | General 3D segmentation (recommended for beginners) | +| Residual Symmetric UNet | `rsunet` | Custom RSUNet with group norm, ELU activation, anisotropic downsampling | Anisotropic EM data (SNEMI, CREMI) | +| MedNeXt Small | `mednext_s` | MedNeXt architecture, size S | Efficient segmentation | +| MedNeXt Base | `mednext_b` | MedNeXt architecture, size B | Balanced performance | +| MedNeXt Medium | `mednext_m` | MedNeXt architecture, size M | Higher capacity | +| MedNeXt Large | `mednext_l` | MedNeXt architecture, size L | Maximum capacity | -## Block Types +### Architecture Profile Details -Set `MODEL.BLOCK_TYPE` to one of: - -| Block Type | Key | Available In | Description | -|-----------|-----|-------------|-------------| -| Residual | `residual` | unet_3d, unet_2d, fpn_3d | Standard residual block with two convolutions and a skip connection | -| Residual + SE | `residual_se` | unet_3d, unet_2d, fpn_3d | Residual block with Squeeze-and-Excitation channel attention. Most popular choice. | -| Residual PA | `residual_pa` | unet_3d | Pre-activation residual block (BN → ReLU → Conv ordering) | -| Residual PA + SE | `residual_pa_se` | unet_3d | Pre-activation block with Squeeze-and-Excitation | - -## Backbones (for FPN and DeepLab) - -Set `MODEL.BACKBONE` to one of: +**MONAI UNet (`monai_unet`)**: +```yaml +model: + arch: + profile: monai_unet + monai: + filters: [32, 64, 128, 256] + dropout: 0.1 +``` -| Backbone | Key | Description | -|----------|-----|-------------| -| ResNet | `resnet` | 3D ResNet backbone. Default for FPN. | -| RepVGG | `repvgg` | Re-parameterizable VGG. Supports deploy mode conversion at inference time. | -| BotNet | `botnet` | Bottleneck Transformer backbone. Requires `fmap_size` parameter. | -| EfficientNet | `effnet` | 3D EfficientNet with inverted residual blocks. | +**RSUNet (`rsunet`)**: +```yaml +model: + arch: + profile: rsunet + rsunet: + width: [18, 36, 48, 64, 80] + norm: group + activation: elu + num_groups: 4 + down_factors: [[1, 2, 2], [1, 2, 2], [1, 2, 2], [1, 2, 2]] + depth_2d: 1 + kernel_2d: [1, 3, 3] +``` ## Model Configuration Options -| Key | Description | Default | +| Key | Description | Example | |-----|-------------|---------| -| `MODEL.FILTERS` | Number of filters at each encoder stage | `[28, 36, 48, 64, 80]` | -| `MODEL.BLOCKS` | Number of residual blocks at each stage | `[2, 2, 2, 2]` | -| `MODEL.IN_PLANES` | Number of input channels (1 for grayscale EM) | `1` | -| `MODEL.OUT_PLANES` | Number of output channels | `1` | -| `MODEL.INPUT_SIZE` | Model input patch size `[z, y, x]` | `[8, 256, 256]` | -| `MODEL.OUTPUT_SIZE` | Model output patch size `[z, y, x]` | `[8, 256, 256]` | -| `MODEL.ISOTROPY` | Per-stage isotropy flags for anisotropic data | `[False, False, False, True, True]` | -| `MODEL.ATTENTION` | Attention mechanism: `'squeeze_excitation'` or None | `'squeeze_excitation'` | -| `MODEL.PAD_MODE` | Convolution padding: `'zeros'`, `'circular'`, `'reflect'`, `'replicate'` | `'replicate'` | -| `MODEL.NORM_MODE` | Normalization: `'bn'` (BatchNorm), `'sync_bn'`, `'in'` (Instance), `'gn'` (Group), `'none'` | `'bn'` | -| `MODEL.ACT_MODE` | Activation: `'relu'`, `'elu'`, `'leaky'` (leaky_relu) | `'elu'` | -| `MODEL.POOLING_LAYER` | Use pooling for downsampling (True) or strided conv (False) | `False` | -| `MODEL.MIXED_PRECESION` | Mixed-precision training (DDP only) | `False` | -| `MODEL.EMBEDDING` | Enable embedding head | `1` | -| `MODEL.HEAD_DEPTH` | Depth of final decoder head | `1` | -| `MODEL.PRE_MODEL` | Path to pretrained model for fine-tuning | `''` | - -### Swin UNETR-Specific Options - -| Key | Default | Description | -|-----|---------|-------------| -| `MODEL.SWIN_UNETR_FEATURE_SIZE` | `48` | Feature dimension | -| `MODEL.DEPTHS` | `(2, 2, 2, 2)` | Layers per stage | -| `MODEL.SWIN_UNETR_NUM_HEADS` | `(3, 6, 12, 24)` | Attention heads per stage | -| `MODEL.SWIN_UNETR_DROPOUT_RATE` | `0.0` | Dropout rate | -| `MODEL.ATTN_DROP_RATE` | `0.0` | Attention dropout | -| `MODEL.USE_V2` | `False` | Use Swin UNETR v2 | -| `MODEL.SPATIAL_DIMS` | `3` | Spatial dimensions | - -### UNETR-Specific Options - -| Key | Default | Description | -|-----|---------|-------------| -| `MODEL.UNETR_FEATURE_SIZE` | `16` | Feature dimension | -| `MODEL.HIDDEN_SIZE` | `768` | Transformer hidden dimension | -| `MODEL.MLP_DIM` | `3072` | Feedforward dimension | -| `MODEL.UNETR_NUM_HEADS` | `12` | Number of attention heads | -| `MODEL.POS_EMBED` | `'perceptron'` | Position embedding type | -| `MODEL.UNETR_DROPOUT_RATE` | `0.0` | Dropout rate | +| `model.arch.profile` | Architecture profile name | `monai_unet`, `rsunet`, `mednext_s` | +| `model.arch.type` | Architecture type (set by profile) | `monai_unet`, `rsunet`, `mednext` | +| `model.out_channels` | Number of output channels | `1` (binary), `3` (BCD), `12` (affinity) | +| `model.input_size` | Model input patch size `[z, y, x]` | `[112, 112, 112]` | +| `model.output_size` | Model output patch size `[z, y, x]` | `[112, 112, 112]` | +| `model.monai.filters` | Filter sizes per stage (MONAI UNet) | `[32, 64, 128, 256]` | +| `model.monai.dropout` | Dropout rate (MONAI UNet) | `0.1` | +| `model.rsunet.width` | Channel widths per stage (RSUNet) | `[18, 36, 48, 64, 80]` | +| `model.rsunet.norm` | Normalization type (RSUNet): `batch`, `group` | `group` | +| `model.rsunet.activation` | Activation function (RSUNet): `elu`, `relu` | `elu` | +| `model.rsunet.num_groups` | Number of groups for group norm | `4` | +| `model.rsunet.down_factors` | Downsampling factors per stage | `[[1,2,2], [1,2,2], ...]` | ## Loss Functions -Set `MODEL.LOSS_OPTION` (a list of lists, one per target): +Losses are configured under `model.loss` using either a profile or explicit list: + +### Loss Profiles + +| Profile | Losses | Use Case | +|---------|--------|----------| +| `loss_binary` | WeightedBCEWithLogitsLoss + DiceLoss | Binary segmentation (mitochondria, synapse) | +| `loss_bcd` | BCE+Dice (foreground) + BCE+Dice (boundary) + WeightedMSE (distance) | Multi-task BCD segmentation | +| `loss_per_channel` | PerChannelBCEWithLogitsLoss (auto pos_weight) | Per-channel affinity prediction | +| `loss_bd` | BCE+Dice (all but last channel) + WeightedMSE (last channel) | Boundary + distance multi-task | -| Loss | Key | Use Case | -|------|-----|----------| -| Weighted Binary Cross-Entropy | `WeightedBCE` | Binary segmentation | -| Weighted BCE with Logits | `WeightedBCEWithLogitsLoss` | Binary segmentation (numerically stable, no activation needed) | -| Weighted BCE Focal Loss | `WeightedBCEFocalLoss` | Binary segmentation with class imbalance | -| Dice Loss | `DiceLoss` | Binary segmentation (overlap-based) | -| Weighted-Sample Dice Loss | `WSDiceLoss` | Per-sample weighted Dice | -| Weighted Cross-Entropy | `WeightedCE` | Multi-class segmentation | -| Weighted MSE | `WeightedMSE` | Regression targets (e.g., distance transforms, flow fields) | -| Weighted MAE | `WeightedMAE` | Regression targets | +### Direct Loss Configuration -Multiple losses can be combined for a single target: ```yaml -MODEL: - LOSS_OPTION: [["WeightedBCEWithLogitsLoss", "DiceLoss"]] - LOSS_WEIGHT: [[1.0, 1.0]] +model: + loss: + losses: + - function: WeightedBCEWithLogitsLoss + weight: 1.0 + kwargs: {reduction: mean} + - function: DiceLoss + weight: 1.0 + kwargs: {sigmoid: true, smooth_nr: 1e-5, smooth_dr: 1e-5} ``` -## Regularization Options +Available loss functions: +- `WeightedBCEWithLogitsLoss` — Binary cross-entropy with logits +- `DiceLoss` — Overlap-based Dice loss +- `WeightedMSELoss` — Mean squared error (for distance/regression targets) +- `PerChannelBCEWithLogitsLoss` — Per-channel BCE with automatic positive weight balancing + +Loss entries support `pred_slice` and `target_slice` to apply different losses to different output channels (e.g., `"0:1"` for first channel only). + +## Pipeline Profiles -Set `MODEL.REGU_OPT`: +Pipeline profiles bundle model output channels, loss, label transforms, and decoding into a single preset. Set via `default.pipeline_profile`: -| Regularization | Key | Description | -|---------------|-----|-------------| -| Binary | `Binary` | Encourages binary predictions | -| Foreground-Contour Consistency | `FgContour` | Enforces consistency between foreground and contour predictions | -| Contour-DT Consistency | `ContourDT` | Consistency between contour and distance transform | -| Foreground-DT Consistency | `FgDT` | Consistency between foreground and distance transform | -| Non-overlap | `Nonoverlap` | Prevents overlapping instance predictions | +| Profile | Out Channels | Description | +|---------|-------------|-------------| +| `binary` | 1 | Binary segmentation (foreground/background) with BCE+Dice loss | +| `bcd` | 3 | Multi-task: binary + boundary + distance transform | +| `affinity-12` | 12 | 12-channel affinity prediction with per-channel BCE and ABISS decoding | -## Target Options (MODEL.TARGET_OPT) +## Label Transforms -The target option string encodes what the model is predicting: +Label transforms convert raw instance labels into training targets. Configured via `data.label_transform.profile`: -| Code | Target Type | OUT_PLANES | Description | -|------|-------------|-----------|-------------| -| `"0"` | Binary mask | 1 | Standard binary segmentation (foreground/background) | -| `"1"` | Synaptic polarity | 3 | Signed polarity prediction for synapses | -| `"2"` | Affinity map | 3 | 3-channel affinity map for instance segmentation | -| `"3"` | Small object mask | 1 | Optimized for small objects | -| `"4"` | Contour map | 1 | Boundary contour prediction | -| `"5"` | Distance transform | 1 | Quantized distance transform for watershed | -| `"7"` | Flow field | 2 | Cellpose-style gradient flow for instance segmentation | -| `"9-N"` | Multi-class | N | N-class semantic segmentation (e.g., `"9-12"` for 12 classes) | +| Profile | Targets | Description | +|---------|---------|-------------| +| `label_bcd` | binary + instance_boundary + instance_edt | Multi-task BCD labels | +| `label_affinity_12` | affinity (12 offsets) | Short and long-range affinity maps with DeepEM crop | -## Weight Options (MODEL.WEIGHT_OPT) +## Activation Profiles -Controls per-voxel loss weighting: +Channel activations applied during TTA and inference. Referenced by pipeline profiles: -| Code | Description | -|------|-------------| -| `"0"` | No weighting (uniform) | -| `"1"` | Binary mask weighting (weight foreground vs. background) | +| Profile | Description | +|---------|-------------| +| `act_binary` | Sigmoid activation for binary channels | +| `act_bcd` | Per-channel activations for BCD outputs | -## Output Activations (MODEL.OUTPUT_ACT) +## Decoding Profiles -Applied to model output during loss computation: +Post-inference decoding to convert raw predictions into segmentation. Referenced by pipeline profiles: -| Activation | When to Use | -|-----------|-------------| -| `"none"` | When loss handles raw logits (e.g., BCEWithLogitsLoss) | -| `"sigmoid"` | Binary segmentation output | -| `"softmax"` | Multi-class segmentation output | -| `"tanh"` | Flow field / regression with [-1, 1] range | +| Profile | Description | +|---------|-------------| +| `decoding_abiss` | ABISS watershed for affinity-to-instance conversion | +| `decoding_bcd` | BCD-based instance segmentation decoding | diff --git a/server_api/chatbot/file_summaries/PyTC-Overview.md b/server_api/chatbot/file_summaries/PyTC-Overview.md index 86a5537c..10059c46 100644 --- a/server_api/chatbot/file_summaries/PyTC-Overview.md +++ b/server_api/chatbot/file_summaries/PyTC-Overview.md @@ -1,51 +1,62 @@ # PyTorch Connectomics (PyTC) — Overview -PyTorch Connectomics (PyTC) is a deep-learning framework for automatic and semi-automatic **semantic and instance segmentation** of volumetric biomedical images, especially electron-microscopy (EM) connectomics data. It is built on PyTorch and maintained by the Visual Computing Group (VCG) at Harvard University. +PyTorch Connectomics (PyTC) is a deep-learning framework for automatic and semi-automatic **semantic and instance segmentation** of volumetric biomedical images, especially electron-microscopy (EM) connectomics data. It is built on PyTorch Lightning and uses Hydra for configuration management. ## What PyTC Can Do PyTC supports the following segmentation tasks out of the box: - **Binary semantic segmentation** — e.g., mitochondria vs. background (Lucchi dataset). -- **Instance segmentation via affinity maps** — e.g., neuron boundary detection (SNEMI, CREMI). -- **Multi-class semantic segmentation** — e.g., labeling 12+ organelle types in a single volume. -- **Nucleus segmentation** — e.g., NucMM mouse and zebrafish nuclei. -- **Cell segmentation with flow fields** — e.g., Cellpose-style gradient prediction. -- **Distance-transform prediction** — boundary distance maps for watershed-based instance segmentation. +- **Instance segmentation via affinity maps** — e.g., neuron boundary detection (SNEMI, CREMI) with ABISS watershed decoding. +- **Multi-task BCD segmentation** — binary + boundary contour + distance transform prediction. +- **Nucleus segmentation** — e.g., NucMM zebrafish nuclei. - **Synaptic cleft detection** — e.g., CREMI synapse detection. -- **Large-scale tile-based processing** — chunked training and inference for datasets that do not fit in memory (MitoEM, Scutoid). +- **Large-scale processing** — training and inference for large datasets. +- **Hyperparameter tuning** — built-in Optuna-based parameter search (e.g., watershed thresholds). ## Supported Data Formats -- **Images**: TIFF stacks (`.tif`), HDF5 (`.h5`), JSON tile manifests (`.json`). +- **Images**: TIFF stacks (`.tif`), HDF5 (`.h5`). - **Labels**: Same formats as images; can be binary masks, instance labels, or multi-class label volumes. -- **Data dimensionality**: Both **3D** volumetric data and **2D** image data are supported. Set `DATASET.DO_2D: True` for 2D mode. +- **Data dimensionality**: Both **3D** volumetric data and **2D** image data are supported. ## Key Capabilities +- **Profile-based configuration**: Reusable presets for architectures, optimizers, augmentations, and pipelines. - **Multi-task learning**: Train a single model with multiple loss functions and target types simultaneously. -- **Distributed training**: Multi-GPU training via DataParallel (DP) or DistributedDataParallel (DDP) with optional mixed-precision. -- **Comprehensive data augmentation**: 11 augmentation types designed specifically for EM data, including missing-section simulation, misalignment, and motion blur. -- **Flexible model zoo**: Multiple encoder-decoder architectures (UNet 3D/2D, UNet++, FPN, UNETR, Swin UNETR, DeepLabV3) with configurable block types and backbones. -- **YACS configuration system**: All settings are controlled via YAML config files. Any default can be overridden on the command line. -- **Stochastic Weight Averaging (SWA)**: Built-in support for improved generalization. -- **Inference augmentation**: Test-time augmentation with configurable modes (mean, min) and number of augmented copies. +- **Automatic distributed training**: Multi-GPU training via PyTorch Lightning with automatic mixed-precision. +- **Comprehensive data augmentation**: Profile-based augmentation system with 10+ augmentation types designed for EM data, including missing-section simulation, misalignment, and motion blur. +- **Flexible model zoo**: MONAI UNet, RSUNet, and MedNeXt architectures with configurable parameters. +- **Hydra configuration system**: All settings are controlled via YAML config files with dot-separated key overrides on the command line. +- **EMA (Exponential Moving Average)**: Built-in support for improved generalization. +- **Test-time augmentation**: Flip and rotation augmentation with configurable ensemble modes (mean, min). +- **Cache-aware inference**: Automatically detects and reuses cached predictions. +- **Sharded inference**: Distribute inference across multiple machines. ## Entry Point All training and inference is launched through a single script: ``` -python scripts/main.py --config-file [OPTIONS] [OVERRIDES] +python scripts/main.py --config [OPTIONS] [key=value overrides] ``` Key flags: -- `--config-file` (required): Path to the YAML configuration file. -- `--config-base`: Optional base config that the main config extends. -- `--inference`: Run inference instead of training. -- `--checkpoint`: Path to a model checkpoint for resuming training or running inference. -- `--distributed`: Enable distributed multi-GPU training (DDP). -- `SECTION.KEY=value` (positional): Override any config option on the command line. +- `--config` (required): Path to the YAML configuration file (relative to PyTC root, e.g., `tutorials/mito_lucchi++.yaml`). +- `--mode`: Mode of operation: `train` (default), `test`, `tune`, `tune-test`. +- `--checkpoint`: Path to a model checkpoint (`.ckpt`) for resuming training or running inference. +- `--fast-dev-run`: Run 1 batch for quick debugging. +- `--demo`: Quick demo using `tutorials/minimal.yaml`. +- `--debug-config`: Print fully resolved config and exit. +- `--reset-max-epochs N`: Reset max_epochs when resuming training. +- `--shard-id` / `--num-shards`: For distributed test sharding. +- `key=value` (positional): Override any config option using Hydra dot-path syntax. + +Example overrides: +``` +python scripts/main.py --config tutorials/mito_lucchi++.yaml \ + optimization.optimizer.lr=0.001 data.dataloader.batch_size=4 +``` ## What PyTC Cannot Do @@ -54,6 +65,6 @@ PyTC is a **segmentation-only** framework. It does not support: - Image classification. - Generative models (GANs, diffusion). - Non-image modalities (text, audio). -- Models or architectures not listed in the model zoo (unet_3d, unet_2d, fpn_3d, unet_plus_3d, unet_plus_2d, deeplabv3a/b/c, unetr, swinunetr). +- Models or architectures not listed in the model zoo (monai_unet, rsunet, mednext). If a user asks to train a model or run a task that is not one of the supported segmentation tasks listed above, the request cannot be fulfilled with PyTC. diff --git a/server_api/chatbot/file_summaries/PyTC-Training.md b/server_api/chatbot/file_summaries/PyTC-Training.md index a0625824..c0091f72 100644 --- a/server_api/chatbot/file_summaries/PyTC-Training.md +++ b/server_api/chatbot/file_summaries/PyTC-Training.md @@ -5,112 +5,138 @@ This document explains how to configure and run training jobs with PyTorch Conne ## Training Command ``` -python scripts/main.py --config-file [OVERRIDES] +python scripts/main.py --config [key=value overrides] ``` Example: ``` -python scripts/main.py --config-file configs/Lucchi-Mitochondria.yaml SOLVER.BASE_LR=0.001 SOLVER.SAMPLES_PER_BATCH=4 +python scripts/main.py --config tutorials/mito_lucchi++.yaml optimization.optimizer.lr=0.001 data.dataloader.batch_size=4 ``` -Overrides use dotted YAML key paths: `SECTION.KEY=value`. Multiple overrides can be appended. +Overrides use Hydra/OmegaConf dot-separated key paths: `section.subsection.key=value`. Multiple overrides can be appended. ## Required Configuration Sections -Every training config must specify at minimum: a model architecture, dataset paths, and solver settings. The YAML files in the `configs/` directory provide complete working examples; users should start from the closest matching config and modify it rather than writing one from scratch. +Every training config must specify at minimum: a model architecture, dataset paths, and optimization settings. The YAML files in the `tutorials/` directory provide complete working examples; users should start from the closest matching config and modify it rather than writing one from scratch. Configs use a profile-based system where base profiles (in `tutorials/bases/`) define reusable presets for architectures, optimizers, augmentations, etc. -### DATASET Section +### Data Section + +| Key | Description | Example | +|-----|-------------|---------| +| `data.train.image` | Path(s) to training image file(s) | `datasets/lucchi++/train_im.h5` | +| `data.train.label` | Path(s) to training label file(s) | `datasets/lucchi++/train_mito.h5` | +| `data.train.resolution` | Voxel resolution `[z, y, x]` in nm | `[5, 5, 5]` | +| `data.dataloader.batch_size` | Number of samples per batch | `8` | +| `data.dataloader.patch_size` | Training patch size `[z, y, x]` | `[112, 112, 112]` | +| `data.dataloader.profile` | Dataloader profile: `cached` or `lazy` | `cached` | +| `data.data_transform.pad_size` | Padding `[z, y, x]` for context | `[8, 128, 128]` | +| `data.augmentation.profile` | Augmentation profile (see PyTC-Augmentation) | `aug_standard` | +| `data.image_transform.clip_percentile_low` | Low percentile for intensity clipping | `0.0` | +| `data.image_transform.clip_percentile_high` | High percentile for intensity clipping | `1.0` | + +### Optimization Section (Training Hyperparameters) | Key | Description | Default | |-----|-------------|---------| -| `DATASET.INPUT_PATH` | Root directory containing images and labels | `'path/to/input'` | -| `DATASET.IMAGE_NAME` | Image filename(s) relative to INPUT_PATH. Multiple files separated by `@` | None | -| `DATASET.LABEL_NAME` | Label filename(s) relative to INPUT_PATH. Multiple files separated by `@` | None | -| `DATASET.OUTPUT_PATH` | Directory where checkpoints and logs are saved | `'path/to/output'` | -| `DATASET.PAD_SIZE` | Padding `[z, y, x]` to avoid border sampling issues | `[2, 64, 64]` | -| `DATASET.IS_ISOTROPIC` | Whether the voxels are cubic (isotropic) | `False` | -| `DATASET.DO_2D` | Enable 2D training mode | `False` | -| `DATASET.LOAD_2D` | Load data as 2D slices | `False` | -| `DATASET.DATA_SCALE` | Scale ratio `[z, y, x]` for resampling input | `[1.0, 1.0, 1.0]` | -| `DATASET.MEAN` | Normalization mean | `0.5` | -| `DATASET.STD` | Normalization std | `0.5` | -| `DATASET.LABEL_BINARY` | Binarize the label volume | `False` | -| `DATASET.LABEL_EROSION` | Erode label masks by N pixels | None | -| `DATASET.DO_CHUNK_TITLE` | Enable tile-based (chunked) training for large datasets | `0` | -| `DATASET.DATA_CHUNK_NUM` | Number of chunks `[z, y, x]` for tiled data | `[1, 1, 1]` | -| `DATASET.DATA_CHUNK_ITER` | Iterations per chunk | `1000` | -| `DATASET.REJECT_SAMPLING.SIZE_THRES` | Reject patches with foreground below this threshold (-1 = disabled) | `-1` | -| `DATASET.REJECT_SAMPLING.P` | Probability of applying reject sampling | `0.95` | - -### SOLVER Section (Training Hyperparameters) +| `optimization.optimizer.name` | Optimizer: `"AdamW"`, `"Adam"`, `"SGD"` | `"AdamW"` | +| `optimization.optimizer.lr` | Learning rate | `0.0003` | +| `optimization.optimizer.weight_decay` | Weight decay | `0.01` | +| `optimization.optimizer.betas` | Adam/AdamW beta parameters | `[0.9, 0.999]` | +| `optimization.optimizer.eps` | Epsilon for numerical stability | `1.0e-08` | +| `optimization.max_steps` | Total training steps (alternative to max_epochs) | varies | +| `optimization.max_epochs` | Total training epochs | `100` | +| `optimization.n_steps_per_epoch` | Steps per epoch | `1000` | +| `optimization.gradient_clip_val` | Gradient clipping value | `1.0` | +| `optimization.accumulate_grad_batches` | Gradient accumulation steps | `1` | +| `optimization.precision` | Training precision: `bf16-mixed`, `16-mixed`, `32` | `bf16-mixed` | +| `optimization.profile` | Optimizer profile preset (see below) | `warmup_cosine_lr` | + +### Optimizer Profiles + +Instead of configuring each optimizer/scheduler field individually, you can use a predefined profile: + +| Profile | Optimizer | Scheduler | Description | +|---------|-----------|-----------|-------------| +| `warmup_cosine_lr` | AdamW (lr=0.0003) | WarmupCosineLR | Warmup then cosine decay (recommended) | +| `cosine_annealing_lr` | AdamW (lr=0.0003) | CosineAnnealingLR | Standard cosine annealing | +| `reduce_on_plateau` | AdamW (lr=0.0003) | ReduceLROnPlateau | Reduce LR when loss plateaus | +| `step_lr` | AdamW (lr=0.0003) | StepLR | Step decay every N epochs | +| `multistep_lr` | AdamW (lr=0.0003) | MultiStepLR | Decay at specific epoch milestones | + +### Scheduler Section | Key | Description | Default | |-----|-------------|---------| -| `SOLVER.NAME` | Optimizer name: `"SGD"`, `"Adam"`, `"AdamW"` | `"SGD"` | -| `SOLVER.BASE_LR` | Base learning rate | `0.001` | -| `SOLVER.MOMENTUM` | SGD momentum | `0.9` | -| `SOLVER.BETAS` | Adam/AdamW beta parameters | `(0.9, 0.999)` | -| `SOLVER.WEIGHT_DECAY` | Weight decay | `0.0001` | -| `SOLVER.LR_SCHEDULER_NAME` | LR scheduler: `"MultiStepLR"`, `"WarmupMultiStepLR"`, `"WarmupCosineLR"`, `"ReduceLROnPlateau"`, `"OneCycle"` | `"MultiStepLR"` | -| `SOLVER.STEPS` | Milestones for MultiStepLR | `(30000, 35000)` | -| `SOLVER.GAMMA` | LR decay factor at each step | `0.1` | -| `SOLVER.WARMUP_ITERS` | Number of warmup iterations | `1000` | -| `SOLVER.WARMUP_FACTOR` | Initial LR multiplier during warmup | `1/1000` | -| `SOLVER.SAMPLES_PER_BATCH` | Number of samples per GPU per batch | `2` | -| `SOLVER.ITERATION_TOTAL` | Total training iterations | `40000` | -| `SOLVER.ITERATION_SAVE` | Save a checkpoint every N iterations | `5000` | -| `SOLVER.ITERATION_VAL` | Run validation every N iterations | `5000` | -| `SOLVER.ITERATION_RESTART` | Restart iteration counter from 0 when loading a pretrained model | `False` | -| `SOLVER.CLIP_GRADIENTS.ENABLED` | Enable gradient clipping | `False` | -| `SOLVER.CLIP_GRADIENTS.CLIP_TYPE` | Clipping type: `"value"` or `"norm"` | `"value"` | -| `SOLVER.CLIP_GRADIENTS.CLIP_VALUE` | Clipping threshold | `1.0` | -| `SOLVER.SWA.ENABLED` | Enable Stochastic Weight Averaging | `False` | -| `SOLVER.SWA.START_ITER` | Iteration to begin SWA | `90000` | -| `SOLVER.SWA.LR_FACTOR` | SWA learning rate = BASE_LR × LR_FACTOR | `0.05` | - -### MONITOR Section +| `optimization.scheduler.name` | LR scheduler name | `WarmupCosineLR` | +| `optimization.scheduler.interval` | Update interval: `epoch` or `step` | `epoch` | +| `optimization.scheduler.warmup_epochs` | Number of warmup epochs | `3` | +| `optimization.scheduler.warmup_start_lr` | LR at start of warmup | `1.0e-05` | +| `optimization.scheduler.min_lr` | Minimum LR | `1.0e-06` | + +### EMA (Exponential Moving Average) | Key | Description | Default | |-----|-------------|---------| -| `MONITOR.LOG_OPT` | Logging options `[loss, lr, gpu_usage]` (1=on, 0=off) | `[1, 1, 0]` | -| `MONITOR.VIS_OPT` | Visualization options `[image, feature_map]` | `[0, 16]` | -| `MONITOR.ITERATION_NUM` | Log every N[0] iters; visualize every N[1] iters | `[20, 200]` | +| `optimization.ema.enabled` | Enable EMA | `true` (in some configs) | +| `optimization.ema.decay` | EMA decay rate | `0.999` | +| `optimization.ema.warmup_steps` | Steps before EMA starts | `500` | +| `optimization.ema.validate_with_ema` | Use EMA weights for validation | `true` | + +### Monitor Section -### SYSTEM Section +| Key | Description | Default | +|-----|-------------|---------| +| `monitor.logging.scalar.loss_every_n_steps` | Log loss every N steps | `50` | +| `monitor.logging.images.log_every_n_epochs` | Log images every N epochs | `10` | +| `monitor.logging.images.max_images` | Max images to log | `8` | +| `monitor.logging.images.num_slices` | Number of slices to visualize | `2` | +| `monitor.checkpoint.save_top_k` | Keep the top K checkpoints | `3` | +| `monitor.checkpoint.monitor` | Metric to monitor for checkpointing | `train_loss_total_epoch` | +| `monitor.checkpoint.save_every_n_epochs` | Save checkpoint every N epochs | `10` | +| `monitor.checkpoint.dirpath` | Directory for saving checkpoints | `outputs//checkpoints/` | +| `monitor.early_stopping.enabled` | Enable early stopping | `false` | + +### System Section | Key | Description | Default | |-----|-------------|---------| -| `SYSTEM.NUM_GPUS` | Number of GPUs | `4` | -| `SYSTEM.NUM_CPUS` | Number of CPU workers for data loading | `4` | -| `SYSTEM.DISTRIBUTED` | Use DistributedDataParallel | `False` | -| `SYSTEM.PARALLEL` | Parallelism mode: `'DP'` (DataParallel) or `'DDP'` | `'DP'` | +| `system.num_gpus` | Number of GPUs (-1 = all available) | `-1` | +| `system.num_workers` | Number of CPU workers (-1 = all available) | `-1` | +| `system.seed` | Random seed for reproducibility | `42` | +| `system.profile` | System profile: `all-gpu-cpu` or `single-gpu-cpu` | `all-gpu-cpu` | -## Distributed Training +## Resuming Training -To run multi-GPU training with DDP: +To resume from a checkpoint: ``` -python -m torch.distributed.launch --nproc_per_node= scripts/main.py --distributed --config-file +python scripts/main.py --config --checkpoint ``` -Mixed-precision training (`MODEL.MIXED_PRECESION: True`) only works with DDP. +To reset max_epochs when resuming: +``` +python scripts/main.py --config --checkpoint --reset-max-epochs 500 +``` -## Resuming Training +## Quick Debug Run -To resume from a checkpoint: ``` -python scripts/main.py --config-file --checkpoint +python scripts/main.py --config tutorials/mito_lucchi++.yaml --fast-dev-run ``` -Training resumes from the saved iteration unless `SOLVER.ITERATION_RESTART: True` is set. +This runs a single training batch for quick validation. Use `--fast-dev-run 2` for 2 batches. -## Transfer Learning / Fine-tuning +## Demo Mode + +``` +python scripts/main.py --demo +``` -Set `MODEL.PRE_MODEL` to the path of a pretrained checkpoint. The model will load those weights before training begins. Use `MODEL.FINETUNE` to add a suffix to saved checkpoint names so they do not overwrite the original. +Uses `tutorials/minimal.yaml` with `--fast-dev-run` for a quick sanity check. ## Tips -- Start from the closest bundled config in `configs/` and override only what you need. -- For anisotropic EM data, use odd input sizes (e.g., `[17, 257, 257]`) when not using pooling layers to avoid feature mismatching. -- For isotropic data, set `DATASET.IS_ISOTROPIC: True` and `AUGMENTOR.FLIP.DO_ZTRANS: 1` to enable z-axis augmentation. -- Use `DATASET.REJECT_SAMPLING.SIZE_THRES` to avoid sampling empty patches in sparse datasets. -- The output directory receives checkpoints (`checkpoint_NNNNN.pth.tar`) and a `config.yaml` snapshot. +- Start from the closest bundled config in `tutorials/` and override only what you need. +- Use optimizer profiles (e.g., `optimization.profile=warmup_cosine_lr`) instead of configuring each field manually. +- Use augmentation profiles (e.g., `data.augmentation.profile=aug_standard`) for recommended augmentation settings. +- For sparse datasets, enable reject sampling in the config: `data.dataloader.reject_sampling.size_thres=1000`. +- The output directory receives checkpoints (`.ckpt` files) and a saved config snapshot. diff --git a/server_api/chatbot/rag_eval.py b/server_api/chatbot/rag_eval.py index 765ddcbe..c994c952 100644 --- a/server_api/chatbot/rag_eval.py +++ b/server_api/chatbot/rag_eval.py @@ -66,7 +66,7 @@ "FileManager.md", "hard", ), - # ── ModelTraining.md ──────────────────────────────────────────────── + # ── ModelTraining.md (UI page) ────────────────────────────────────── ("How do I start a training job?", "ModelTraining.md", "easy"), ("What are the three steps to configure training?", "ModelTraining.md", "easy"), ("What inputs are required for model training?", "ModelTraining.md", "easy"), @@ -78,14 +78,9 @@ ), ("How do I stop a running training job?", "ModelTraining.md", "easy"), ("Can I edit the raw YAML for training?", "ModelTraining.md", "medium"), - ( - "What model architectures are available for training?", - "ModelTraining.md", - "hard", - ), ("Do training config fields have AI help buttons?", "ModelTraining.md", "medium"), ( - "I want to train a segmentation model - where do I configure that?", + "I want to train a segmentation model - where do I configure that in the UI?", "ModelTraining.md", "hard", ), @@ -99,12 +94,12 @@ "ModelTraining.md", "hard", ), - # ── ModelInference.md ─────────────────────────────────────────────── - ("How do I run inference?", "ModelInference.md", "easy"), - ("What inputs do I need for inference?", "ModelInference.md", "easy"), + # ── ModelInference.md (UI page) ───────────────────────────────────── + ("How do I run inference from the UI?", "ModelInference.md", "easy"), + ("What inputs do I need for inference in the app?", "ModelInference.md", "easy"), ("What is the checkpoint path for inference?", "ModelInference.md", "medium"), ("How do I stop an inference job?", "ModelInference.md", "easy"), - ("What is the augmentations setting in inference?", "ModelInference.md", "medium"), + ("What is the augmentations slider in inference?", "ModelInference.md", "medium"), ("Is there an Input Label field for inference?", "ModelInference.md", "hard"), ( "I have a trained model and want to run predictions on new data - which page?", @@ -113,7 +108,7 @@ ), ("Do I need ground truth labels to run inference?", "ModelInference.md", "hard"), ( - "Where do I specify which trained weights to use for prediction?", + "Where do I specify which trained weights to use for prediction in the UI?", "ModelInference.md", "hard", ), @@ -247,177 +242,292 @@ ("What data formats does PyTC accept?", "PyTC-Overview.md", "easy"), ("Does PyTC support object detection?", "PyTC-Overview.md", "medium"), ("What is the main entry point script for PyTC?", "PyTC-Overview.md", "medium"), + ("Can I use PyTC for image classification?", "PyTC-Overview.md", "hard"), + ("What configuration system does PyTC use?", "PyTC-Overview.md", "medium"), + ("Does PyTC use Hydra or YACS for configs?", "PyTC-Overview.md", "medium"), + ("What is the --config flag for?", "PyTC-Overview.md", "medium"), + ("What modes does PyTC support? train, test, tune?", "PyTC-Overview.md", "medium"), ( - "Can I use PyTC for image classification?", + "I want to do a quick test to see if PyTC is working, is there a demo mode?", + "PyTC-Overview.md", + "hard", + ), + ( + "What is the --fast-dev-run flag?", "PyTC-Overview.md", "hard", ), # ── PyTC-Training.md ─────────────────────────────────────────────── + # Easy: direct key/concept lookup ("How do I run a training job with PyTC?", "PyTC-Training.md", "easy"), ("What learning rate schedulers are available?", "PyTC-Training.md", "easy"), ("How do I resume training from a checkpoint?", "PyTC-Training.md", "easy"), + ("What is optimization.optimizer.lr?", "PyTC-Training.md", "easy"), + ("What is the training command for PyTC?", "PyTC-Training.md", "easy"), + # Medium: requires understanding of config structure ("What optimizer options does PyTC support?", "PyTC-Training.md", "medium"), ("How do I enable gradient clipping?", "PyTC-Training.md", "medium"), - ("What is reject sampling in PyTC?", "PyTC-Training.md", "medium"), + ("What optimizer profiles are available?", "PyTC-Training.md", "medium"), + ("What is the warmup_cosine_lr profile?", "PyTC-Training.md", "medium"), + ("How do I set optimization.max_epochs?", "PyTC-Training.md", "medium"), + ("What does optimization.precision control?", "PyTC-Training.md", "medium"), + ("How do I change data.dataloader.batch_size?", "PyTC-Training.md", "medium"), + ("What is monitor.checkpoint.save_top_k?", "PyTC-Training.md", "medium"), + ("How do I set the system.num_gpus?", "PyTC-Training.md", "medium"), + ("What is the EMA feature in PyTC training?", "PyTC-Training.md", "medium"), + # Real-user training questions (natural language) + ("Can I train my model for a longer period of time?", "PyTC-Training.md", "medium"), + ("How do I lower the learning rate?", "PyTC-Training.md", "easy"), + ("How often does the model save checkpoints?", "PyTC-Training.md", "medium"), + ("Can I use AdamW instead of SGD?", "PyTC-Training.md", "medium"), ( - "How do I run distributed multi-GPU training?", + "How do I use cosine learning rate decay?", "PyTC-Training.md", - "hard", + "medium", ), ( - "What does SOLVER.ITERATION_TOTAL control?", + "Can I override config values from the command line?", "PyTC-Training.md", - "hard", + "medium", ), - # Real-user training questions - ("Can I train my model for a longer period of time?", "PyTC-Training.md", "medium"), - ("How do I change the batch size?", "PyTC-Training.md", "medium"), - ("How do I lower the learning rate?", "PyTC-Training.md", "easy"), - ("How often does the model save checkpoints?", "PyTC-Training.md", "medium"), - ("Can I use Adam instead of SGD?", "PyTC-Training.md", "medium"), - ("What is stochastic weight averaging in PyTC?", "PyTC-Training.md", "medium"), ( - "I want to fine-tune a pretrained model on my own data, how?", + "Where are training checkpoints saved?", "PyTC-Training.md", - "hard", + "medium", ), + # Hard: requires synthesis or non-obvious mapping ( - "My training is crashing because of NaN gradients, what should I try?", + "How do I train for 200 epochs instead of the default?", "PyTC-Training.md", "hard", ), ( - "How do I use cosine learning rate decay?", + "My training is crashing because of NaN gradients, what should I try?", "PyTC-Training.md", - "medium", + "hard", ), ( - "Can I override config values from the command line?", + "What is mixed precision training and how do I enable bf16?", "PyTC-Training.md", - "medium", + "hard", ), ( - "How do I set the number of GPUs for training?", + "How do I use the ReduceLROnPlateau scheduler?", "PyTC-Training.md", - "medium", + "hard", ), ( - "What is mixed precision training and how do I enable it?", + "How do I enable early stopping during training?", "PyTC-Training.md", "hard", ), ( - "I want to train for 200K iterations instead of the default, how?", + "What is the --reset-max-epochs flag for when resuming?", "PyTC-Training.md", "hard", ), ( - "Where are training checkpoints saved?", + "I want to accumulate gradients over multiple batches, how?", "PyTC-Training.md", - "medium", + "hard", ), # ── PyTC-Inference.md ────────────────────────────────────────────── + # Easy ("How do I run inference with a trained model?", "PyTC-Inference.md", "easy"), ("What is test-time augmentation in PyTC?", "PyTC-Inference.md", "easy"), - ("How does inference stride affect results?", "PyTC-Inference.md", "medium"), - ("What blending modes are available for inference?", "PyTC-Inference.md", "medium"), + ("What does --mode test do?", "PyTC-Inference.md", "easy"), + ("How do I run predictions on my test data?", "PyTC-Inference.md", "easy"), + # Medium + ("What blending modes are available for sliding window inference?", "PyTC-Inference.md", "medium"), + ("How do I increase the sliding window overlap?", "PyTC-Inference.md", "medium"), + ("Can I increase inference.batch_size for faster inference?", "PyTC-Inference.md", "medium"), + ("What is inference.sliding_window.sw_batch_size?", "PyTC-Inference.md", "medium"), + ("How do I enable TTA flips during inference?", "PyTC-Inference.md", "medium"), + ("What is Gaussian blending in sliding window inference?", "PyTC-Inference.md", "medium"), ( - "How do I run chunked inference on a large volume?", + "How do I specify where to save inference output?", "PyTC-Inference.md", - "hard", + "medium", ), ( - "What does DO_SINGLY do in PyTC inference?", + "Can I use test-time augmentation to improve accuracy?", "PyTC-Inference.md", - "hard", + "medium", ), - # Real-user inference questions - ("How do I make my inference predictions smoother?", "PyTC-Inference.md", "medium"), - ("Can I increase the batch size for faster inference?", "PyTC-Inference.md", "medium"), ("What output format does inference produce?", "PyTC-Inference.md", "medium"), + # Hard ( "I'm seeing tile boundary artifacts in my predictions, how do I fix that?", "PyTC-Inference.md", "hard", ), ( - "My volume is too large to fit in memory for inference, what should I do?", + "How do I run sharded inference across multiple machines?", "PyTC-Inference.md", "hard", ), ( - "How do I run inference on multiple test volumes one at a time?", + "What is cache-aware inference in PyTC?", "PyTC-Inference.md", "hard", ), ( - "Can I use test-time augmentation to improve accuracy?", - "PyTC-Inference.md", - "medium", - ), - ( - "What does the --inference flag do?", + "Should I use the same config file for inference as I used for training?", "PyTC-Inference.md", - "easy", + "hard", ), ( - "How do I specify where to save inference output?", + "How do I configure postprocessing after inference?", "PyTC-Inference.md", - "medium", + "hard", ), ( - "Should I use the same config file for inference as I used for training?", + "What is inference.test_time_augmentation.ensemble_mode?", "PyTC-Inference.md", "hard", ), ( - "How do I set the output activation for inference predictions?", + "How do I use the --shard-id and --num-shards flags?", "PyTC-Inference.md", - "medium", + "hard", ), ( - "What is Gaussian blending in PyTC?", + "What does inference.evaluation.enabled control?", "PyTC-Inference.md", - "medium", + "hard", ), # ── PyTC-Models.md ───────────────────────────────────────────────── + # Easy ("What model architectures does PyTC support?", "PyTC-Models.md", "easy"), - ("What is the difference between unet_3d and unet_2d?", "PyTC-Models.md", "easy"), - ("What block types are available?", "PyTC-Models.md", "easy"), + ("What is the MONAI UNet architecture?", "PyTC-Models.md", "easy"), ("What loss functions does PyTC support?", "PyTC-Models.md", "easy"), - ("What is residual_se block type?", "PyTC-Models.md", "medium"), - ("What backbones are available for FPN?", "PyTC-Models.md", "medium"), - ("What are the Swin UNETR specific options?", "PyTC-Models.md", "medium"), - ("What does TARGET_OPT control?", "PyTC-Models.md", "hard"), + ("What is the RSUNet architecture?", "PyTC-Models.md", "easy"), + # Medium + ("What is the difference between monai_unet and rsunet?", "PyTC-Models.md", "medium"), + ("What MedNeXt model sizes are available?", "PyTC-Models.md", "medium"), + ("What is the loss_binary profile?", "PyTC-Models.md", "medium"), + ("How do I set model.arch.profile?", "PyTC-Models.md", "medium"), + ("What pipeline profiles are available?", "PyTC-Models.md", "medium"), + ("What does the binary pipeline profile do?", "PyTC-Models.md", "medium"), + ("How do I combine BCE loss with Dice loss?", "PyTC-Models.md", "medium"), + ("What is PerChannelBCEWithLogitsLoss?", "PyTC-Models.md", "medium"), + # Hard + ( + "What is the affinity-12 pipeline profile?", + "PyTC-Models.md", + "hard", + ), + ( + "How do I configure loss functions with pred_slice and target_slice?", + "PyTC-Models.md", + "hard", + ), + ( + "What label transforms are available for instance segmentation?", + "PyTC-Models.md", + "hard", + ), + ( + "What is the ABISS decoding profile?", + "PyTC-Models.md", + "hard", + ), + ( + "How do I set up a model with 12 output channels for affinity prediction?", + "PyTC-Models.md", + "hard", + ), + ( + "What is the difference between DiceLoss and WeightedBCEWithLogitsLoss?", + "PyTC-Models.md", + "hard", + ), ( - "How do I configure a multi-class segmentation model?", + "What activation profiles are used during TTA?", + "PyTC-Models.md", + "hard", + ), + ( + "How does the BCD pipeline combine boundary, contour, and distance targets?", "PyTC-Models.md", "hard", ), # ── PyTC-Augmentation.md ─────────────────────────────────────────── + # Easy ("What augmentations does PyTC support?", "PyTC-Augmentation.md", "easy"), ("What is CutBlur augmentation?", "PyTC-Augmentation.md", "easy"), + ("What augmentation profiles are available?", "PyTC-Augmentation.md", "easy"), + ("What is the aug_standard profile?", "PyTC-Augmentation.md", "easy"), + # Medium ("How do I disable elastic deformation?", "PyTC-Augmentation.md", "medium"), ( - "What augmentation settings should I use for isotropic data?", + "What augmentation profile should I use for isotropic EM data?", + "PyTC-Augmentation.md", + "medium", + ), + ( + "What is the aug_em_neuron profile for?", + "PyTC-Augmentation.md", + "medium", + ), + ( + "How do I set data.augmentation.profile?", + "PyTC-Augmentation.md", + "medium", + ), + ( + "What is the intensity augmentation in PyTC?", + "PyTC-Augmentation.md", + "medium", + ), + ( + "What is the difference between aug_light and aug_strong?", "PyTC-Augmentation.md", "medium", ), + # Hard ( "How do I simulate missing sections in my training data?", "PyTC-Augmentation.md", "hard", ), ( - "What does AUGMENTOR.FLIP.DO_ZTRANS do?", + "What is the defect_mutex feature in augmentation?", + "PyTC-Augmentation.md", + "hard", + ), + ( + "How does the copy-paste augmentation work for instance segmentation?", + "PyTC-Augmentation.md", + "hard", + ), + ( + "What augmentation profile matches the DeepEM recipe?", + "PyTC-Augmentation.md", + "hard", + ), + ( + "How do I add mixup augmentation to my training?", + "PyTC-Augmentation.md", + "hard", + ), + ( + "What does data.augmentation.misalignment.displacement control?", "PyTC-Augmentation.md", "hard", ), # ── PyTC-Configs.md ──────────────────────────────────────────────── + # Easy ("What bundled configs does PyTC have?", "PyTC-Configs.md", "easy"), ("Which config should I use for mitochondria segmentation?", "PyTC-Configs.md", "easy"), + ("Where are the tutorial configs located?", "PyTC-Configs.md", "easy"), + # Medium ("What config is used for CREMI synapse detection?", "PyTC-Configs.md", "medium"), - ("Which config uses the Swin UNETR architecture?", "PyTC-Configs.md", "medium"), + ("What architecture does neuron_snemi.yaml use?", "PyTC-Configs.md", "medium"), + ("What is the mito_lucchi++.yaml config for?", "PyTC-Configs.md", "medium"), + ("What base profiles do the tutorial configs inherit from?", "PyTC-Configs.md", "medium"), + ("Are there configs for nucleus segmentation?", "PyTC-Configs.md", "medium"), + # Hard ( "What is the recommended config for neuron instance segmentation?", "PyTC-Configs.md", @@ -428,83 +538,79 @@ "PyTC-Configs.md", "hard", ), + ( + "What is the profile-based inheritance system for configs?", + "PyTC-Configs.md", + "hard", + ), + ( + "What misc configs are available for specialized tasks?", + "PyTC-Configs.md", + "hard", + ), + ( + "How do I override data paths in a tutorial config?", + "PyTC-Configs.md", + "hard", + ), # ── PyTC-Evaluation.md ───────────────────────────────────────────── + # Easy ("How do I evaluate segmentation results in PyTC?", "PyTC-Evaluation.md", "easy"), ("What is adapted Rand error?", "PyTC-Evaluation.md", "easy"), + # Medium ("What evaluation metric should I use for instance segmentation?", "PyTC-Evaluation.md", "medium"), ("How do I compute IoU for binary segmentation?", "PyTC-Evaluation.md", "medium"), + ("How do I know if my model is good?", "PyTC-Evaluation.md", "medium"), ( - "What is variation of information in segmentation evaluation?", + "How do I compare my prediction against ground truth?", "PyTC-Evaluation.md", - "hard", + "medium", ), ( - "Does PyTC automatically compute metrics after inference?", + "I have a binary segmentation prediction, how do I get precision and recall?", "PyTC-Evaluation.md", - "hard", + "medium", ), - # Real-user evaluation questions - ("How do I know if my model is good?", "PyTC-Evaluation.md", "medium"), - ("What metric should I report for the SNEMI3D challenge?", "PyTC-Evaluation.md", "hard"), ( - "How do I compare my prediction against ground truth?", + "What Python functions does PyTC provide for evaluation?", "PyTC-Evaluation.md", "medium", ), + # Hard ( - "Is my model over-segmenting or under-segmenting? How can I tell?", + "What is variation of information in segmentation evaluation?", "PyTC-Evaluation.md", "hard", ), ( - "What is panoptic quality and how do I compute it?", + "Does PyTC automatically compute metrics after inference?", "PyTC-Evaluation.md", "hard", ), + ("What metric should I report for the SNEMI3D challenge?", "PyTC-Evaluation.md", "hard"), ( - "I have a binary segmentation prediction, how do I get precision and recall?", + "Is my model over-segmenting or under-segmenting? How can I tell?", "PyTC-Evaluation.md", - "medium", + "hard", ), ( - "How do I evaluate my CREMI synapse detection results?", + "What is panoptic quality and how do I compute it?", "PyTC-Evaluation.md", "hard", ), ( - "What Python functions does PyTC provide for evaluation?", + "How do I evaluate my CREMI synapse detection results?", "PyTC-Evaluation.md", - "medium", - ), - # ── Cross-doc / Models real-user questions ───────────────────────── - ( - "Can I use a transformer-based model in PyTC?", - "PyTC-Models.md", - "medium", - ), - ( - "How do I combine BCE loss with Dice loss?", - "PyTC-Models.md", - "medium", - ), - ( - "What normalization options are available for my model?", - "PyTC-Models.md", - "medium", - ), - ( - "I want to train an affinity model for neuron segmentation, what target option do I use?", - "PyTC-Models.md", "hard", ), ( - "How do I set up a model with 12 output classes?", - "PyTC-Models.md", + "What is the cremi_distance metric?", + "PyTC-Evaluation.md", "hard", ), ( - "What is the difference between DiceLoss and WeightedBCEWithLogitsLoss?", - "PyTC-Models.md", + "Can I enable automatic evaluation during inference with inference.evaluation.enabled?", + "PyTC-Evaluation.md", "hard", ), ] diff --git a/server_api/chatbot/tools.py b/server_api/chatbot/tools.py index c8e13bf1..a125d390 100644 --- a/server_api/chatbot/tools.py +++ b/server_api/chatbot/tools.py @@ -19,30 +19,6 @@ def get_pytc_root() -> Path: return Path(__file__).parent.parent.parent / "pytorch_connectomics" -CONFIG_DESCRIPTIONS = { - "Lucchi-Mitochondria.yaml": ( - "3D mitochondria segmentation on Lucchi EM dataset. " - "Good baseline for beginners learning 3D EM segmentation." - ), - "CREMI-Base.yaml": ( - "Base config for CREMI neuronal segmentation benchmark. " - "Standard UNet for boundary/affinity prediction." - ), - "SNEMI-Affinity-UNet.yaml": ( - "Affinity-based neuron segmentation on SNEMI3D. " - "Predicts affinity maps for instance segmentation." - ), - "MitoEM-R-Base.yaml": ( - "Base config for MitoEM mitochondria benchmark. " - "Large-scale 3D mitochondria segmentation." - ), - "NucMM-Zebrafish-Base.yaml": ( - "Nucleus segmentation on NucMM Zebrafish dataset. " - "Multi-task learning with boundary and distance maps." - ), -} - - @tool def list_training_configs() -> List[Dict[str, str]]: """ @@ -53,40 +29,40 @@ def list_training_configs() -> List[Dict[str, str]]: List of configs with name, path, model info, and description """ pytc_root = get_pytc_root() - configs_dir = pytc_root / "configs" + tutorials_dir = pytc_root / "tutorials" - if not configs_dir.exists(): - return [{"error": f"Configs directory not found at {configs_dir}"}] + if not tutorials_dir.exists(): + return [{"error": f"Tutorials directory not found at {tutorials_dir}"}] configs = [] - for yaml_file in configs_dir.rglob("*.yaml"): + for yaml_file in tutorials_dir.rglob("*.yaml"): + # Skip base profile files — they are not standalone configs + if yaml_file.parent.name == "bases": + continue + rel_path = yaml_file.relative_to(pytc_root) name = yaml_file.name - # Get description from known configs or try to extract from file - description = CONFIG_DESCRIPTIONS.get(name, "") + description = "" model_arch = "unknown" try: with open(yaml_file, "r") as f: config_data = yaml.safe_load(f) - # Extract model architecture if config_data: - model_arch = config_data.get("MODEL", {}).get("ARCHITECTURE", "unknown") - - # Try to get description from comments if not in known list - if not description: - with open(yaml_file, "r") as f: - for line in f: - if line.startswith("#"): - description = line.lstrip("# ").strip() - break - elif not line.strip(): - continue - else: - break + # Extract model architecture from Hydra config format + default = config_data.get("default", {}) + if isinstance(default, dict): + model_cfg = default.get("model", {}) + if isinstance(model_cfg, dict): + arch_cfg = model_cfg.get("arch", {}) + if isinstance(arch_cfg, dict): + model_arch = arch_cfg.get("profile", arch_cfg.get("type", "unknown")) + + # Read description from the YAML's own description field + description = config_data.get("description", "") except Exception: pass @@ -162,9 +138,9 @@ def list_checkpoints(experiment_name: Optional[str] = None) -> List[Dict[str, st else: return [{"info": f"Experiment '{experiment_name}' not found in outputs/"}] - # Search for .pth and .ckpt files + # Search for .ckpt (Lightning) and .pth files for search_dir in search_dirs: - for pattern in ["*.pth", "*.ckpt"]: + for pattern in ["*.ckpt", "*.pth"]: for ckpt_file in search_dir.rglob(pattern): rel_path = ckpt_file.relative_to(pytc_root)