54 lines
1.4 KiB
Python
54 lines
1.4 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from PIL import Image, ImageOps
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ShapeOnlyTensor:
|
|
shape: tuple[int, ...]
|
|
|
|
|
|
def _torch_module() -> Any | None:
|
|
try:
|
|
import torch
|
|
except ImportError:
|
|
return None
|
|
return torch
|
|
|
|
|
|
def _image_tensor(image: Image.Image) -> Any:
|
|
torch = _torch_module()
|
|
width, height = image.size
|
|
if torch is None:
|
|
return ShapeOnlyTensor((1, height, width, 3))
|
|
|
|
rgb_image = image.convert("RGB")
|
|
data = list(rgb_image.getdata())
|
|
return torch.tensor(data, dtype=torch.float32).reshape(1, height, width, 3) / 255.0
|
|
|
|
|
|
def _mask_tensor(image: Image.Image) -> Any:
|
|
torch = _torch_module()
|
|
width, height = image.size
|
|
if torch is None:
|
|
return ShapeOnlyTensor((1, height, width))
|
|
|
|
if "A" in image.getbands():
|
|
alpha = image.getchannel("A")
|
|
mask_data = [1.0 - (value / 255.0) for value in alpha.getdata()]
|
|
else:
|
|
mask_data = [0.0] * (width * height)
|
|
return torch.tensor(mask_data, dtype=torch.float32).reshape(1, height, width)
|
|
|
|
|
|
def load_image_and_mask(image_path: str | Path) -> tuple[Any, Any, int, int]:
|
|
path = Path(image_path)
|
|
with Image.open(path) as loaded:
|
|
image = ImageOps.exif_transpose(loaded)
|
|
width, height = image.size
|
|
return _image_tensor(image), _mask_tensor(image), width, height
|