pixtreme-train 0.1.0__py3-none-any.whl

pixtreme/train/common/__init__.py

@@ -0,0 +1,18 @@
+"""Shared configuration, image metrics, and checkpoint saving for model trainers."""
+
+from .checkpoint import save_checkpoint
+from .configuration import Option, plain, write_yaml
+from .metrics import EvaluationMetrics, FrameMetrics, MaterialMetrics, json_metric, measure, rgb_frame
+
+__all__ = (
+    "Option",
+    "plain",
+    "write_yaml",
+    "FrameMetrics",
+    "MaterialMetrics",
+    "EvaluationMetrics",
+    "json_metric",
+    "rgb_frame",
+    "measure",
+    "save_checkpoint",
+)

pixtreme/train/common/checkpoint.py

@@ -0,0 +1,31 @@
+"""Save checkpoint payloads without replacing a previous file on write failure."""
+
+from __future__ import annotations
+
+import os
+from collections.abc import Mapping
+from pathlib import Path
+from tempfile import NamedTemporaryFile
+from typing import Any
+
+import torch
+
+
+def save_checkpoint(path: str | Path, payload: Mapping[str, Any], *, prefix: str = ".checkpoint-") -> None:
+    """Flush and sync a sibling temporary file, then atomically replace the destination.
+
+    The caller assembles and validates the payload. Its values are passed unchanged
+    to torch.save; callers needing weights_only loading must supply compatible values.
+    """
+    path = Path(path)
+    temporary: Path | None = None
+    try:
+        with NamedTemporaryFile(prefix=prefix, suffix=".tmp", dir=path.parent, delete=False) as stream:
+            temporary = Path(stream.name)
+            torch.save(payload, stream)
+            stream.flush()
+            os.fsync(stream.fileno())
+        os.replace(temporary, path)
+    finally:
+        if temporary is not None:
+            temporary.unlink(missing_ok=True)

pixtreme/train/common/configuration.py

@@ -0,0 +1,61 @@
+"""Validate option values and write reusable YAML configurations."""
+
+from __future__ import annotations
+
+import argparse
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+
+@dataclass(frozen=True)
+class Option:
+    """Describe an option's accepted YAML types, parser, and optional choices."""
+
+    name: str
+    parse: Callable[[str], Any]
+    default: Any
+    kinds: tuple[type, ...]
+    description: str
+    choices: tuple[str, ...] = ()
+    repeat: bool = False
+    required: bool = False
+
+    def yaml_value(self, value: Any) -> Any:
+        if value is None and self.default is None:
+            return None
+        if self.repeat:
+            if not isinstance(value, list):
+                raise ValueError("must be a list")
+            return [self._scalar(item) for item in value]
+        return self._scalar(value)
+
+    def _scalar(self, value: Any) -> Any:
+        if type(value) not in self.kinds:
+            raise ValueError(f"must be {self.description}")
+        try:
+            result = self.parse(str(value))
+        except argparse.ArgumentTypeError as error:
+            raise ValueError(f"must be {self.description}: {error}") from error
+        if self.choices and result not in self.choices:
+            raise ValueError(f"must be one of {self.choices}")
+        return result
+
+
+def plain(value: Any) -> Any:
+    """Resolve paths and recursively turn mappings and sequences into YAML values."""
+    if isinstance(value, Path):
+        return str(value.resolve())
+    if isinstance(value, Mapping):
+        return {key: plain(item) for key, item in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [plain(item) for item in value]
+    return value
+
+
+def write_yaml(path: Path, document: Mapping[str, Any]) -> None:
+    """Write a configuration with sorted keys and readable Unicode values."""
+    path.write_text(yaml.safe_dump(plain(document), sort_keys=True, allow_unicode=True))

pixtreme/train/common/metrics.py

@@ -0,0 +1,128 @@
+"""Measure RGB frames with pixtreme and aggregate PSNR and SSIM by frame count."""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+
+import torch
+from pixtreme.core import Frame
+
+
+@dataclass(frozen=True)
+class FrameMetrics:
+    """PSNR and SSIM for one indexed frame."""
+
+    index: int
+    psnr: float
+    ssim: float
+
+
+@dataclass(frozen=True)
+class MaterialMetrics:
+    """Per-frame measurements and their sums and means for one material."""
+
+    path: str
+    frames: tuple[FrameMetrics, ...]
+
+    @property
+    def frame_count(self) -> int:
+        return len(self.frames)
+
+    @property
+    def psnr_sum(self) -> float:
+        return math.fsum(frame.psnr for frame in self.frames)
+
+    @property
+    def psnr_mean(self) -> float:
+        return self.psnr_sum / self.frame_count
+
+    @property
+    def ssim_sum(self) -> float:
+        return math.fsum(frame.ssim for frame in self.frames)
+
+    @property
+    def ssim_mean(self) -> float:
+        return self.ssim_sum / self.frame_count
+
+    def to_dict(self) -> dict[str, object]:
+        return {
+            "path": self.path,
+            "frames": [
+                {"index": frame.index, "psnr": json_metric(frame.psnr), "ssim": json_metric(frame.ssim)}
+                for frame in self.frames
+            ],
+            "frame_count": self.frame_count,
+            "psnr_sum": json_metric(self.psnr_sum),
+            "psnr_mean": json_metric(self.psnr_mean),
+            "ssim_sum": json_metric(self.ssim_sum),
+            "ssim_mean": json_metric(self.ssim_mean),
+        }
+
+
+@dataclass(frozen=True)
+class EvaluationMetrics:
+    """Frame-weighted pixtreme PSNR and SSIM, with optional interpolation results."""
+
+    materials: tuple[MaterialMetrics, ...]
+    baselines: dict[str, EvaluationMetrics] = field(default_factory=dict)
+
+    @property
+    def frame_count(self) -> int:
+        return sum(material.frame_count for material in self.materials)
+
+    @property
+    def psnr_sum(self) -> float:
+        return math.fsum(material.psnr_sum for material in self.materials)
+
+    @property
+    def psnr_mean(self) -> float:
+        return self.psnr_sum / self.frame_count
+
+    @property
+    def ssim_sum(self) -> float:
+        return math.fsum(material.ssim_sum for material in self.materials)
+
+    @property
+    def ssim_mean(self) -> float:
+        return self.ssim_sum / self.frame_count
+
+    def to_dict(self) -> dict[str, object]:
+        result: dict[str, object] = {
+            "materials": [material.to_dict() for material in self.materials],
+            "frame_count": self.frame_count,
+            "psnr_sum": json_metric(self.psnr_sum),
+            "psnr_mean": json_metric(self.psnr_mean),
+            "ssim_sum": json_metric(self.ssim_sum),
+            "ssim_mean": json_metric(self.ssim_mean),
+        }
+
+        if self.baselines:
+            result["baselines"] = {name: value.to_dict() for name, value in self.baselines.items()}
+        return result
+
+
+def json_metric(value: float) -> float | str:
+    """Represent positive infinity as a string accepted by strict JSON encoders."""
+    return "Infinity" if value == math.inf else value
+
+
+def rgb_frame(image: torch.Tensor) -> Frame:
+    """Convert a CHW RGB CPU tensor into a Rec.709 Frame stored on the CUDA device."""
+    import cupy as cp
+    from pixtreme.io import from_array
+
+    return from_array(
+        cp.asarray(image.permute(1, 2, 0).contiguous().numpy()), colorspace="Rec.709", gamma="Rec.709", channels="RGB"
+    )
+
+
+def measure(index: int, reference: Frame, candidate: Frame) -> FrameMetrics:
+    """Measure RGB PSNR and SSIM with pixtreme using a value range of 1.0."""
+    from pixtreme.metrics import psnr, ssim
+
+    return FrameMetrics(
+        index,
+        float(psnr(reference, candidate, data_range=1.0)),
+        float(ssim(reference, candidate, data_range=1.0)),
+    )

pixtreme_train-0.1.0.dist-info/METADATA

@@ -0,0 +1,50 @@
+Metadata-Version: 2.5
+Name: pixtreme-train
+Version: 0.1.0
+Summary: Shared training components for pixtreme model trainers
+Project-URL: repository, https://github.com/sync-dev-org/pixtreme-train
+Author-email: minamik <mia@sync.dev>
+License-Expression: MIT
+License-File: LICENSE
+Classifier: Development Status :: 2 - Pre-Alpha
+Classifier: Environment :: GPU :: NVIDIA CUDA :: 13
+Classifier: Intended Audience :: Developers
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development :: Libraries
+Requires-Python: >=3.12
+Requires-Dist: numpy
+Requires-Dist: pixtreme>=2.0.0
+Requires-Dist: pyyaml>=6
+Requires-Dist: torch>=2.11
+Description-Content-Type: text/markdown
+
+# pixtreme.train.common
+
+The pure Python distribution `pixtreme-train` provides the import package
+`pixtreme.train.common`: training components that do not depend on a particular
+model, shared by the pixtreme trainer modules.
+
+The wheel owns only the `pixtreme/train/common/` subtree. The namespace root
+`pixtreme/__init__.py` is owned by [pixtreme](https://github.com/sync-dev-org/pixtreme),
+and each trainer module owns its own `pixtreme/train/<module>/` subtree.
+
+## Installation
+
+Use Linux x86_64 with CPython 3.12 or later. Install the distribution from PyPI:
+
+```console
+pip install pixtreme-train
+```
+
+For a project managed with uv:
+
+```console
+uv add pixtreme-train
+```
+
+## License
+
+MIT. See [LICENSE](LICENSE).

pixtreme_train-0.1.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+pixtreme/train/common/__init__.py,sha256=Ot1tWhucFZij1E4ZkLrhRYCGcuT6Aj--EEMU71mTEDA,484
+pixtreme/train/common/checkpoint.py,sha256=LWPBkHLXWR4XJKQoLllO6T2zUIfo0mt_Lp3EgDtZU-I,1092
+pixtreme/train/common/configuration.py,sha256=NLPYr3mJu02WBlBY6Sm2knx2LKNwa89DfM_E7necbow,2017
+pixtreme/train/common/metrics.py,sha256=0l7ntRk9C9Tl0xMVbJVRBpuJzMdz4-Qa6ba86OPeB9E,3867
+pixtreme_train-0.1.0.dist-info/METADATA,sha256=8h8Q7T9P3pZqSlsl_vXjy_6ZJ6Oi4JzL5E1YeBo0CKA,1585
+pixtreme_train-0.1.0.dist-info/WHEEL,sha256=W3fkpkm7-wf9vBI5Z-7s0eWkeM-spu78I8Neb98DeEg,87
+pixtreme_train-0.1.0.dist-info/licenses/LICENSE,sha256=C3kcuSn77D1YmBPwywbEXaYrjkyWWcMvubpMt-2bi3Y,1064
+pixtreme_train-0.1.0.dist-info/RECORD,,

pixtreme_train-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.32.4
+Root-Is-Purelib: true
+Tag: py3-none-any

pixtreme_train-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 minamik
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.