File size: 2,593 Bytes
c645013 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | """File output and panoptic image utilities."""
from __future__ import annotations
import json
import os
import tempfile
from collections.abc import Iterable
from pathlib import Path
from typing import Any
def panoptic_ids_from_image(image):
"""Decode integer or RGB(A) panoptic IDs without rescaling."""
import numpy as np
array = np.asarray(image)
if array.ndim == 2:
if not np.issubdtype(array.dtype, np.integer):
raise ValueError("panoptic map must contain integer IDs")
return array.astype(np.int64)
if array.ndim == 3 and array.shape[2] in (3, 4):
rgb = array[..., :3].astype(np.int64)
return rgb[..., 0] + 256 * rgb[..., 1] + 65536 * rgb[..., 2]
raise ValueError("panoptic PNG must be integer or RGB")
def external_output(
path: str | Path,
dataset_root: str | Path,
*,
protected_roots: Iterable[str | Path] = (),
) -> Path:
"""Keep generated files outside dataset and annotation directories."""
output = Path(path).expanduser().resolve()
requested_roots = (dataset_root, *tuple(protected_roots))
immutable_roots: list[Path] = []
for raw_root in requested_roots:
root = Path(raw_root).expanduser().resolve(strict=True)
if not root.is_dir():
raise ValueError(f"Protected input root is not a directory: {root}")
if root not in immutable_roots:
immutable_roots.append(root)
for root in immutable_roots:
if output == root or root in output.parents or output in root.parents:
raise ValueError(
"Derived output must be outside immutable input roots and not an ancestor of one"
)
return output
def atomic_json(path: str | Path, value: Any) -> None:
"""Write a complete UTF-8 JSON document, then atomically replace its target."""
path = Path(path)
encoded = (
json.dumps(value, sort_keys=True, indent=2, ensure_ascii=False, allow_nan=False)
+ "\n"
)
path.parent.mkdir(parents=True, exist_ok=True)
temporary = None
try:
with tempfile.NamedTemporaryFile(
mode="w",
encoding="utf-8",
dir=path.parent,
prefix=f".{path.name}.",
suffix=".tmp",
delete=False,
) as stream:
temporary = Path(stream.name)
stream.write(encoded)
stream.flush()
os.fsync(stream.fileno())
os.replace(temporary, path)
finally:
if temporary is not None and temporary.exists():
temporary.unlink()
|