File size: 6,253 Bytes
205a7af |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 |
"""Image preprocessing utilities."""
import collections.abc as collections
from pathlib import Path
from typing import Optional, Tuple
import cv2
import kornia
import numpy as np
import torch
import torchvision
from omegaconf import OmegaConf
from siclib.utils.tensor import fit_features_to_multiple
# mypy: ignore-errors
class ImagePreprocessor:
"""Preprocess images for calibration."""
default_conf = {
"resize": 320, # target edge length, None for no resizing
"edge_divisible_by": None,
"side": "short",
"interpolation": "bilinear",
"align_corners": None,
"antialias": True,
"square_crop": False,
"add_padding_mask": False,
"resize_backend": "kornia", # torchvision, kornia
}
def __init__(self, conf) -> None:
"""Initialize the image preprocessor."""
super().__init__()
default_conf = OmegaConf.create(self.default_conf)
OmegaConf.set_struct(default_conf, True)
self.conf = OmegaConf.merge(default_conf, conf)
def __call__(self, img: torch.Tensor, interpolation: Optional[str] = None) -> dict:
"""Resize and preprocess an image, return image and resize scale."""
h, w = img.shape[-2:]
size = h, w
if self.conf.square_crop:
min_size = min(h, w)
offset = (h - min_size) // 2, (w - min_size) // 2
img = img[:, offset[0] : offset[0] + min_size, offset[1] : offset[1] + min_size]
size = img.shape[-2:]
if self.conf.resize is not None:
if interpolation is None:
interpolation = self.conf.interpolation
size = self.get_new_image_size(h, w)
img = self.resize(img, size, interpolation)
scale = torch.Tensor([img.shape[-1] / w, img.shape[-2] / h]).to(img)
T = np.diag([scale[0].cpu(), scale[1].cpu(), 1])
data = {
"scales": scale,
"image_size": np.array(size[::-1]),
"transform": T,
"original_image_size": np.array([w, h]),
}
if self.conf.edge_divisible_by is not None:
# crop to make the edge divisible by a number
w_, h_ = img.shape[-1], img.shape[-2]
img, _ = fit_features_to_multiple(img, self.conf.edge_divisible_by, crop=True)
crop_pad = torch.Tensor([img.shape[-1] - w_, img.shape[-2] - h_]).to(img)
data["crop_pad"] = crop_pad
data["image_size"] = np.array([img.shape[-1], img.shape[-2]])
data["image"] = img
return data
def resize(self, img: torch.Tensor, size: Tuple[int, int], interpolation: str) -> torch.Tensor:
"""Resize an image using the specified backend."""
if self.conf.resize_backend == "kornia":
return kornia.geometry.transform.resize(
img,
size,
side=self.conf.side,
antialias=self.conf.antialias,
align_corners=self.conf.align_corners,
interpolation=interpolation,
)
elif self.conf.resize_backend == "torchvision":
return torchvision.transforms.Resize(size, antialias=self.conf.antialias)(img)
else:
raise ValueError(f"{self.conf.resize_backend} not implemented.")
def load_image(self, image_path: Path) -> dict:
"""Load an image from a path and preprocess it."""
return self(load_image(image_path))
def get_new_image_size(self, h: int, w: int) -> Tuple[int, int]:
"""Get the new image size after resizing."""
side = self.conf.side
if isinstance(self.conf.resize, collections.Iterable):
assert len(self.conf.resize) == 2
return tuple(self.conf.resize)
side_size = self.conf.resize
aspect_ratio = w / h
if side not in ("short", "long", "vert", "horz"):
raise ValueError(
f"side can be one of 'short', 'long', 'vert', and 'horz'. Got '{side}'"
)
return (
(side_size, int(side_size * aspect_ratio))
if side == "vert" or (side != "horz" and (side == "short") ^ (aspect_ratio < 1.0))
else (int(side_size / aspect_ratio), side_size)
)
def numpy_image_to_torch(image: np.ndarray) -> torch.Tensor:
"""Normalize the image tensor and reorder the dimensions."""
if image.ndim == 3:
image = image.transpose((2, 0, 1)) # HxWxC to CxHxW
elif image.ndim == 2:
image = image[None] # add channel axis
else:
raise ValueError(f"Not an image: {image.shape}")
return torch.tensor(image / 255.0, dtype=torch.float)
def torch_image_to_numpy(image: torch.Tensor) -> np.ndarray:
"""Normalize and reorder the dimensions of an image tensor."""
if image.ndim == 3:
image = image.permute((1, 2, 0)) # CxHxW to HxWxC
elif image.ndim == 2:
image = image[None] # add channel axis
else:
raise ValueError(f"Not an image: {image.shape}")
return (image.cpu().detach().numpy() * 255).astype(np.uint8)
def read_image(path: Path, grayscale: bool = False) -> np.ndarray:
"""Read an image from path as RGB or grayscale."""
if not Path(path).exists():
raise FileNotFoundError(f"No image at path {path}.")
mode = cv2.IMREAD_GRAYSCALE if grayscale else cv2.IMREAD_COLOR
image = cv2.imread(str(path), mode)
if image is None:
raise IOError(f"Could not read image at {path}.")
if not grayscale:
image = image[..., ::-1]
return image
def write_image(img: torch.Tensor, path: Path):
"""Write an image tensor to a file."""
img = torch_image_to_numpy(img) if isinstance(img, torch.Tensor) else img
cv2.imwrite(str(path), img[..., ::-1])
def load_image(path: Path, grayscale: bool = False, return_tensor: bool = True) -> torch.Tensor:
"""Load an image from a path and return as a tensor."""
image = read_image(path, grayscale=grayscale)
if return_tensor:
return numpy_image_to_torch(image)
assert image.ndim in [2, 3], f"Not an image: {image.shape}"
image = image[None] if image.ndim == 2 else image
return torch.tensor(image.copy(), dtype=torch.uint8)
|