File size: 4,997 Bytes
8813b45 a3afcbc 724cc25 a3afcbc 724cc25 a3afcbc a6d4f44 2808776 a6d4f44 a3afcbc 4415e8d a3afcbc 35c91c9 a3afcbc 8813b45 724cc25 |
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 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 |
import trimesh
import numpy as np
from copy import deepcopy
from PIL import Image
from . import color_mappings
def line(p1, p2, c=(255,0,0), resolution=10, radius=0.05):
'''draws a 3d cylinder along the line (p1, p2)'''
# check colors
if len(c) == 1:
c = [c[0]]*4
elif len(c) == 3:
c = [*c, 255]
elif len(c) != 4:
raise ValueError(f'{c} is not a valid color (must have 1,3, or 4 elements).')
# compute length and direction of segment
p1, p2 = np.asarray(p1), np.asarray(p2)
l = np.linalg.norm(p2-p1)
direction = (p2 - p1) / l
# point z along direction of segment
T = np.eye(4)
T[:3, 2] = direction
T[:3, 3] = (p1+p2)/2
#reorthogonalize basis
b0, b1 = T[:3, 0], T[:3, 1]
if np.abs(np.dot(b0, direction)) < np.abs(np.dot(b1, direction)):
T[:3, 1] = -np.cross(b0, direction)
else:
T[:3, 0] = np.cross(b1, direction)
# generate and transform mesh
mesh = trimesh.primitives.Cylinder(radius=radius, height=l, transform=T)
# apply uniform color
mesh.visual.vertex_colors = np.ones_like(mesh.visual.vertex_colors)*c
return mesh
def show_wf(row, radius=10):
EDGE_CLASSES = ['eave',
'ridge',
'step_flashing',
'rake',
'flashing',
'post',
'valley',
'hip',
'transition_line']
return [line(a,b, radius=radius, c=color_mappings.gestalt_color_mapping[EDGE_CLASSES[cls_id]]) for (a,b), cls_id in zip(np.stack([*row['wf_vertices']])[np.stack(row['wf_edges'])], row['edge_semantics'])]
# return [line(a,b, radius=radius, c=color_mappings.edge_colors[cls_id]) for (a,b), cls_id in zip(np.stack([*row['wf_vertices']])[np.stack(row['wf_edges'])], row['edge_semantics'])]
def show_grid(edges, meshes=None, row_length=5):
'''
edges: list of list of meshes
meshes: optional corresponding list of meshes
row_length: number of meshes per row
returns trimesh.Scene()
'''
T = np.eye(4)
out = []
edges = [sum(e[1:], e[0]) for e in edges]
row_height = 1.1 * max((e.extents for e in edges), key=lambda e: e[1])[1]
col_width = 1.1 * max((e.extents for e in edges), key=lambda e: e[0])[0]
# print(row_height, col_width)
if meshes is None:
meshes = [None]*len(edges)
for i, (gt, mesh) in enumerate(zip(edges, meshes), start=0):
mesh = deepcopy(mesh)
gt = deepcopy(gt)
if i%row_length != 0:
T[0, 3] += col_width
else:
T[0, 3] = 0
T[1, 3] += row_height
# print(T[0,3]/col_width, T[2,3]/row_height)
if mesh is not None:
mesh.apply_transform(T)
out.append(mesh)
gt.apply_transform(T)
out.append(gt)
out.extend([mesh, gt])
return trimesh.Scene(out)
def visualize_order_images(row_order):
return create_image_grid(row_order['ade20k'] + row_order['gestalt'] + [visualize_depth(dm) for dm in row_order['depthcm']], num_per_row=len(row_order['ade20k']))
def create_image_grid(images, target_length=312, num_per_row=2):
# Calculate the target size for the first image
first_img = images[0]
aspect_ratio = first_img.width / first_img.height
new_width = int((target_length ** 2 * aspect_ratio) ** 0.5)
new_height = int((target_length ** 2 / aspect_ratio) ** 0.5)
# Resize the first image
resized_images = [img.resize((new_width, new_height), Image.Resampling.LANCZOS) for img in images]
# Calculate the grid size
num_rows = (len(resized_images) + num_per_row - 1) // num_per_row
grid_width = new_width * num_per_row
grid_height = new_height * num_rows
# Create a new image for the grid
grid_img = Image.new('RGB', (grid_width, grid_height))
# Paste the images into the grid
for i, img in enumerate(resized_images):
x_offset = (i % num_per_row) * new_width
y_offset = (i // num_per_row) * new_height
grid_img.paste(img, (x_offset, y_offset))
return grid_img
import matplotlib
def visualize_depth(depth, min_depth=None, max_depth=None, cmap='rainbow'):
depth = np.array(depth)
if min_depth is None:
min_depth = np.min(depth)
if max_depth is None:
max_depth = np.max(depth)
# Normalize the depth to be between 0 and 1
depth = (depth - min_depth) / (max_depth - min_depth)
depth = np.clip(depth, 0, 1)
# Use the matplotlib colormap to convert the depth to an RGB image
cmap = matplotlib.cm.get_cmap(cmap)
depth_image = (cmap(depth) * 255).astype(np.uint8)
# Convert the depth image to a PIL image
depth_image = Image.fromarray(depth_image)
return depth_image |