315 lines
11 KiB
Python
315 lines
11 KiB
Python
import os
|
|
import re
|
|
import shutil
|
|
|
|
import cv2
|
|
import imageio
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
import torch
|
|
from matplotlib import cm
|
|
from matplotlib.colors import LinearSegmentedColormap
|
|
from PIL import Image, ImageDraw
|
|
|
|
from .typing import *
|
|
|
|
class SaverMixin:
|
|
_save_dir: Optional[str] = None
|
|
|
|
def set_save_dir(self, save_dir: str):
|
|
self._save_dir = save_dir
|
|
|
|
def get_save_dir(self):
|
|
if self._save_dir is None:
|
|
raise ValueError("Save dir is not set")
|
|
return self._save_dir
|
|
|
|
def convert_data(self, data):
|
|
if data is None:
|
|
return None
|
|
elif isinstance(data, np.ndarray):
|
|
return data
|
|
elif isinstance(data, torch.Tensor):
|
|
return data.detach().cpu().numpy()
|
|
elif isinstance(data, list):
|
|
return [self.convert_data(d) for d in data]
|
|
elif isinstance(data, dict):
|
|
return {k: self.convert_data(v) for k, v in data.items()}
|
|
else:
|
|
raise TypeError(
|
|
"Data must be in type numpy.ndarray, torch.Tensor, list or dict, getting",
|
|
type(data),
|
|
)
|
|
|
|
def get_save_path(self, filename):
|
|
save_path = os.path.join(self.get_save_dir(), filename)
|
|
os.makedirs(os.path.dirname(save_path), exist_ok=True)
|
|
return save_path
|
|
|
|
DEFAULT_RGB_KWARGS = {"data_format": "HWC", "data_range": (0, 1)}
|
|
DEFAULT_UV_KWARGS = {
|
|
"data_format": "HWC",
|
|
"data_range": (0, 1),
|
|
"cmap": "checkerboard",
|
|
}
|
|
DEFAULT_GRAYSCALE_KWARGS = {"data_range": None, "cmap": "jet"}
|
|
DEFAULT_GRID_KWARGS = {"align": "max"}
|
|
|
|
def get_rgb_image_(self, img, data_format, data_range, rgba=False):
|
|
img = self.convert_data(img)
|
|
assert data_format in ["CHW", "HWC"]
|
|
if data_format == "CHW":
|
|
img = img.transpose(1, 2, 0)
|
|
if img.dtype != np.uint8:
|
|
img = img.clip(min=data_range[0], max=data_range[1])
|
|
img = (
|
|
(img - data_range[0]) / (data_range[1] - data_range[0]) * 255.0
|
|
).astype(np.uint8)
|
|
nc = 4 if rgba else 3
|
|
imgs = [img[..., start : start + nc] for start in range(0, img.shape[-1], nc)]
|
|
imgs = [
|
|
img_
|
|
if img_.shape[-1] == nc
|
|
else np.concatenate(
|
|
[
|
|
img_,
|
|
np.zeros(
|
|
(img_.shape[0], img_.shape[1], nc - img_.shape[2]),
|
|
dtype=img_.dtype,
|
|
),
|
|
],
|
|
axis=-1,
|
|
)
|
|
for img_ in imgs
|
|
]
|
|
img = np.concatenate(imgs, axis=1)
|
|
if rgba:
|
|
img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGRA)
|
|
else:
|
|
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
|
|
return img
|
|
|
|
def _save_rgb_image(
|
|
self,
|
|
filename,
|
|
img,
|
|
data_format,
|
|
data_range
|
|
):
|
|
img = self.get_rgb_image_(img, data_format, data_range)
|
|
cv2.imwrite(filename, img)
|
|
|
|
def save_rgb_image(
|
|
self,
|
|
filename,
|
|
img,
|
|
data_format=DEFAULT_RGB_KWARGS["data_format"],
|
|
data_range=DEFAULT_RGB_KWARGS["data_range"],
|
|
) -> str:
|
|
save_path = self.get_save_path(filename)
|
|
self._save_rgb_image(save_path, img, data_format, data_range)
|
|
return save_path
|
|
|
|
def get_grayscale_image_(self, img, data_range, cmap):
|
|
img = self.convert_data(img)
|
|
img = np.nan_to_num(img)
|
|
if data_range is None:
|
|
img = (img - img.min()) / (img.max() - img.min())
|
|
else:
|
|
img = img.clip(data_range[0], data_range[1])
|
|
img = (img - data_range[0]) / (data_range[1] - data_range[0])
|
|
assert cmap in [None, "jet", "magma", "spectral"]
|
|
if cmap == None:
|
|
img = (img * 255.0).astype(np.uint8)
|
|
img = np.repeat(img[..., None], 3, axis=2)
|
|
elif cmap == "jet":
|
|
img = (img * 255.0).astype(np.uint8)
|
|
img = cv2.applyColorMap(img, cv2.COLORMAP_JET)
|
|
elif cmap == "magma":
|
|
img = 1.0 - img
|
|
base = cm.get_cmap("magma")
|
|
num_bins = 256
|
|
colormap = LinearSegmentedColormap.from_list(
|
|
f"{base.name}{num_bins}", base(np.linspace(0, 1, num_bins)), num_bins
|
|
)(np.linspace(0, 1, num_bins))[:, :3]
|
|
a = np.floor(img * 255.0)
|
|
b = (a + 1).clip(max=255.0)
|
|
f = img * 255.0 - a
|
|
a = a.astype(np.uint16).clip(0, 255)
|
|
b = b.astype(np.uint16).clip(0, 255)
|
|
img = colormap[a] + (colormap[b] - colormap[a]) * f[..., None]
|
|
img = (img * 255.0).astype(np.uint8)
|
|
elif cmap == "spectral":
|
|
colormap = plt.get_cmap("Spectral")
|
|
|
|
def blend_rgba(image):
|
|
image = image[..., :3] * image[..., -1:] + (
|
|
1.0 - image[..., -1:]
|
|
) # blend A to RGB
|
|
return image
|
|
|
|
img = colormap(img)
|
|
img = blend_rgba(img)
|
|
img = (img * 255).astype(np.uint8)
|
|
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
|
|
return img
|
|
|
|
def _save_grayscale_image(
|
|
self,
|
|
filename,
|
|
img,
|
|
data_range,
|
|
cmap,
|
|
):
|
|
img = self.get_grayscale_image_(img, data_range, cmap)
|
|
cv2.imwrite(filename, img)
|
|
|
|
def save_grayscale_image(
|
|
self,
|
|
filename,
|
|
img,
|
|
data_range=DEFAULT_GRAYSCALE_KWARGS["data_range"],
|
|
cmap=DEFAULT_GRAYSCALE_KWARGS["cmap"],
|
|
) -> str:
|
|
save_path = self.get_save_path(filename)
|
|
self._save_grayscale_image(save_path, img, data_range, cmap)
|
|
return save_path
|
|
|
|
def get_image_grid_(self, imgs, align):
|
|
if isinstance(imgs[0], list):
|
|
return np.concatenate(
|
|
[self.get_image_grid_(row, align) for row in imgs], axis=0
|
|
)
|
|
cols = []
|
|
for col in imgs:
|
|
assert col["type"] in ["rgb", "uv", "grayscale"]
|
|
if col["type"] == "rgb":
|
|
rgb_kwargs = self.DEFAULT_RGB_KWARGS.copy()
|
|
rgb_kwargs.update(col["kwargs"])
|
|
cols.append(self.get_rgb_image_(col["img"], **rgb_kwargs))
|
|
elif col["type"] == "uv":
|
|
uv_kwargs = self.DEFAULT_UV_KWARGS.copy()
|
|
uv_kwargs.update(col["kwargs"])
|
|
cols.append(self.get_uv_image_(col["img"], **uv_kwargs))
|
|
elif col["type"] == "grayscale":
|
|
grayscale_kwargs = self.DEFAULT_GRAYSCALE_KWARGS.copy()
|
|
grayscale_kwargs.update(col["kwargs"])
|
|
cols.append(self.get_grayscale_image_(col["img"], **grayscale_kwargs))
|
|
|
|
if align == "max":
|
|
h = max([col.shape[0] for col in cols])
|
|
w = max([col.shape[1] for col in cols])
|
|
elif align == "min":
|
|
h = min([col.shape[0] for col in cols])
|
|
w = min([col.shape[1] for col in cols])
|
|
elif isinstance(align, int):
|
|
h = align
|
|
w = align
|
|
elif (
|
|
isinstance(align, tuple)
|
|
and isinstance(align[0], int)
|
|
and isinstance(align[1], int)
|
|
):
|
|
h, w = align
|
|
else:
|
|
raise ValueError(
|
|
f"Unsupported image grid align: {align}, should be min, max, int or (int, int)"
|
|
)
|
|
|
|
for i in range(len(cols)):
|
|
if cols[i].shape[0] != h or cols[i].shape[1] != w:
|
|
cols[i] = cv2.resize(cols[i], (w, h), interpolation=cv2.INTER_LINEAR)
|
|
return np.concatenate(cols, axis=1)
|
|
|
|
def save_image_grid(
|
|
self,
|
|
filename,
|
|
imgs,
|
|
align=DEFAULT_GRID_KWARGS["align"],
|
|
texts: Optional[List[float]] = None,
|
|
):
|
|
save_path = self.get_save_path(filename)
|
|
img = self.get_image_grid_(imgs, align=align)
|
|
|
|
if texts is not None:
|
|
img = Image.fromarray(img)
|
|
draw = ImageDraw.Draw(img)
|
|
black, white = (0, 0, 0), (255, 255, 255)
|
|
for i, text in enumerate(texts):
|
|
draw.text((2, (img.size[1] // len(texts)) * i + 1), f"{text}", white)
|
|
draw.text((0, (img.size[1] // len(texts)) * i + 1), f"{text}", white)
|
|
draw.text((2, (img.size[1] // len(texts)) * i - 1), f"{text}", white)
|
|
draw.text((0, (img.size[1] // len(texts)) * i - 1), f"{text}", white)
|
|
draw.text((1, (img.size[1] // len(texts)) * i), f"{text}", black)
|
|
img = np.asarray(img)
|
|
|
|
cv2.imwrite(save_path, img)
|
|
return save_path
|
|
|
|
def save_image(self, filename, img) -> str:
|
|
save_path = self.get_save_path(filename)
|
|
img = self.convert_data(img)
|
|
assert img.dtype == np.uint8 or img.dtype == np.uint16
|
|
if img.ndim == 3 and img.shape[-1] == 3:
|
|
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
|
|
elif img.ndim == 3 and img.shape[-1] == 4:
|
|
img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGRA)
|
|
cv2.imwrite(save_path, img)
|
|
return save_path
|
|
|
|
def save_img_sequence(
|
|
self,
|
|
filename,
|
|
img_dir,
|
|
matcher,
|
|
save_format="mp4",
|
|
fps=30,
|
|
) -> str:
|
|
assert save_format in ["gif", "mp4"]
|
|
if not filename.endswith(save_format):
|
|
filename += f".{save_format}"
|
|
save_path = self.get_save_path(filename)
|
|
matcher = re.compile(matcher)
|
|
img_dir = os.path.join(self.get_save_dir(), img_dir)
|
|
imgs = []
|
|
for f in os.listdir(img_dir):
|
|
if matcher.search(f):
|
|
imgs.append(f)
|
|
imgs = sorted(imgs, key=lambda f: int(matcher.search(f).groups()[0]))
|
|
imgs = [cv2.imread(os.path.join(img_dir, f)) for f in imgs]
|
|
|
|
if save_format == "gif":
|
|
imgs = [cv2.cvtColor(i, cv2.COLOR_BGR2RGB) for i in imgs]
|
|
imageio.mimsave(save_path, imgs, fps=fps, palettesize=256)
|
|
elif save_format == "mp4":
|
|
imgs = [cv2.cvtColor(i, cv2.COLOR_BGR2RGB) for i in imgs]
|
|
imageio.mimsave(save_path, imgs, fps=fps)
|
|
return save_path
|
|
|
|
def save_img_sequences(
|
|
self,
|
|
seq_dir,
|
|
matcher,
|
|
save_format="mp4",
|
|
fps=30,
|
|
delete=True
|
|
):
|
|
seq_dir_ = os.path.join(self.get_save_dir(), seq_dir)
|
|
for f in os.listdir(seq_dir_):
|
|
img_dir_ = os.path.join(seq_dir_, f)
|
|
if not os.path.isdir(img_dir_):
|
|
continue
|
|
try:
|
|
self.save_img_sequence(
|
|
os.path.join(seq_dir, f),
|
|
os.path.join(seq_dir, f),
|
|
matcher,
|
|
save_format=save_format,
|
|
fps=fps
|
|
)
|
|
except:
|
|
raise ValueError(f"Video saving for directory {seq_dir_} failed!")
|
|
|
|
if delete:
|
|
shutil.rmtree(img_dir_)
|