user
47 lines
1.7 KiB
Python
47 lines
1.7 KiB
Python
from __future__ import annotations
|
|
from dataclasses import dataclass
|
|
from typing import Iterable
|
|
import numpy as np
|
|
from PIL import Image
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Label64:
|
|
"""Бинарная метка 64x64."""
|
|
|
|
data: np.ndarray # uint8, shape (64, 64), значения {0,1}
|
|
|
|
def __post_init__(self):
|
|
arr = np.asarray(self.data, dtype=np.uint8)
|
|
if arr.shape != (64, 64):
|
|
raise ValueError("Метка должна быть 64x64.")
|
|
# Нормализация к {0,1}
|
|
arr = (arr > 0).astype(np.uint8)
|
|
object.__setattr__(self, "data", arr)
|
|
|
|
@classmethod
|
|
def from_image(cls, path: str, threshold: int = 128) -> "Label64":
|
|
"""Создать из изображения 64x64."""
|
|
img = Image.open(path).convert("L")
|
|
if img.size != (64, 64):
|
|
raise ValueError("Изображение метки должно быть 64x64.")
|
|
arr = np.array(img, dtype=np.uint8)
|
|
arr = (arr >= threshold).astype(np.uint8)
|
|
return cls(arr)
|
|
|
|
@classmethod
|
|
def from_vector(cls, vec: Iterable[int]) -> "Label64":
|
|
"""Создать из вектора длиной 4096."""
|
|
v = np.fromiter(vec, dtype=np.uint8, count=4096)
|
|
if v.size != 4096:
|
|
raise ValueError("Длина вектора должна быть 4096.")
|
|
return cls(v.reshape(64, 64))
|
|
|
|
def to_vector(self) -> np.ndarray:
|
|
"""Вернуть вектор длиной 4096."""
|
|
return self.data.reshape(-1).astype(np.uint8)
|
|
|
|
def to_image(self, path: str) -> None:
|
|
"""Сохранить PNG 64x64."""
|
|
Image.fromarray(self.data * 255, mode="L").save(path)
|