pytorch
126 строк · 4.1 Кб
1# mypy: allow-untyped-defs
2import numpy as np
3
4
5# Functions for converting
6def figure_to_image(figures, close=True):
7"""Render matplotlib figure to numpy format.
8
9Note that this requires the ``matplotlib`` package.
10
11Args:
12figures (matplotlib.pyplot.figure or list of figures): figure or a list of figures
13close (bool): Flag to automatically close the figure
14
15Returns:
16numpy.array: image in [CHW] order
17"""
18import matplotlib.pyplot as plt
19import matplotlib.backends.backend_agg as plt_backend_agg
20
21def render_to_rgb(figure):
22canvas = plt_backend_agg.FigureCanvasAgg(figure)
23canvas.draw()
24data: np.ndarray = np.frombuffer(canvas.buffer_rgba(), dtype=np.uint8)
25w, h = figure.canvas.get_width_height()
26image_hwc = data.reshape([h, w, 4])[:, :, 0:3]
27image_chw = np.moveaxis(image_hwc, source=2, destination=0)
28if close:
29plt.close(figure)
30return image_chw
31
32if isinstance(figures, list):
33images = [render_to_rgb(figure) for figure in figures]
34return np.stack(images)
35else:
36image = render_to_rgb(figures)
37return image
38
39
40def _prepare_video(V):
41"""
42Convert a 5D tensor into 4D tensor.
43
44Convesrion is done from [batchsize, time(frame), channel(color), height, width] (5D tensor)
45to [time(frame), new_width, new_height, channel] (4D tensor).
46
47A batch of images are spreaded to a grid, which forms a frame.
48e.g. Video with batchsize 16 will have a 4x4 grid.
49"""
50b, t, c, h, w = V.shape
51
52if V.dtype == np.uint8:
53V = np.float32(V) / 255.0
54
55def is_power2(num):
56return num != 0 and ((num & (num - 1)) == 0)
57
58# pad to nearest power of 2, all at once
59if not is_power2(V.shape[0]):
60len_addition = int(2 ** V.shape[0].bit_length() - V.shape[0])
61V = np.concatenate((V, np.zeros(shape=(len_addition, t, c, h, w))), axis=0)
62
63n_rows = 2 ** ((b.bit_length() - 1) // 2)
64n_cols = V.shape[0] // n_rows
65
66V = np.reshape(V, newshape=(n_rows, n_cols, t, c, h, w))
67V = np.transpose(V, axes=(2, 0, 4, 1, 5, 3))
68V = np.reshape(V, newshape=(t, n_rows * h, n_cols * w, c))
69
70return V
71
72
73def make_grid(I, ncols=8):
74# I: N1HW or N3HW
75assert isinstance(I, np.ndarray), "plugin error, should pass numpy array here"
76if I.shape[1] == 1:
77I = np.concatenate([I, I, I], 1)
78assert I.ndim == 4 and I.shape[1] == 3
79nimg = I.shape[0]
80H = I.shape[2]
81W = I.shape[3]
82ncols = min(nimg, ncols)
83nrows = int(np.ceil(float(nimg) / ncols))
84canvas = np.zeros((3, H * nrows, W * ncols), dtype=I.dtype)
85i = 0
86for y in range(nrows):
87for x in range(ncols):
88if i >= nimg:
89break
90canvas[:, y * H : (y + 1) * H, x * W : (x + 1) * W] = I[i]
91i = i + 1
92return canvas
93
94# if modality == 'IMG':
95# if x.dtype == np.uint8:
96# x = x.astype(np.float32) / 255.0
97
98
99def convert_to_HWC(tensor, input_format): # tensor: numpy array
100assert len(set(input_format)) == len(
101input_format
102), f"You can not use the same dimension shordhand twice. input_format: {input_format}"
103assert len(tensor.shape) == len(
104input_format
105), f"size of input tensor and input format are different. \
106tensor shape: {tensor.shape}, input_format: {input_format}"
107input_format = input_format.upper()
108
109if len(input_format) == 4:
110index = [input_format.find(c) for c in "NCHW"]
111tensor_NCHW = tensor.transpose(index)
112tensor_CHW = make_grid(tensor_NCHW)
113return tensor_CHW.transpose(1, 2, 0)
114
115if len(input_format) == 3:
116index = [input_format.find(c) for c in "HWC"]
117tensor_HWC = tensor.transpose(index)
118if tensor_HWC.shape[2] == 1:
119tensor_HWC = np.concatenate([tensor_HWC, tensor_HWC, tensor_HWC], 2)
120return tensor_HWC
121
122if len(input_format) == 2:
123index = [input_format.find(c) for c in "HW"]
124tensor = tensor.transpose(index)
125tensor = np.stack([tensor, tensor, tensor], 2)
126return tensor
127