nnU-Net dataset conversion Agent

An intelligent dataset conversion tool built on the OpenHands Agent SDK that automatically converts raw medical-image segmentation datasets into the nnU-Net v2 format based on natural-language instructions.

Features

• React mode: the Agent decides its next action autonomously rather than following a predefined workflow.
• Automatic retry: on validation failure, the Agent automatically cleans up and regenerates improved code.
• Multi-format support: supports .nii.gz, .png, .nrrd, .mha, and other medical-image formats.
• Long-running tasks: validation timeout is set to 2 hours to accommodate large datasets.
• Detailed logging: shows the Agent's reasoning and tool calls.

Prerequisites

1. Python environment: Python 3.8+.
2. Dependencies:

   pip install openhands-ai openhands-sdk

3. nnUNet v2: installed with environment variables configured.

   export nnUNet_raw="/path/to/nnunet_raw_data"
   export nnUNet_preprocessed="/path/to/nnunet_preprocessed"

4. API key: OpenRouter API key.

   export OPENROUTER_API_KEY="your-api-key"

Quick start

Basic usage

from nnunet_agent import DataConversionAgent

# Create the Agent
agent = DataConversionAgent()

# Run a conversion task
instruction = """
Please preprocess the dataset located at /data/raw_dataset and convert it to Dataset700_MyDataset.

Raw dataset structure:
- images/ directory contains medical images in .nii.gz format
- labels/ directory contains the corresponding segmentation labels

This is a single-channel 3D medical-image segmentation task.
"""

result = agent.run(instruction)
print(result)

Command-line usage

# Run the example directly
.venv/bin/python nnunet_agent.py

# Or use interactive mode
.venv/bin/python
>>> from nnunet_agent import DataConversionAgent
>>> agent = DataConversionAgent()
>>> agent.run("your instruction...")

How it works

React loop

1. Thought    -> analyze the current state and decide the next step
2. Action     -> select and invoke a tool
3. Observation-> inspect the tool's output
4. Repeat until the task is complete -> output Final Answer

Tool set

Tool	Purpose	Usage
read_directory	Read directory structure	Analyze the raw dataset layout
read_format_guide	Read the format specification	Learn the nnUNet v2 requirements
execute_python_code	Execute Python code	Run the dataset-conversion script
run_validation	Run the validation command	Check whether the conversion is correct
cleanup_dataset	Clean up a failed conversion	Remove bad data before retrying

Typical flow

graph LR
    A[User instruction] --> B[Analyze raw dataset]
    B --> C[Learn nnUNet format]
    C --> D[Generate conversion code]
    D --> E[Execute code]
    E --> F[Validate result]
    F --> G{Success?}
    G -->|Yes| H[Done]
    G -->|No| I[Clean up data]
    I --> D

Loading

Configuration

Edit config.py to change configuration:

# Timeout (seconds)
VALIDATION_TIMEOUT = 7200  # 2 hours

# Maximum Agent iterations
MAX_ITERATIONS = 20

# LLM model
LLM_MODEL = "deepseek/deepseek-v3.2"

Project structure

nnprep/
├── __init__.py                    # Public API (convert_dataset, DataConversionAgent)
├── cli.py                         # Command-line entry point
├── config.py                      # Default LLM / runtime configuration
├── openhands_agent.py             # OpenHands Agent main program
├── openhands_tools.py             # OpenHands-registered custom tools
├── tools.py                       # Tool-function collection
├── timeout_decorator.py           # Generic timeout helper
├── dataset_format.md              # nnU-Net v2 format specification
└── skills/
    └── system_prompt.md           # System prompt template for the Agent

Example scenarios

Scenario 1: single-channel 2D dataset

agent.run("""
Convert /data/cell_segmentation to Dataset800_CellSeg.
Raw data: 2D cell images in PNG format.
""")

Scenario 2: multi-channel 3D dataset

agent.run("""
Convert /data/brain_mri to Dataset801_BrainMRI.
Contains 4 channels: T1, T1ce, T2, FLAIR.
File format: .nii.gz.
""")

Scenario 3: custom-structure dataset

agent.run("""
Dataset path: /data/custom_dataset
Target: Dataset802_Custom

Notes:
- Images are under the scans/ directory.
- Labels are under the annotations/ directory.
- Filename format: patient_001.nrrd.
""")

Troubleshooting

Issue 1: validation timeout

Solution: increase VALIDATION_TIMEOUT or test with a smaller dataset.

Issue 2: code generation error

Solution: the Agent retries automatically, typically resolving the issue within 2-3 iterations.

Issue 3: dataset_format.md not found

Solution: make sure dataset_format.md is in the same directory as the script.

Notes

Important notes:

• The Agent automatically deletes Dataset directories whose validation failed; make sure you do not accidentally lose important data.
• On first use, run on a test dataset first.
• Validation can take a long time (depending on dataset size).

Contributing

Issues and pull requests are welcome.

License

MIT License