Spaces:
Running
Running
| """Processing Module Framework — pluggable pipeline nodes. | |
| 每个功能(相机输入、骨骼识别、手套、标注、审核、导出)是一个独立模块: | |
| 一个文件 = 一个模块(元数据 + run()),放在 app/processing/modules/ 下, | |
| 自动发现注册。前端画布与 worker 执行都从注册表取模块定义。 | |
| 新增自定义模块三步: | |
| 1. 复制 modules/example 模板为 app/processing/modules/<slug>.py | |
| 2. 填写 slug/label/category/inputs/outputs/config_schema 等元数据,实现 run(ctx) | |
| 3. 重启后端与 worker —— 前端面板自动出现可拖拽卡片,worker 自动执行 | |
| """ | |
| from __future__ import annotations | |
| from dataclasses import dataclass, field as dc_field | |
| from pathlib import Path | |
| from typing import Any, Callable | |
| def field(name: str, kind: str, label: str, default=None, **extra) -> dict: | |
| """config_schema 单字段定义: {"name","type","label","default",**extra}。""" | |
| value = {"name": name, "type": kind, "label": label, "default": default} | |
| value.update(extra) | |
| return value | |
| class ModuleResult: | |
| """Returned by a ProcessingModule after execution (legacy, kept for compat).""" | |
| success: bool | |
| data: dict[str, Any] = dc_field(default_factory=dict) | |
| outputs: dict[str, "ArtifactRef"] = dc_field(default_factory=dict) | |
| metrics: dict[str, Any] = dc_field(default_factory=dict) | |
| error: str | None = None | |
| class ArtifactRef: | |
| """Portable reference to a workflow input/output artifact.""" | |
| kind: str | |
| path: str | None = None | |
| source_key: str | None = None | |
| schema_version: str = "1.0" | |
| metadata: dict[str, Any] = dc_field(default_factory=dict) | |
| def to_dict(self) -> dict[str, Any]: | |
| return { | |
| "kind": self.kind, | |
| "path": self.path, | |
| "source_key": self.source_key, | |
| "schema_version": self.schema_version, | |
| "metadata": self.metadata, | |
| } | |
| def from_dict(cls, value: dict[str, Any]) -> "ArtifactRef": | |
| return cls( | |
| kind=str(value.get("kind", "unknown")), | |
| path=value.get("path"), | |
| source_key=value.get("source_key"), | |
| schema_version=str(value.get("schema_version", "1.0")), | |
| metadata=dict(value.get("metadata") or {}), | |
| ) | |
| class ModuleDescriptor: | |
| """Metadata shared by the backend module registry and card palette.""" | |
| slug: str | |
| version: str | |
| category: str | |
| label: str | |
| icon: str | |
| color: str | |
| inputs: tuple[dict[str, str], ...] = () | |
| outputs: tuple[dict[str, str], ...] = () | |
| default_config: dict[str, Any] = dc_field(default_factory=dict) | |
| config_schema: tuple[dict[str, Any], ...] = () | |
| execution_target: str = "worker" | |
| capabilities: tuple[str, ...] = () | |
| def to_dict(self) -> dict[str, Any]: | |
| return { | |
| "type": self.slug, | |
| "slug": self.slug, | |
| "version": self.version, | |
| "category": self.category, | |
| "label": self.label, | |
| "icon": self.icon, | |
| "color": self.color, | |
| "inputs": list(self.inputs), | |
| "outputs": list(self.outputs), | |
| "defaultConfig": self.default_config, | |
| "default_config": self.default_config, | |
| "configSchema": list(self.config_schema), | |
| "config_schema": list(self.config_schema), | |
| "executionTarget": self.execution_target, | |
| "execution_target": self.execution_target, | |
| "capabilities": list(self.capabilities), | |
| } | |
| class ModuleSkip(Exception): | |
| """模块主动跳过(无有效输入等)—— 节点标记 skipped,不视为失败。""" | |
| def __init__(self, reason: str): | |
| self.reason = reason | |
| super().__init__(reason) | |
| class JobContext: | |
| """Worker 与模块之间的统一执行上下文。""" | |
| node_id: str | |
| node_type: str | |
| config: dict[str, Any] | |
| job: dict[str, Any] | |
| input_root: Path # 批次解压目录(源文件扫描) | |
| output_root: Path # 输出根(outputs/) | |
| incoming: dict[str, ArtifactRef] # targetHandle → 上游 ArtifactRef(源节点为空 dict) | |
| progress: Callable[[float], None] # 本节点进度 0..1 | |
| node_states: dict # 节点状态累加器 | |
| def output_dir(self) -> Path: | |
| d = self.output_root / self.node_id | |
| d.mkdir(parents=True, exist_ok=True) | |
| return d | |
| def ref(self, kind: str, path: Path, source_key: str | None = None, | |
| metadata: dict | None = None) -> ArtifactRef: | |
| """构造 ArtifactRef,路径自动相对化:input_root 下 → 相对 input_root,否则相对 output_root。""" | |
| try: | |
| rel = path.resolve().relative_to(self.input_root.resolve()) | |
| except ValueError: | |
| rel = path.resolve().relative_to(self.output_root.resolve()) | |
| return ArtifactRef(kind=kind, path=str(rel).replace("\\", "/"), | |
| source_key=source_key, metadata=metadata or {}) | |
| def resolve(self, ref: ArtifactRef) -> Path | None: | |
| """按 input_root/output_root 依次解析 ArtifactRef 路径。""" | |
| if not ref or not ref.path: | |
| return None | |
| p = self.input_root / ref.path | |
| if p.exists(): | |
| return p | |
| p = self.output_root / ref.path | |
| return p if p.exists() else None | |
| def find_videos(self) -> list[Path]: | |
| """按 config 的 source_keys/source_key/position 匹配批次内视频(见 batch.py)。""" | |
| from app.processing.batch import find_videos | |
| return find_videos(self.input_root, self.config) | |
| def find_parquet(self) -> list[Path]: | |
| """批次内全部 parquet(排除 auto_labels/hand_kpts 等合并产物)。""" | |
| from app.processing.batch import find_parquet | |
| return find_parquet(self.input_root) | |
| def skip(self, reason: str): | |
| raise ModuleSkip(reason) | |
| class ProcessingModule: | |
| """Base class for a pipeline processing node. | |
| Subclass, fill in the metadata class attributes, implement ``run()``, | |
| and decorate with ``@register`` from app.processing.registry — the | |
| module appears in the Workflow Studio palette and runs in the worker | |
| automatically. | |
| """ | |
| slug: str = "" | |
| version: str = "1.0" | |
| category: str = "process" # input | process | review | export | |
| label: str = "" | |
| icon: str = "ant-design:appstore-outlined" | |
| color: str = "#64748b" | |
| inputs: tuple[dict[str, str], ...] = () | |
| outputs: tuple[dict[str, str], ...] = () | |
| default_config: dict[str, Any] = dc_field(default_factory=dict) | |
| config_schema: tuple[dict[str, Any], ...] = () | |
| execution_target: str = "worker" | |
| capabilities: tuple[str, ...] = () | |
| def descriptor(cls) -> ModuleDescriptor: | |
| return ModuleDescriptor( | |
| slug=cls.slug, version=cls.version, category=cls.category, label=cls.label, | |
| icon=cls.icon, color=cls.color, inputs=tuple(cls.inputs), | |
| outputs=tuple(cls.outputs), default_config=dict(cls.default_config), | |
| config_schema=tuple(cls.config_schema), execution_target=cls.execution_target, | |
| capabilities=tuple(cls.capabilities), | |
| ) | |
| def run(self, ctx: JobContext) -> dict[str, ArtifactRef] | None: | |
| """执行入口。返回输出 handle → ArtifactRef;None 或 raise ModuleSkip → 节点 skipped。""" | |
| raise NotImplementedError | |