Spaces:
Running
Running
File size: 1,277 Bytes
94f04b7 |
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 |
# source: huggingface: fashn-ai/sapiens-body-part-segmentation
import colorsys
import matplotlib.colors as mcolors
import numpy as np
from PIL import Image
def get_palette(num_cls):
palette = [0] * (256 * 3)
palette[0:3] = [0, 0, 0]
for j in range(1, num_cls):
hue = (j - 1) / (num_cls - 1)
saturation = 1.0
value = 1.0 if j % 2 == 0 else 0.5
rgb = colorsys.hsv_to_rgb(hue, saturation, value)
r, g, b = [int(x * 255) for x in rgb]
palette[j * 3 : j * 3 + 3] = [r, g, b]
return palette
def create_colormap(palette):
colormap = np.array(palette).reshape(-1, 3) / 255.0
return mcolors.ListedColormap(colormap)
def visualize_mask_with_overlay(img: Image.Image, mask: Image.Image, labels_to_ids: dict[str, int], alpha=0.5):
img_np = np.array(img.convert("RGB"))
mask_np = np.array(mask)
num_cls = len(labels_to_ids)
palette = get_palette(num_cls)
colormap = create_colormap(palette)
overlay = np.zeros((*mask_np.shape, 3), dtype=np.uint8)
for label, idx in labels_to_ids.items():
if idx != 0:
overlay[mask_np == idx] = np.array(colormap(idx)[:3]) * 255
blended = Image.fromarray(np.uint8(img_np * (1 - alpha) + overlay * alpha))
return blended |