File size: 1,878 Bytes
00c3521 |
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 |
import cv2
import numpy as np
import torch
def save_img(img, name, gamma=False):
if gamma:
img = np.power(img, 1/2.2)
img = np.clip(img, 0, 1)
img = (img * 65535).astype(np.uint16)
if img.ndim == 3:
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
cv2.imwrite(name, img)
def compute_y(img: np.ndarray) -> np.ndarray:
y = 0.299 * img[:, :, 0] + 0.587 * img[:, :, 1] + 0.114 * img[:, :, 2]
return y
def compute_raw_y(img: np.ndarray) -> np.ndarray:
g1 = img[..., 1]
g2 = img[..., 2]
ret = (g1 + g2) / 2
return ret
def pack_raw(im):
# pack Bayer image to 4 channels
if isinstance(im, torch.Tensor):
im = torch.unsqueeze(im, dim=-1)
img_shape = im.shape
H = img_shape[0]
W = img_shape[1]
out = torch.cat((im[0:H:2, 0:W:2, :], im[0:H:2, 1:W:2, :],
im[1:H:2, 1:W:2, :], im[1:H:2, 0:W:2, :]),
dim=-1)
elif isinstance(im, np.ndarray):
im = np.expand_dims(im, axis=-1)
img_shape = im.shape
H = img_shape[0]
W = img_shape[1]
out = np.concatenate((im[0:H:2, 0:W:2, :], im[0:H:2, 1:W:2, :],
im[1:H:2, 1:W:2, :], im[1:H:2, 0:W:2, :]),
axis=-1)
return out
def depack_raw(im):
# unpack 4 channels to Bayer image
img_shape = im.shape
H = img_shape[0]
W = img_shape[1]
if isinstance(im, torch.Tensor):
output = torch.zeros((H * 2, W * 2), dtype=im.dtype)
elif isinstance(im, np.ndarray):
output = np.zeros((H * 2, W * 2), dtype=im.dtype)
img_shape = output.shape
H = img_shape[0]
W = img_shape[1]
output[0:H:2, 0:W:2] = im[:, :, 0]
output[0:H:2, 1:W:2] = im[:, :, 1]
output[1:H:2, 1:W:2] = im[:, :, 2]
output[1:H:2, 0:W:2] = im[:, :, 3]
return output
|