ssitu
5d3887fad7
- Higher denoise - Larger band pass size - Make image MAE thresholds more strict: from 0.05 -> 0.01 - Use a helper function for building test image names
27 lines
990 B
Python
27 lines
990 B
Python
import pathlib
|
|
from PIL import Image
|
|
import usdu_utils
|
|
|
|
|
|
def save_image(tensor, path: pathlib.Path):
|
|
"""The goto function to save a tensor image to the sampled images directory."""
|
|
assert tensor.ndim == 3 or (tensor.ndim == 4 and tensor.shape[0] == 1), (
|
|
f"Expected a 3D tensor (H, W, C) or (1, H, W, C), got {tensor.ndim=}"
|
|
)
|
|
if tensor.ndim == 3:
|
|
tensor = tensor.unsqueeze(0)
|
|
image = usdu_utils.tensor_to_pil(tensor.cpu())
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
image.save(path, quality=75, optimize=True)
|
|
|
|
|
|
def load_image(path: pathlib.Path, device=None):
|
|
"""Load an image from disk and convert it to a tensor."""
|
|
return usdu_utils.pil_to_tensor(Image.open(path)).to(device=device)
|
|
|
|
|
|
def image_name_format(name: str, extension: str, batch_size: int = 1) -> str:
|
|
"""Helper for building image names for tests."""
|
|
batch_suffix = f"_batch{batch_size}" if batch_size > 1 else ""
|
|
return f"{name}{batch_suffix}{extension}"
|