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