39 lines
1.3 KiB
Python
39 lines
1.3 KiB
Python
|
|
from __future__ import annotations
|
||
|
|
from dataclasses import dataclass
|
||
|
|
import numpy as np
|
||
|
|
from PIL import Image
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True)
|
||
|
|
class RgbImage512:
|
||
|
|
"""RGB-изображение 512x512."""
|
||
|
|
|
||
|
|
data: np.ndarray # uint8, shape (512, 512, 3)
|
||
|
|
|
||
|
|
def __post_init__(self):
|
||
|
|
arr = np.asarray(self.data, dtype=np.uint8)
|
||
|
|
if arr.shape != (512, 512, 3):
|
||
|
|
raise ValueError("Изображение должно быть 512x512x3 (RGB).")
|
||
|
|
object.__setattr__(self, "data", arr)
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def from_file(cls, path: str) -> "RgbImage512":
|
||
|
|
"""Загрузить из файла."""
|
||
|
|
img = Image.open(path).convert("RGB")
|
||
|
|
if img.size != (512, 512):
|
||
|
|
raise ValueError("Размер изображения должен быть 512x512.")
|
||
|
|
return cls(np.array(img, dtype=np.uint8))
|
||
|
|
|
||
|
|
@classmethod
|
||
|
|
def from_array(cls, arr: np.ndarray) -> "RgbImage512":
|
||
|
|
"""Создать из массива 512x512x3."""
|
||
|
|
return cls(arr)
|
||
|
|
|
||
|
|
def to_array(self) -> np.ndarray:
|
||
|
|
"""Вернуть массив uint8 512x512x3."""
|
||
|
|
return self.data
|
||
|
|
|
||
|
|
def save(self, path: str) -> None:
|
||
|
|
"""Сохранить в файл."""
|
||
|
|
Image.fromarray(self.data, mode="RGB").save(path)
|