scripts/pixelmodel_robustness/conditions.py

"""Weight image codec conditions used by the study."""

from __future__ import annotations

from collections.abc import Callable, Mapping
from dataclasses import dataclass
from io import BytesIO

import numpy as np
from PIL import Image


@dataclass(frozen=True)
class Condition:
    """一つの画像変換条件とその capability を表す。"""

    name: str
    family: str
    level: float | int | str | None
    capability: bool
    transform: Callable[[Image.Image, np.random.Generator], bytes]
    description: str
    metadata: Mapping[str, object] | None = None

    def apply(self, image: Image.Image, seed: int = 0) -> bytes:
        """固定 seed で条件を適用し、PNG/JPEG等の bytes を返す。"""
        return self.transform(image.copy(), np.random.default_rng(seed))


def _encode(image: Image.Image, fmt: str, **kwargs: object) -> bytes:
    stream = BytesIO()
    image.convert("RGB").save(stream, format=fmt, **kwargs)
    return stream.getvalue()


def _pixel_png(array: np.ndarray) -> bytes:
    """RGB arrayを有効なPNG streamとして保存する。"""
    return _encode(Image.fromarray(np.asarray(array, dtype=np.uint8), "RGB"), "PNG")


def _noise(image: Image.Image, rng: np.random.Generator, kind: str, level: float) -> bytes:
    array = np.asarray(image.convert("RGB"), dtype=np.int16).copy()
    if kind == "gaussian":
        array += rng.normal(0, level, array.shape).round().astype(np.int16)
    elif kind == "salt_pepper":
        mask = rng.random(array.shape[:2]) < level
        salt = rng.random(mask.sum()) < 0.5
        pixels = array[mask]
        pixels[:] = np.where(salt[:, None], 255, 0)
        array[mask] = pixels
    elif kind == "pixel_dropout":
        mask = rng.random(array.shape[:2]) < level
        array[mask] = 0
    elif kind == "bit_flip":
        mask = rng.random(array.shape[:2] + (2,)) < level
        bits = rng.integers(0, 8, mask.shape, dtype=np.int16)
        channels = array[..., :2]
        channels[mask] ^= (1 << bits[mask])
    return _encode(Image.fromarray(np.clip(array, 0, 255).astype(np.uint8)), "PNG")


def _bit_flip(image: Image.Image, rng: np.random.Generator, level: float) -> bytes:
    """R/G channelだけをbit flipし、Bを不変にしたPNGを返す。"""
    return _noise(image, rng, "bit_flip", level)


def _byte_corrupt(image: Image.Image, rng: np.random.Generator, level: float, byte: str = "low") -> bytes:
    """PNG streamを壊さず、RGB arrayの指定byteのみ非重複indexで変更する。"""
    array = np.asarray(image.convert("RGB"), dtype=np.uint8).copy()
    channel = 0 if byte == "high" else 1 if byte == "low" else None
    if channel is None:
        raise ValueError("byte must be high or low")
    flat = array[..., channel].reshape(-1)
    count = max(1, min(len(flat), round(len(flat) * level)))
    indices = rng.choice(len(flat), size=count, replace=False)
    old = flat[indices].copy()
    replacement = rng.integers(0, 256, size=count, dtype=np.uint8)
    replacement[replacement == old] ^= np.uint8(1)
    flat[indices] = replacement
    return _pixel_png(array)


def _jpeg(image: Image.Image, quality: int, subsampling: int | str = 0) -> bytes:
    return _encode(image, "JPEG", quality=quality, subsampling=subsampling, optimize=False)


def build_registry(profile: str = "full") -> list[Condition]:
    """仕様で固定された codec/corruption registry を返す。"""
    qualities = [100, 95, 90, 85, 80, 70, 60, 50, 40, 30, 20, 10]
    registry: list[Condition] = [Condition("png_baseline", "png", None, True, lambda i, r: _encode(i, "PNG"), "lossless PNG", {})]
    if profile == "quick":
        qualities = [100, 95, 80, 50, 10]
        gaussian = [2.0]
        salt_pepper = [0.01]
        bit_flips = [1e-4]
        dropouts = [0.01]
        bytes_levels = [1e-4]
    else:
        gaussian = [1.0, 2.0, 5.0, 10.0]
        salt_pepper = [0.001, 0.01, 0.05, 0.1]
        bit_flips = [1e-6, 1e-5, 1e-4, 1e-3]
        dropouts = [0.001, 0.01, 0.05, 0.1]
        bytes_levels = [1e-5, 1e-4, 1e-3]
    for quality in qualities:
        registry.append(Condition(f"jpeg_q{quality}", "jpeg", quality, True, lambda i, r, q=quality: _jpeg(i, q), f"JPEG quality {quality}"))
    for subsampling, label in [(0, "444"), (1, "422"), (2, "420")]:
        registry.append(Condition(f"jpeg_q90_{label}", "jpeg_subsampling", label, True, lambda i, r, s=subsampling: _jpeg(i, 90, s), f"JPEG 90 {label}"))
    try:
        Image.registered_extensions()
        webp_capable = "WEBP" in Image.SAVE
    except (AttributeError, KeyError):
        webp_capable = False
    registry.extend([
        Condition("webp_lossless", "webp", "lossless", webp_capable, lambda i, r: _encode(i, "WEBP", lossless=True), "WebP lossless"),
        Condition("webp_lossy_q80", "webp", 80, webp_capable, lambda i, r: _encode(i, "WEBP", quality=80), "WebP lossy quality 80"),
    ])
    for value in gaussian:
        registry.append(Condition(f"gaussian_{value:g}", "gaussian", value, True, lambda i, r, v=value: _noise(i, r, "gaussian", v), f"Gaussian sigma {value}"))
    for value in salt_pepper:
        registry.append(Condition(f"salt_pepper_{value:g}", "salt_pepper", value, True, lambda i, r, v=value: _noise(i, r, "salt_pepper", v), f"salt-pepper rate {value}"))
    for value in bit_flips:
        registry.append(Condition(f"random_bit_flip_{value:g}", "bit_flip", value, True, lambda i, r, v=value: _noise(i, r, "bit_flip", v), f"bit flip probability {value}"))
    for value in dropouts:
        registry.append(Condition(f"pixel_dropout_{value:g}", "pixel_dropout", value, True, lambda i, r, v=value: _noise(i, r, "pixel_dropout", v), f"pixel dropout rate {value}"))
    for value in bytes_levels:
        for level_name in ("low", "high"):
            scale = value * (0.25 if level_name == "low" else 4.0)
            registry.append(Condition(f"random_byte_{level_name}_{value:g}", "byte_corruption", value, True, lambda i, r, v=scale, b=level_name: _byte_corrupt(i, r, v, b), f"random byte corruption {level_name}", {"nominal_rate": value, "effective_rate": scale, "channel": "R" if level_name == "high" else "G", "rate_multiplier": 4.0 if level_name == "high" else 0.25}))
    return registry