| 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): |
| |
| 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): |
| |
| 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 |
|
|
