def to_image(t):
"""
:param t: tensor or np.ndarray, may be of shape NCHW / CHW with C=1 or 3 / HW, dtype float32 or uint8. If float32:
must be in [0, 1]
:return: HW3 uint8 np.ndarray
"""
if not isinstance(t, np.ndarray):
t = pe.tensor_to_np(t)
# - t is numpy array
if t.ndim == 4:
# - t has batch dimension, only use first
t = t[0, ...]
elif t.ndim == 2:
t = np.expand_dims(t, 0) # Now 1HW
assert_exc(t.ndim == 3, 'Invalid shape: {}'.format(t.shape))
# - t is 3 dimensional CHW numpy array
if t.dtype != np.uint8:
assert_exc(t.dtype == np.float32, 'Expected either uint8 or float32, got {}'.format(t.dtype))
_check_range(t, 0, 1)
t = (t * 255.).astype(np.uint8)
# - t is uint8 numpy array
num_channels = t.shape[0]
if num_channels == 3:
t = np.transpose(t, (1, 2, 0))
elif num_channels == 1:
t = np.stack([t[0, :, :] for _ in range(3)], -1)
else:
raise ValueError('Expected CHW, got {}'.format(t.shape))
assert_exc(t.ndim == 3 and t.shape[2] == 3, str(t.shape))
# - t is uint8 numpy array of shape HW3
return t
def add(self, v):
v = pe.tensor_to_np(v)
num_values = np.prod(v.shape)
if self._buffer is None:
print(f'Creating {v.dtype} buffer for {self._name}: {self._buffer_size}x{num_values}')
self._buffer = np.zeros((self._buffer_size, num_values), dtype=v.dtype)
assert_exc(self._buffer.shape[1] == num_values, (self._buffer.shape, v.shape, num_values), BufferSizeMismatch)
self._buffer[self._idx, :] = v.flatten()
self._idx = (self._idx + 1) % self._buffer_size
self._filled_idx = min(self._filled_idx + 1, self._buffer_size)
def _write_img(decoded, png_out_p):
"""
:param decoded: 1CHW tensor
:param png_out_p: str
"""
assert decoded.shape[0] == 1 and decoded.shape[1] == 3, decoded.shape
img = pe.tensor_to_np(decoded.squeeze(0)) # CHW
img = img.transpose(1, 2, 0).astype(np.uint8) # Make HW3
img = Image.fromarray(img)
img.save(png_out_p)
def get_p_y(y):
"""
:param y: NLCHW float, logits
:return: L dimensional vector p
"""
Ldim = 1
L = y.shape[Ldim]
y = y.detach()
p = F.softmax(y, dim=Ldim)
p = p.transpose(Ldim, -1)
p = p.contiguous().view(-1, L) # nL
p = torch.mean(p, dim=0) # L
return pe.tensor_to_np(p)