| import numpy as np |
| import torch |
| import cv2 |
|
|
|
|
| def HWC3(x): |
| assert x.dtype == np.uint8 |
| if x.ndim == 2: |
| x = x[:, :, None] |
| assert x.ndim == 3 |
| H, W, C = x.shape |
| assert C == 1 or C == 3 or C == 4 |
| if C == 3: |
| return x |
| if C == 1: |
| return np.concatenate([x, x, x], axis=2) |
| if C == 4: |
| color = x[:, :, 0:3].astype(np.float32) |
| alpha = x[:, :, 3:4].astype(np.float32) / 255.0 |
| y = color * alpha + 255.0 * (1.0 - alpha) |
| y = y.clip(0, 255).astype(np.uint8) |
| return y |
|
|
|
|
| def resize_image(input_image, resolution): |
| H, W, C = input_image.shape |
| H = float(H) |
| W = float(W) |
| k = float(resolution) / min(H, W) |
| H *= k |
| W *= k |
| H = int(np.round(H / 64.0)) * 64 |
| W = int(np.round(W / 64.0)) * 64 |
| img = cv2.resize( |
| input_image, |
| (W, H), |
| interpolation=cv2.INTER_LANCZOS4 if k > 1 else cv2.INTER_AREA, |
| ) |
| return img |
|
|
|
|
| |
| def norm_normalize(norm_out): |
| norm_x, norm_y, norm_z = torch.split(norm_out, 1, dim=0) |
| norm = torch.sqrt(norm_x**2.0 + norm_y**2.0 + norm_z**2.0) + 1e-10 |
| final_out = torch.cat([norm_x / norm, norm_y / norm, norm_z / norm], dim=0) |
| fg_mask = torch.ones_like(norm).repeat(3, 1, 1) |
| fg_mask[norm.repeat(3, 1, 1) < 0.5] = 0.0 |
| fg_mask[norm.repeat(3, 1, 1) > 1.5] = 0.0 |
|
|
| final_out[norm.repeat(3, 1, 1) < 0.5] = -1 |
| final_out[norm.repeat(3, 1, 1) > 1.5] = -1 |
| return final_out, fg_mask.bool() |
|
|
|
|
| def center_crop(input_image): |
| height, width = input_image.shape[:2] |
|
|
| if height < width: |
| min_dim = height |
| else: |
| min_dim = width |
|
|
| center_x = width // 2 |
| center_y = height // 2 |
| half_length = min_dim // 2 |
|
|
| crop_x1 = center_x - half_length |
| crop_x2 = center_x + half_length |
| crop_y1 = center_y - half_length |
| crop_y2 = center_y + half_length |
|
|
| center_cropped_image = input_image[crop_y1:crop_y2, crop_x1:crop_x2] |
|
|
| return center_cropped_image |
|
|
|
|
| def flip_x(normal): |
|
|
| if isinstance(normal, np.ndarray): |
| return normal.dot(np.array([[-1, 0, 0], [0, 1, 0], [0, 0, 1]])).astype( |
| np.float32 |
| ) |
| else: |
| trans = torch.tensor([[-1, 0, 0], [0, 1, 0], [0, 0, 1]]).float() |
| return normal @ trans |
|
|
