276 lines
8.0 KiB
Markdown
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 |
|