Spaces:
Running
Running
| # 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 | |
| def resize_image(pil_image, target_size): | |
| """ | |
| Resize a PIL image while maintaining its aspect ratio. | |
| Args: | |
| pil_image (PIL.Image): The input image. | |
| target_size (tuple): The target size as (width, height). | |
| Returns: | |
| PIL.Image: The resized image. | |
| """ | |
| original_width, original_height = pil_image.size | |
| target_width, target_height = target_size | |
| # Calculate aspect ratios | |
| aspect_ratio = original_width / original_height | |
| target_aspect = target_width / target_height | |
| if aspect_ratio > target_aspect: | |
| # Image is wider than target, scale based on width | |
| new_width = target_width | |
| new_height = int(new_width / aspect_ratio) | |
| else: | |
| # Image is taller than target, scale based on height | |
| new_height = target_height | |
| new_width = int(new_height * aspect_ratio) | |
| # Resize the image | |
| resized_image = pil_image.resize((new_width, new_height), Image.LANCZOS) | |
| # Create a new image with the target size and paste the resized image | |
| new_image = Image.new('RGB', target_size, (0, 0, 0)) | |
| paste_x = (target_width - new_width) // 2 | |
| paste_y = (target_height - new_height) // 2 | |
| new_image.paste(resized_image, (paste_x, paste_y)) | |
| return new_image | |