Files
imagepipeline/docs/MODULE_DEVELOPMENT.md
2026-05-30 11:33:07 +02:00

276 lines
8.0 KiB
Markdown

# Module development guide
This document explains how to add new processing modules to the image pipeline framework.
## Overview
Each module is a Python class that:
1. Declares a unique `name`
2. Defines accepted parameters via `parameters()`
3. Implements `run(ctx)` to process images
4. Registers itself with `@register`
Pipeline scripts reference modules by name and pass parameters as kwargs.
## Quick Start
Create a new file in `imagepipeline/modules/`, e.g. `my_module.py`:
```python
from imagepipeline.core.context import ModuleContext
from imagepipeline.core.params import Param
from imagepipeline.modules.base import BaseModule
from imagepipeline.modules.registry import register
@register
class MyModule(BaseModule):
name = "my_module"
description = "Short description for documentation"
@classmethod
def parameters(cls) -> dict[str, Param]:
return {
"strength": Param(
"float",
default=1.0,
help="Processing strength from 0.0 to 1.0",
),
}
def run(self, ctx: ModuleContext) -> None:
ctx.output_dir.mkdir(parents=True, exist_ok=True)
for src in ctx.input_paths:
dst = ctx.output_dir / src.name
# ... write dst ...
```
Import the module in `imagepipeline/modules/__init__.py` so it registers on load:
```python
import imagepipeline.modules.my_module # noqa: F401
```
Use it in a pipeline script:
```python
with Pipeline(name="my_run", input_dir=INPUT) as p:
result = p.step("my_module", inputs="input", strength=0.8)
p.run()
```
## ModuleContext
Every `run()` receives a `ModuleContext`:
| Field | Description |
|-------|-------------|
| `input_paths` | Flat list of input images for this step |
| `matched_groups` | For multi-input steps: list of path groups matched by stem |
| `output_dir` | Directory where outputs must be written |
| `params` | Validated parameters |
| `pipeline_output_root` | Root folder of the current run |
| `step_id` | Internal step identifier (e.g. `my_module_01`) |
### Single-input modules
Use `ctx.input_paths` — one path per image in the batch:
```python
for src in ctx.input_paths:
dst = ctx.output_dir / src.name
process(src, dst)
```
### Multi-input modules
When a step receives multiple input references (e.g. background + foreground), the runner matches images by filename stem across all sources.
Use `ctx.matched_groups` — each entry is a list of paths with the same stem:
```python
for group in ctx.matched_groups:
background, foreground = group
dst = ctx.output_dir / foreground.name
composite(background, foreground, dst)
```
If a stem is missing in any source, the pipeline fails with a clear error before your module runs.
## Parameters
Define parameters with `Param`:
```python
Param("string", default="value", help="Description")
Param("int", required=True, help="Required integer")
Param("float", default=0.5)
Param("bool", default=False)
Param("path")
Param("list", default=[])
Param("string", choices=("a", "b", "c"))
```
Supported types: `string`, `int`, `float`, `bool`, `path`, `list`.
Unknown kwargs passed to `p.step()` raise a validation error. Missing required parameters raise as well.
## Base Classes
### `BaseModule`
Use for pure Python processing or custom logic.
Optional class attributes:
- `supported_input_formats` — tuple of extensions (informational)
- `description` — short module description
Optional methods:
- `check_dependencies()` — raise `DependencyError` if tools are missing
- `list_output_images(ctx)` — default: all images in `output_dir`
### `SubprocessModule`
Use for CLI tools (ImageMagick, gmic, darktable-cli, rembg):
```python
from imagepipeline.modules.base import SubprocessModule
from imagepipeline.utils.subprocess import run_command
@register
class MyCliModule(SubprocessModule):
name = "my_cli_module"
command_candidates = ("my-tool", "my-tool-fallback")
@classmethod
def check_dependencies(cls) -> None:
super().check_dependencies() # verifies command_candidates
def run(self, ctx: ModuleContext) -> None:
tool = self.resolve_command()
for src in ctx.input_paths:
dst = ctx.output_dir / src.name
run_command([tool, str(src), str(dst)])
```
`run_command()` captures stderr and raises `RuntimeError` with the command output on failure.
## Output Conventions
- Write one output file per input, keeping the **same filename** (especially the stem) so downstream steps can match batches.
- Only write image files into `ctx.output_dir` — the runner discovers outputs by scanning for known image extensions.
- Create `ctx.output_dir` if your tool does not do so automatically.
## Step Folder Naming
The runner assigns output folders automatically: `{module_name}_{nn}`.
Using the same module twice in one pipeline produces separate folders:
```
imagemagick_grayscale_01/
imagemagick_grayscale_02/
darktable_style_01/
darktable_style_02/
```
You do not choose folder names in the module.
## Format Warnings
Declare `supported_input_formats` on your module. When chaining modules with incompatible formats (e.g. JPEG without alpha after rembg), document expected behavior in your module's `description` and consider converting explicitly in `run()`.
## Testing
Add tests under `tests/`:
1. **Unit tests** — parameter validation, matching logic (no external tools)
2. **Integration tests** — run the module against a real image; skip if CLI tool missing:
```python
import shutil
import pytest
pytestmark = pytest.mark.skipif(
not shutil.which("magick") and not shutil.which("convert"),
reason="ImageMagick not installed",
)
```
## Checklist for New Modules
- [ ] Unique `name` (snake_case)
- [ ] `@register` decorator
- [ ] `parameters()` documents all kwargs
- [ ] `run()` writes outputs preserving stems
- [ ] `check_dependencies()` if external tools required
- [ ] Import in `imagepipeline/modules/__init__.py`
- [ ] Test added (unit and/or integration)
## AI Modules
Optional dependencies: `pip install -e ".[ai]"` (numpy, Pillow, torch).
All AI modules inherit from `AIModule` (`imagepipeline/modules/ai_base.py`), which adds:
| Parameter | Default | Description |
|-----------|---------|-------------|
| `skip_existing` | `True` | Skip when output file already exists |
| `max_edge` | `2048` | Downscale long edge before inference (`0` = full resolution) |
| `device` | `"cpu"` | Inference device (v1: CPU only) |
Use `iter_input_images(ctx, processor)` in `run()` — it handles skip, resize, upscale, and per-image ETA logging.
### Built-in AI modules
| Module | Backend | Notes |
|--------|---------|-------|
| `ai_exposure` | Zero-DCE++ (PyTorch) | Weights auto-download on first run |
| `ai_tone_map` | HDRNet or CLAHE fallback | Set `checkpoint` to a creotiv-format `.pth`; empty uses CLAHE |
| `openrouter_edit` | OpenRouter API | Requires `OPENROUTER_API_KEY`; default model Flux Klein 4B |
| `comfy_flux_edit` | ComfyUI HTTP API | Experimental; export workflow to `workflows/comfy/flux_klein_edit_api.json` |
Example pipeline: `pipelines/example_ai.py`.
### Adding a new AI module
```python
from imagepipeline.modules.ai_base import AIModule
@register
class MyAIModule(AIModule):
name = "my_ai_module"
@classmethod
def parameters(cls) -> dict[str, Param]:
params = dict(super().parameters())
params["strength"] = Param("float", default=1.0)
return params
def run(self, ctx: ModuleContext) -> None:
self.configure_torch(ctx.params["device"])
def process(src: Path, dst: Path, index: int, total: int) -> None:
...
self.iter_input_images(ctx, process)
```
Put shared inference code under `imagepipeline/ai/`.
## Planned Modules
These slots follow the same interface — no pipeline API changes needed:
| Module | Base class | Tool |
|--------|-----------|------|
| `imagemagick_resize` | SubprocessModule | ImageMagick |
| `darktable_style` | SubprocessModule | darktable-cli |
| `rembg` | SubprocessModule | rembg |
| `gmic` | SubprocessModule | gmic |
| `composite` | BaseModule | Pillow |