File size: 6,909 Bytes
ace9173 |
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 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 |
"""Image preprocessing utilities."""
"""Adapted from https://github.com/cvg/GeoCalib"""
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 PIL import Image
from tensor import fit_features_to_multiple
# mypy: ignore-errors
class ImagePreprocessor:
"""Preprocess images for calibration."""
default_conf = {
"resize": None, # target edge length (320), 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 == "PIL":
device = img.device
imgs = []
has_batch_dim = img.ndim == 4
img = img if has_batch_dim else img[None]
for im in img:
im = (im.permute(1, 2, 0) * 255).cpu().numpy().astype(np.uint8)
im = Image.fromarray(im).resize(size[::-1], Image.BILINEAR)
im = torch.tensor(np.array(im)).permute(2, 0, 1) / 255.0
imgs.append(im.to(device))
imgs = torch.stack(imgs)
return imgs if has_batch_dim else imgs[0]
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)
|