74 lines
2.2 KiB
Python
74 lines
2.2 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 _numpy_module() -> Any | None:
|
|
try:
|
|
import numpy
|
|
except ImportError:
|
|
return None
|
|
return numpy
|
|
|
|
|
|
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")
|
|
numpy = _numpy_module()
|
|
if numpy is not None and hasattr(torch, "from_numpy"):
|
|
array = numpy.asarray(rgb_image, dtype=numpy.float32)
|
|
return torch.from_numpy(array.reshape(1, height, width, 3)) / 255.0
|
|
|
|
data = list(rgb_image.tobytes())
|
|
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))
|
|
|
|
numpy = _numpy_module()
|
|
if "A" in image.getbands():
|
|
alpha = image.getchannel("A")
|
|
if numpy is not None and hasattr(torch, "from_numpy"):
|
|
alpha_array = numpy.asarray(alpha, dtype=numpy.float32)
|
|
mask_array = (1.0 - (alpha_array / 255.0)).reshape(1, height, width)
|
|
return torch.from_numpy(mask_array)
|
|
mask_data = [1.0 - (value / 255.0) for value in alpha.tobytes()]
|
|
else:
|
|
if numpy is not None and hasattr(torch, "from_numpy"):
|
|
return torch.from_numpy(numpy.zeros((1, height, width), dtype=numpy.float32))
|
|
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
|