Spaces:
Running
on
Zero
Running
on
Zero
File size: 3,686 Bytes
daa6779 b703ea2 daa6779 |
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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 |
"""
-----------------------------------------------------------------------------
Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.
NVIDIA CORPORATION and its licensors retain all intellectual property
and proprietary rights in and to this software, related documentation
and any modifications thereto. Any use, reproduction, disclosure or
distribution of this software and related documentation without an express
license agreement from NVIDIA CORPORATION is strictly prohibited.
-----------------------------------------------------------------------------
"""
from typing import Optional
import cv2
import numpy as np
def recenter_foreground(image, mask, border_ratio: float = 0.1):
"""recenter an image to leave some empty space at the image border.
Args:
image (ndarray): input image, float/uint8 [H, W, 3/4]
mask (ndarray): alpha mask, bool [H, W]
border_ratio (float, optional): border ratio, image will be resized to (1 - border_ratio). Defaults to 0.1.
Returns:
ndarray: output image, float/uint8 [H, W, 3/4]
"""
# empty foreground: just return
if mask.sum() == 0:
return image
return_int = False
if image.dtype == np.uint8:
image = image.astype(np.float32) / 255
return_int = True
H, W, C = image.shape
size = max(H, W)
# default to white bg
result = np.ones((size, size, C), dtype=np.float32)
# if rgba, set alpha to 0
if C == 4:
result[..., -1] = 0
coords = np.nonzero(mask)
x_min, x_max = coords[0].min(), coords[0].max()
y_min, y_max = coords[1].min(), coords[1].max()
h = x_max - x_min
w = y_max - y_min
desired_size = int(size * (1 - border_ratio))
scale = desired_size / max(h, w)
h2 = int(h * scale)
w2 = int(w * scale)
x2_min = (size - h2) // 2
x2_max = x2_min + h2
y2_min = (size - w2) // 2
y2_max = y2_min + w2
result[x2_min:x2_max, y2_min:y2_max] = cv2.resize(
image[x_min:x_max, y_min:y_max], (w2, h2), interpolation=cv2.INTER_AREA
)
if return_int:
result = (result * 255).astype(np.uint8)
return result
def get_random_color(index: Optional[int] = None, use_float: bool = False):
# some pleasing colors
# matplotlib.colormaps['Set3'].colors + matplotlib.colormaps['Set2'].colors + matplotlib.colormaps['Set1'].colors
palette = np.array(
[
[141, 211, 199, 255],
[255, 255, 179, 255],
[190, 186, 218, 255],
[251, 128, 114, 255],
[128, 177, 211, 255],
[253, 180, 98, 255],
[179, 222, 105, 255],
[252, 205, 229, 255],
[217, 217, 217, 255],
[188, 128, 189, 255],
[204, 235, 197, 255],
[255, 237, 111, 255],
[102, 194, 165, 255],
[252, 141, 98, 255],
[141, 160, 203, 255],
[231, 138, 195, 255],
[166, 216, 84, 255],
[255, 217, 47, 255],
[229, 196, 148, 255],
[179, 179, 179, 255],
[228, 26, 28, 255],
[55, 126, 184, 255],
[77, 175, 74, 255],
[152, 78, 163, 255],
[255, 127, 0, 255],
[255, 255, 51, 255],
[166, 86, 40, 255],
[247, 129, 191, 255],
[153, 153, 153, 255],
],
dtype=np.uint8,
)
if index is None:
index = np.random.randint(0, len(palette))
if index >= len(palette):
index = index % len(palette)
if use_float:
return palette[index].astype(np.float32) / 255
else:
return palette[index]
|