DA-2 / da2 /utils /vis.py
haodongli's picture
update
d82e7f9
import torch
from PIL import Image
import numpy as np
import matplotlib
import cv2
def concatenate_images(*image_lists):
max_width = 0
total_height = 0
row_widths = []
row_heights = []
for i, image_list in enumerate(image_lists):
width = sum(img.width for img in image_list)
max_width = max(max_width, width)
row_widths.append(width)
# Assuming all images in the list have the same height
height = image_list[0].height
total_height += height
row_heights.append(height)
new_image = Image.new('RGB', (max_width, total_height))
y_offset = 0
for i, image_list in enumerate(image_lists):
x_offset = 0
for img in image_list:
new_image.paste(img, (x_offset, y_offset))
x_offset += img.width
y_offset += row_heights[i]
return new_image
def colorize_distance(distance, mask, cmap='Spectral'):
if distance.ndim >= 3: distance = distance.squeeze()
cm = matplotlib.colormaps[cmap]
valid_distance = distance[mask]
max_distance = np.quantile(valid_distance, 0.98)
min_distance = np.quantile(valid_distance, 0.02)
distance[~mask] = max_distance
distance = ((distance - min_distance) / (max_distance - min_distance))
distance = np.clip(distance, 0, 1)
img_colored_np = cm(distance, bytes=False)[:, :, 0:3]
distance_colored = (img_colored_np * 255).astype(np.uint8)
return Image.fromarray(distance_colored)