Frank Schwenk
a60a18a253
Introduce ruff lint/format config, expand .gitignore, and reformat Python sources. Co-authored-by: Cursor <cursoragent@cursor.com>
126 lines
4.1 KiB
Python
126 lines
4.1 KiB
Python
from __future__ import annotations
|
|
|
|
import shutil
|
|
import time
|
|
from collections.abc import Callable
|
|
from pathlib import Path
|
|
from typing import Any, ClassVar
|
|
|
|
from imagepipeline.core.context import ModuleContext
|
|
from imagepipeline.core.params import Param
|
|
from imagepipeline.modules.base import BaseModule
|
|
|
|
|
|
def _format_eta(seconds: float) -> str:
|
|
seconds = max(0, int(seconds))
|
|
minutes, secs = divmod(seconds, 60)
|
|
hours, minutes = divmod(minutes, 60)
|
|
if hours:
|
|
return f"{hours}h{minutes:02d}m"
|
|
if minutes:
|
|
return f"{minutes}m{secs:02d}s"
|
|
return f"{secs}s"
|
|
|
|
|
|
class AIModule(BaseModule):
|
|
"""Base class for AI pipeline modules (CPU-first, batch-friendly)."""
|
|
|
|
ai_common_parameters: ClassVar[dict[str, Param]] = {
|
|
"skip_existing": Param(
|
|
"bool",
|
|
default=True,
|
|
help="Skip images whose output file already exists",
|
|
),
|
|
"max_edge": Param(
|
|
"int",
|
|
default=2048,
|
|
help="Max long edge for inference (0 = full resolution)",
|
|
),
|
|
"device": Param(
|
|
"string",
|
|
default="cpu",
|
|
choices=("cpu",),
|
|
help="Inference device (v1 supports cpu only)",
|
|
),
|
|
}
|
|
|
|
@classmethod
|
|
def parameters(cls) -> dict[str, Param]:
|
|
return dict(cls.ai_common_parameters)
|
|
|
|
@classmethod
|
|
def all_parameters(cls) -> dict[str, Param]:
|
|
merged = dict(cls.ai_common_parameters)
|
|
merged.update(cls.parameters())
|
|
return merged
|
|
|
|
@classmethod
|
|
def validate_module_params(cls, raw: dict[str, Any]) -> dict[str, Any]:
|
|
from imagepipeline.core.params import validate_params
|
|
|
|
return validate_params(cls.all_parameters(), raw)
|
|
|
|
def output_path(self, src: Path, ctx: ModuleContext) -> Path:
|
|
return ctx.output_dir / src.name
|
|
|
|
def configure_torch(self, device_name: str) -> None:
|
|
import torch
|
|
|
|
if device_name != "cpu":
|
|
raise ValueError(f"Unsupported device: {device_name!r} (only 'cpu' in v1)")
|
|
torch.set_num_threads(24)
|
|
|
|
def iter_input_images(
|
|
self,
|
|
ctx: ModuleContext,
|
|
processor: Callable[[Path, Path, int, int], None],
|
|
) -> None:
|
|
import tempfile
|
|
|
|
from imagepipeline.ai.imaging import resize_max_edge, resize_to_size
|
|
|
|
ctx.output_dir.mkdir(parents=True, exist_ok=True)
|
|
skip_existing = ctx.params["skip_existing"]
|
|
max_edge = ctx.params["max_edge"]
|
|
total = len(ctx.input_paths)
|
|
elapsed_times: list[float] = []
|
|
|
|
for index, src in enumerate(ctx.input_paths, start=1):
|
|
dst = self.output_path(src, ctx)
|
|
if skip_existing and dst.is_file():
|
|
if ctx.logger is not None:
|
|
ctx.logger.skipped(self.name, index, total, dst.name)
|
|
continue
|
|
|
|
self.log_image(ctx, index, total, src)
|
|
started = time.perf_counter()
|
|
|
|
if max_edge > 0:
|
|
with tempfile.TemporaryDirectory(prefix="imagepipeline_ai_") as tmp:
|
|
tmp_path = Path(tmp)
|
|
work_src = tmp_path / f"work_{src.name}"
|
|
work_out = tmp_path / f"out_{src.name}"
|
|
orig_w, orig_h, _, _ = resize_max_edge(src, work_src, max_edge)
|
|
processor(work_src, work_out, index, total)
|
|
if (orig_w, orig_h) != self._image_size(work_out):
|
|
resize_to_size(work_out, dst, orig_w, orig_h)
|
|
else:
|
|
shutil.copy2(work_out, dst)
|
|
else:
|
|
processor(src, dst, index, total)
|
|
|
|
elapsed = time.perf_counter() - started
|
|
elapsed_times.append(elapsed)
|
|
remaining = total - index
|
|
avg = sum(elapsed_times) / len(elapsed_times)
|
|
eta = avg * remaining if remaining else 0.0
|
|
if ctx.logger is not None:
|
|
ctx.logger.image_done(self.name, index, total, dst.name, elapsed, eta_seconds=eta)
|
|
|
|
@staticmethod
|
|
def _image_size(path: Path) -> tuple[int, int]:
|
|
from PIL import Image
|
|
|
|
with Image.open(path) as image:
|
|
return image.size
|