86 lines
3.7 KiB
Python
86 lines
3.7 KiB
Python
"""Exact six-type values for internal sparse End exports; no pack imports.
|
|
|
|
Tensor/text payloads reuse the existing media codec. Scalars have distinct
|
|
encodings: BOOLEAN is one byte, INT is minimal signed little-endian, FLOAT is
|
|
IEEE binary64. No scalar passes through float() or a lossy JSON number parser.
|
|
"""
|
|
import math
|
|
import struct
|
|
|
|
import torch
|
|
|
|
|
|
KINDS = frozenset(("IMAGE", "MASK", "STRING", "BOOLEAN", "INT", "FLOAT"))
|
|
_SCALAR_ENCODINGS = {"BOOLEAN": "bool-v1", "INT": "int-le-v1", "FLOAT": "float64-le-v1"}
|
|
|
|
|
|
def value_kind(value):
|
|
if isinstance(value, torch.Tensor):
|
|
if value.ndim == 4 and value.shape[-1] in (3, 4):
|
|
return "IMAGE"
|
|
if value.ndim == 3:
|
|
return "MASK"
|
|
raise ValueError("VALUE tensor requires IMAGE [B,H,W,3|4] or MASK [B,H,W]")
|
|
kind = {bool: "BOOLEAN", int: "INT", float: "FLOAT", str: "STRING"}.get(type(value))
|
|
if kind is None:
|
|
raise ValueError("VALUE requires IMAGE, MASK, STRING, INT, FLOAT or BOOLEAN")
|
|
if kind == "FLOAT" and not math.isfinite(value):
|
|
raise ValueError("VALUE numbers must be finite")
|
|
return kind
|
|
|
|
|
|
def encode_value(value, *, max_bytes):
|
|
from .codec import encode_payload
|
|
kind = value_kind(value)
|
|
if kind in ("IMAGE", "MASK", "STRING"):
|
|
descriptor, payload = encode_payload(kind, value, max_bytes=max_bytes)
|
|
else:
|
|
descriptor = {"encoding": _SCALAR_ENCODINGS[kind]}
|
|
size = ((value if value >= 0 else ~value).bit_length() + 8) // 8 if kind == "INT" else (8 if kind == "FLOAT" else 1)
|
|
if size > max_bytes:
|
|
raise ValueError("VALUE payload exceeds byte limit")
|
|
if kind == "BOOLEAN":
|
|
payload = b"\x01" if value else b"\x00"
|
|
elif kind == "INT":
|
|
payload = value.to_bytes(size, "little", signed=True)
|
|
else:
|
|
payload = struct.pack("<d", value)
|
|
return {"encoding": "value-v1", "type": kind, "descriptor": descriptor}, payload
|
|
|
|
|
|
def decode_value(descriptor, payload, *, max_bytes):
|
|
from .codec import decode_payload
|
|
if (set(descriptor) != {"encoding", "type", "descriptor"}
|
|
or descriptor["encoding"] != "value-v1"
|
|
or type(descriptor["type"]) is not str or descriptor["type"] not in KINDS):
|
|
raise ValueError("Invalid VALUE descriptor")
|
|
kind, inner = descriptor["type"], descriptor["descriptor"]
|
|
if kind in ("IMAGE", "MASK", "STRING"):
|
|
# Generic values always stay exact even if End.images is compressed.
|
|
if not isinstance(inner, dict) or inner.get("encoding") != (
|
|
"utf8-v1" if kind == "STRING" else "tensor-le-v1"):
|
|
raise ValueError("VALUE requires exact typed payloads")
|
|
value = decode_payload(kind, inner, payload, max_bytes=max_bytes)
|
|
if value_kind(value) != kind:
|
|
raise ValueError("VALUE tensor/type mismatch")
|
|
return value
|
|
if inner != {"encoding": _SCALAR_ENCODINGS[kind]}:
|
|
raise ValueError("Invalid VALUE scalar encoding")
|
|
if kind == "BOOLEAN":
|
|
if payload not in (b"\x00", b"\x01"):
|
|
raise ValueError("Invalid BOOLEAN payload")
|
|
return payload == b"\x01"
|
|
if kind == "FLOAT":
|
|
if len(payload) != 8:
|
|
raise ValueError("Invalid FLOAT payload length")
|
|
value = struct.unpack("<d", payload)[0]
|
|
if not math.isfinite(value):
|
|
raise ValueError("VALUE numbers must be finite")
|
|
return value
|
|
if not payload:
|
|
raise ValueError("Invalid INT payload length")
|
|
value = int.from_bytes(payload, "little", signed=True)
|
|
if ((value if value >= 0 else ~value).bit_length() + 8) // 8 != len(payload):
|
|
raise ValueError("Invalid INT payload length")
|
|
return value
|