Spaces:
Sleeping
Sleeping
File size: 3,482 Bytes
8cc8ac9 | 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 122 123 124 125 126 127 128 129 130 | """
Visualization utilities for depth map rendering.
Provides colormapped depth images, side-by-side comparisons, and overlay modes.
"""
import numpy as np
from PIL import Image
import matplotlib
import matplotlib.cm as cm
# Available colormaps for depth visualization
COLORMAPS = {
"inferno": cm.inferno,
"magma": cm.magma,
"viridis": cm.viridis,
"plasma": cm.plasma,
}
def depth_to_colormap(depth: np.ndarray, colormap: str = "inferno") -> Image.Image:
"""
Apply a scientific colormap to a normalized depth array.
Args:
depth: Normalized depth array (H, W), values in [0, 1]
colormap: Name of colormap ('inferno', 'magma', 'viridis', 'plasma')
Returns:
Colormapped PIL Image (RGB)
"""
if colormap not in COLORMAPS:
raise ValueError(f"Unknown colormap '{colormap}'. Choose from: {list(COLORMAPS.keys())}")
cmap = COLORMAPS[colormap]
colored = cmap(depth) # Returns (H, W, 4) RGBA float array
colored = (colored[:, :, :3] * 255).astype(np.uint8) # Drop alpha, convert to uint8
return Image.fromarray(colored)
def create_side_by_side(
original: Image.Image,
depth_colored: Image.Image,
gap: int = 4,
bg_color: tuple = (20, 24, 30),
) -> Image.Image:
"""
Create a side-by-side comparison of original image and depth map.
Args:
original: Original PIL Image
depth_colored: Colormapped depth PIL Image
gap: Pixel gap between images
bg_color: Background color for the gap
Returns:
Combined PIL Image
"""
# Resize depth to match original dimensions
depth_resized = depth_colored.resize(original.size, Image.LANCZOS)
w, h = original.size
canvas = Image.new("RGB", (w * 2 + gap, h), bg_color)
canvas.paste(original, (0, 0))
canvas.paste(depth_resized, (w + gap, 0))
return canvas
def create_overlay(
original: Image.Image,
depth_colored: Image.Image,
alpha: float = 0.5,
) -> Image.Image:
"""
Blend the depth map on top of the original image.
Args:
original: Original PIL Image
depth_colored: Colormapped depth PIL Image
alpha: Blend factor (0 = only original, 1 = only depth)
Returns:
Blended PIL Image
"""
depth_resized = depth_colored.resize(original.size, Image.LANCZOS)
original_rgb = original.convert("RGB")
blended = Image.blend(original_rgb, depth_resized, alpha)
return blended
def add_depth_legend(
image: Image.Image,
colormap: str = "inferno",
bar_height: int = 24,
padding: int = 12,
) -> Image.Image:
"""
Add a color legend bar at the bottom of the image.
Args:
image: Input PIL Image
colormap: Colormap name for the legend
bar_height: Height of the legend bar
padding: Padding around the bar
Returns:
Image with legend appended at the bottom
"""
w, h = image.size
total_height = h + bar_height + padding * 2
canvas = Image.new("RGB", (w, total_height), (20, 24, 30))
canvas.paste(image, (0, 0))
# Create gradient bar
cmap = COLORMAPS.get(colormap, cm.inferno)
gradient = np.linspace(0, 1, w - padding * 2).reshape(1, -1)
gradient = np.repeat(gradient, bar_height, axis=0)
colored = cmap(gradient)[:, :, :3]
colored = (colored * 255).astype(np.uint8)
bar_img = Image.fromarray(colored)
canvas.paste(bar_img, (padding, h + padding))
return canvas
|