alivemachine
7ad72028b7
The error "clamp_scalar_cpu" not implemented for 'Half' occurs because a tensor of type float16 (or Half precision) is being passed to the .clamp() function, which does not support this operation on the CPU. To resolve this issue, you need to ensure that the tensor is in a format compatible with the .clamp() function.
99 lines
2.6 KiB
Python
99 lines
2.6 KiB
Python
from typing import List, Optional, Tuple, Union
|
|
|
|
import numpy as np
|
|
import PIL.Image
|
|
import torch
|
|
import torchvision
|
|
|
|
|
|
def denormalize(images: Union[torch.Tensor, np.ndarray]) -> torch.Tensor:
|
|
"""
|
|
Denormalize an image array to [0,1].
|
|
"""
|
|
return (images.float() / 2 + 0.5).clamp(0, 1)
|
|
|
|
|
|
def pt_to_numpy(images: torch.Tensor) -> np.ndarray:
|
|
"""
|
|
Convert a PyTorch tensor to a NumPy image.
|
|
"""
|
|
images = images.cpu().permute(0, 2, 3, 1).float().numpy()
|
|
return images
|
|
|
|
|
|
def numpy_to_pil(images: np.ndarray) -> PIL.Image.Image:
|
|
"""
|
|
Convert a NumPy image or a batch of images to a PIL image.
|
|
"""
|
|
if images.ndim == 3:
|
|
images = images[None, ...]
|
|
images = (images * 255).round().astype("uint8")
|
|
if images.shape[-1] == 1:
|
|
# special case for grayscale (single channel) images
|
|
pil_images = [
|
|
PIL.Image.fromarray(image.squeeze(), mode="L") for image in images
|
|
]
|
|
else:
|
|
pil_images = [PIL.Image.fromarray(image) for image in images]
|
|
|
|
return pil_images
|
|
|
|
|
|
def postprocess_image(
|
|
image: torch.Tensor,
|
|
output_type: str = "pil",
|
|
do_denormalize: Optional[List[bool]] = None,
|
|
) -> Union[torch.Tensor, np.ndarray, PIL.Image.Image]:
|
|
if not isinstance(image, torch.Tensor):
|
|
raise ValueError(
|
|
f"Input for postprocessing is in incorrect format: {type(image)}. We only support pytorch tensor"
|
|
)
|
|
|
|
if output_type == "latent":
|
|
return image
|
|
|
|
do_normalize_flg = True
|
|
if do_denormalize is None:
|
|
do_denormalize = [do_normalize_flg] * image.shape[0]
|
|
|
|
image = torch.stack(
|
|
[
|
|
denormalize(image[i]) if do_denormalize[i] else image[i]
|
|
for i in range(image.shape[0])
|
|
]
|
|
)
|
|
|
|
if output_type == "pt":
|
|
return image
|
|
|
|
image = pt_to_numpy(image)
|
|
|
|
if output_type == "np":
|
|
return image
|
|
|
|
if output_type == "pil":
|
|
return numpy_to_pil(image)
|
|
|
|
|
|
def process_image(
|
|
image_pil: PIL.Image.Image, range: Tuple[int, int] = (-1, 1)
|
|
) -> Tuple[torch.Tensor, PIL.Image.Image]:
|
|
image = torchvision.transforms.ToTensor()(image_pil)
|
|
r_min, r_max = range[0], range[1]
|
|
image = image * (r_max - r_min) + r_min
|
|
return image[None, ...], image_pil
|
|
|
|
|
|
def pil2tensor(image_pil: PIL.Image.Image) -> torch.Tensor:
|
|
height = image_pil.height
|
|
width = image_pil.width
|
|
imgs = []
|
|
img, _ = process_image(image_pil)
|
|
imgs.append(img)
|
|
imgs = torch.vstack(imgs)
|
|
images = torch.nn.functional.interpolate(
|
|
imgs, size=(height, width), mode="bilinear"
|
|
)
|
|
image_tensors = images.to(torch.float16)
|
|
return image_tensors
|