Spaces:
Runtime error
Runtime error
| 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 | |
| # normalize | |
| 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. | |
| fg_mask[norm.repeat(3, 1, 1) > 1.5] = 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 | |