Spaces:
Runtime error
Runtime error
File size: 1,348 Bytes
c19ca42 |
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 |
from enum import Enum
import numpy as np
from ...utils.utils import get_h_w_c
class HeightSource(Enum):
AVERAGE_RGB = 0
MAX_RGB = 1
# 1 - ((1-r) * (1-g) * (1-b))
SCREEN_RGB = 2
RED = 3
GREEN = 4
BLUE = 5
ALPHA = 6
def get_height_map(img: np.ndarray, source: HeightSource) -> np.ndarray:
"""
Converts the given color/grayscale image to a height map.
"""
h, w, c = get_h_w_c(img)
assert c in (1, 3, 4), "Only grayscale, RGB, and RGBA images are supported"
if source == HeightSource.ALPHA:
if c < 4:
return np.ones((h, w), dtype=np.float32)
return img[:, :, 3]
if c == 1:
if source == HeightSource.SCREEN_RGB:
x = 1 - img
return 1 - x * x * x
return img
r = img[:, :, 2]
g = img[:, :, 1]
b = img[:, :, 0]
if source == HeightSource.RED:
return r
elif source == HeightSource.GREEN:
return g
elif source == HeightSource.BLUE:
return b
elif source == HeightSource.MAX_RGB:
return np.maximum(np.maximum(r, g), b)
elif source == HeightSource.AVERAGE_RGB:
return (r + g + b) / 3
elif source == HeightSource.SCREEN_RGB:
return 1 - ((1 - r) * (1 - g) * (1 - b))
else:
raise AssertionError(f"Invalid height source {source}.")
|