File size: 6,237 Bytes
6debdcc | 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 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | """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
|