wehub-resource-sync
565 lines
21 KiB
Python
565 lines
21 KiB
Python
from typing import *
|
|
import hashlib
|
|
import numpy as np
|
|
import cv2
|
|
|
|
|
|
def get_file_hash(file: str) -> str:
|
|
sha256 = hashlib.sha256()
|
|
# Read the file from the path
|
|
with open(file, "rb") as f:
|
|
# Update the hash with the file content
|
|
for byte_block in iter(lambda: f.read(4096), b""):
|
|
sha256.update(byte_block)
|
|
return sha256.hexdigest()
|
|
|
|
# ===============LOW DISCREPANCY SEQUENCES================
|
|
|
|
PRIMES = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53]
|
|
|
|
def radical_inverse(base, n):
|
|
val = 0
|
|
inv_base = 1.0 / base
|
|
inv_base_n = inv_base
|
|
while n > 0:
|
|
digit = n % base
|
|
val += digit * inv_base_n
|
|
n //= base
|
|
inv_base_n *= inv_base
|
|
return val
|
|
|
|
def halton_sequence(dim, n):
|
|
return [radical_inverse(PRIMES[dim], n) for dim in range(dim)]
|
|
|
|
def hammersley_sequence(dim, n, num_samples):
|
|
return [n / num_samples] + halton_sequence(dim - 1, n)
|
|
|
|
def sphere_hammersley_sequence(n, num_samples, offset=(0, 0)):
|
|
u, v = hammersley_sequence(2, n, num_samples)
|
|
u += offset[0] / num_samples
|
|
v += offset[1]
|
|
u = 2 * u if u < 0.25 else 2 / 3 * u + 1 / 3
|
|
theta = np.arccos(1 - 2 * u) - np.pi / 2
|
|
phi = v * 2 * np.pi
|
|
return [phi, theta]
|
|
|
|
# ==============PLY IO===============
|
|
import struct
|
|
import re
|
|
import torch
|
|
|
|
def read_ply(filename):
|
|
"""
|
|
Read a PLY file and return vertices, triangle faces, and quad faces.
|
|
|
|
Args:
|
|
filename (str): The file path to read from.
|
|
|
|
Returns:
|
|
vertices (torch.Tensor): Tensor of shape [N, 3] containing vertex positions.
|
|
tris (torch.Tensor): Tensor of shape [M, 3] containing triangle face indices (empty if none).
|
|
quads (torch.Tensor): Tensor of shape [K, 4] containing quad face indices (empty if none).
|
|
"""
|
|
with open(filename, 'rb') as f:
|
|
# Read the header until 'end_header' is encountered
|
|
header_bytes = b""
|
|
while True:
|
|
line = f.readline()
|
|
if not line:
|
|
raise ValueError("PLY header not found")
|
|
header_bytes += line
|
|
if b"end_header" in line:
|
|
break
|
|
header = header_bytes.decode('utf-8')
|
|
|
|
# Determine if the file is in ASCII or binary format
|
|
is_ascii = "ascii" in header
|
|
|
|
# Extract the number of vertices and faces from the header using regex
|
|
vertex_match = re.search(r'element vertex (\d+)', header)
|
|
if vertex_match:
|
|
num_vertices = int(vertex_match.group(1))
|
|
else:
|
|
raise ValueError("Vertex count not found in header")
|
|
|
|
face_match = re.search(r'element face (\d+)', header)
|
|
if face_match:
|
|
num_faces = int(face_match.group(1))
|
|
else:
|
|
raise ValueError("Face count not found in header")
|
|
|
|
vertices = []
|
|
tris = []
|
|
quads = []
|
|
|
|
if is_ascii:
|
|
# For ASCII format, read each line of vertex data (each line contains 3 floats)
|
|
for _ in range(num_vertices):
|
|
line = f.readline().decode('utf-8').strip()
|
|
if not line:
|
|
continue
|
|
parts = line.split()
|
|
vertices.append([float(parts[0]), float(parts[1]), float(parts[2])])
|
|
|
|
# Read face data, where the first number indicates the number of vertices for the face
|
|
for _ in range(num_faces):
|
|
line = f.readline().decode('utf-8').strip()
|
|
if not line:
|
|
continue
|
|
parts = line.split()
|
|
count = int(parts[0])
|
|
indices = list(map(int, parts[1:]))
|
|
if count == 3:
|
|
tris.append(indices)
|
|
elif count == 4:
|
|
quads.append(indices)
|
|
else:
|
|
# Skip faces with other numbers of vertices (can be extended as needed)
|
|
pass
|
|
else:
|
|
# For binary format: read directly from the binary stream
|
|
# Each vertex consists of 3 floats (12 bytes per vertex)
|
|
for _ in range(num_vertices):
|
|
data = f.read(12)
|
|
if len(data) < 12:
|
|
raise ValueError("Insufficient vertex data")
|
|
v = struct.unpack('<fff', data)
|
|
vertices.append(v)
|
|
|
|
# Read face data from the binary stream
|
|
for _ in range(num_faces):
|
|
# First, read 1 byte indicating the number of vertices in the face
|
|
count_data = f.read(1)
|
|
if len(count_data) < 1:
|
|
raise ValueError("Failed to read face vertex count")
|
|
count = struct.unpack('<B', count_data)[0]
|
|
if count == 3:
|
|
data = f.read(12) # 3 * 4 bytes
|
|
if len(data) < 12:
|
|
raise ValueError("Insufficient data for triangle face")
|
|
indices = struct.unpack('<3i', data)
|
|
tris.append(indices)
|
|
elif count == 4:
|
|
data = f.read(16) # 4 * 4 bytes
|
|
if len(data) < 16:
|
|
raise ValueError("Insufficient data for quad face")
|
|
indices = struct.unpack('<4i', data)
|
|
quads.append(indices)
|
|
else:
|
|
# For faces with a different number of vertices, read count*4 bytes
|
|
data = f.read(count * 4)
|
|
# Skip or extend processing as needed
|
|
raise ValueError(f"Unsupported face with {count} vertices")
|
|
|
|
# Convert lists to torch.Tensor
|
|
vertices = torch.tensor(vertices, dtype=torch.float32)
|
|
tris = torch.tensor(tris, dtype=torch.int32) if len(tris) > 0 else torch.empty((0, 3), dtype=torch.int32)
|
|
quads = torch.tensor(quads, dtype=torch.int32) if len(quads) > 0 else torch.empty((0, 4), dtype=torch.int32)
|
|
|
|
return vertices, tris, quads
|
|
|
|
|
|
def write_ply(filename, vertices, tris, quads, ascii=False):
|
|
"""
|
|
Write a mesh to a PLY file, with the option to save in ASCII or binary format.
|
|
|
|
Args:
|
|
filename (str): The filename to write to.
|
|
vertices (torch.Tensor): [N, 3] The vertex positions.
|
|
tris (torch.Tensor): [M, 3] The triangle indices.
|
|
quads (torch.Tensor): [K, 4] The quad indices.
|
|
ascii (bool): If True, write in ASCII format. If False, write in binary format.
|
|
"""
|
|
# Convert torch tensors to numpy arrays
|
|
vertices = vertices.numpy()
|
|
tris = tris.numpy()
|
|
quads = quads.numpy()
|
|
|
|
# Prepare the header
|
|
num_vertices = len(vertices)
|
|
num_faces = len(tris) + len(quads)
|
|
|
|
# Vertex properties
|
|
vertex_header = "property float x\nproperty float y\nproperty float z"
|
|
|
|
# Face properties (the number of vertices per face is variable)
|
|
face_header = "property list uchar int vertex_index"
|
|
|
|
# Start writing the PLY header
|
|
header = f"ply\n"
|
|
header += f"format {'ascii 1.0' if ascii else 'binary_little_endian 1.0'}\n"
|
|
header += f"element vertex {num_vertices}\n"
|
|
header += vertex_header + "\n"
|
|
header += f"element face {num_faces}\n"
|
|
header += face_header + "\n"
|
|
header += "end_header\n"
|
|
|
|
# Open the file for writing
|
|
with open(filename, 'wb' if not ascii else 'w') as f:
|
|
# Write the header
|
|
f.write(header if ascii else header.encode('utf-8'))
|
|
|
|
# Write the vertex data
|
|
if ascii:
|
|
for v in vertices:
|
|
f.write(f"{v[0]} {v[1]} {v[2]}\n")
|
|
else:
|
|
for v in vertices:
|
|
f.write(struct.pack('<fff', *v))
|
|
|
|
# Write the face data
|
|
if ascii:
|
|
for tri in tris:
|
|
f.write(f"3 {tri[0]} {tri[1]} {tri[2]}\n")
|
|
for quad in quads:
|
|
f.write(f"4 {quad[0]} {quad[1]} {quad[2]} {quad[3]}\n")
|
|
else:
|
|
for tri in tris:
|
|
f.write(struct.pack('<B3i', 3, *tri)) # 3 indices for triangle
|
|
for quad in quads:
|
|
f.write(struct.pack('<B4i', 4, *quad)) # 4 indices for quad
|
|
|
|
|
|
# ==============IMAGE UTILS===============
|
|
|
|
def make_grid(images, nrow=None, ncol=None, aspect_ratio=None):
|
|
num_images = len(images)
|
|
if nrow is None and ncol is None:
|
|
if aspect_ratio is not None:
|
|
nrow = int(np.round(np.sqrt(num_images / aspect_ratio)))
|
|
else:
|
|
nrow = int(np.sqrt(num_images))
|
|
ncol = (num_images + nrow - 1) // nrow
|
|
elif nrow is None and ncol is not None:
|
|
nrow = (num_images + ncol - 1) // ncol
|
|
elif nrow is not None and ncol is None:
|
|
ncol = (num_images + nrow - 1) // nrow
|
|
else:
|
|
assert nrow * ncol >= num_images, 'nrow * ncol must be greater than or equal to the number of images'
|
|
|
|
if images[0].ndim == 2:
|
|
grid = np.zeros((nrow * images[0].shape[0], ncol * images[0].shape[1]), dtype=images[0].dtype)
|
|
else:
|
|
grid = np.zeros((nrow * images[0].shape[0], ncol * images[0].shape[1], images[0].shape[2]), dtype=images[0].dtype)
|
|
for i, img in enumerate(images):
|
|
row = i // ncol
|
|
col = i % ncol
|
|
grid[row * img.shape[0]:(row + 1) * img.shape[0], col * img.shape[1]:(col + 1) * img.shape[1]] = img
|
|
return grid
|
|
|
|
|
|
def notes_on_image(img, notes=None):
|
|
img = np.pad(img, ((0, 32), (0, 0), (0, 0)), 'constant', constant_values=0)
|
|
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
|
|
if notes is not None:
|
|
img = cv2.putText(img, notes, (0, img.shape[0] - 4), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 1)
|
|
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
|
|
return img
|
|
|
|
|
|
|
|
def text_image(text, resolution=(512, 512), max_size=0.5, h_align="left", v_align="center"):
|
|
"""
|
|
Draw text on an image of the given resolution. The text is automatically wrapped
|
|
and scaled so that it fits completely within the image while preserving any explicit
|
|
line breaks and original spacing. Horizontal and vertical alignment can be controlled
|
|
via flags.
|
|
|
|
Parameters:
|
|
text (str): The input text. Newline characters and spacing are preserved.
|
|
resolution (tuple): The image resolution as (width, height).
|
|
max_size (float): The maximum font size.
|
|
h_align (str): Horizontal alignment. Options: "left", "center", "right".
|
|
v_align (str): Vertical alignment. Options: "top", "center", "bottom".
|
|
|
|
Returns:
|
|
numpy.ndarray: The resulting image (BGR format) with the text drawn.
|
|
"""
|
|
width, height = resolution
|
|
# Create a white background image
|
|
img = np.full((height, width, 3), 255, dtype=np.uint8)
|
|
|
|
# Set margins and compute available drawing area
|
|
margin = 10
|
|
avail_width = width - 2 * margin
|
|
avail_height = height - 2 * margin
|
|
|
|
# Choose OpenCV font and text thickness
|
|
font = cv2.FONT_HERSHEY_SIMPLEX
|
|
thickness = 1
|
|
# Ratio for additional spacing between lines (relative to the height of "A")
|
|
line_spacing_ratio = 0.5
|
|
|
|
def wrap_line(line, max_width, font, thickness, scale):
|
|
"""
|
|
Wrap a single line of text into multiple lines such that each line's
|
|
width (measured at the given scale) does not exceed max_width.
|
|
This function preserves the original spacing by splitting the line into tokens
|
|
(words and whitespace) using a regular expression.
|
|
|
|
Parameters:
|
|
line (str): The input text line.
|
|
max_width (int): Maximum allowed width in pixels.
|
|
font (int): OpenCV font identifier.
|
|
thickness (int): Text thickness.
|
|
scale (float): The current font scale.
|
|
|
|
Returns:
|
|
List[str]: A list of wrapped lines.
|
|
"""
|
|
# Split the line into tokens (words and whitespace), preserving spacing
|
|
tokens = re.split(r'(\s+)', line)
|
|
if not tokens:
|
|
return ['']
|
|
|
|
wrapped_lines = []
|
|
current_line = ""
|
|
for token in tokens:
|
|
candidate = current_line + token
|
|
candidate_width = cv2.getTextSize(candidate, font, scale, thickness)[0][0]
|
|
if candidate_width <= max_width:
|
|
current_line = candidate
|
|
else:
|
|
# If current_line is empty, the token itself is too wide;
|
|
# break the token character by character.
|
|
if current_line == "":
|
|
sub_token = ""
|
|
for char in token:
|
|
candidate_char = sub_token + char
|
|
if cv2.getTextSize(candidate_char, font, scale, thickness)[0][0] <= max_width:
|
|
sub_token = candidate_char
|
|
else:
|
|
if sub_token:
|
|
wrapped_lines.append(sub_token)
|
|
sub_token = char
|
|
current_line = sub_token
|
|
else:
|
|
wrapped_lines.append(current_line)
|
|
current_line = token
|
|
if current_line:
|
|
wrapped_lines.append(current_line)
|
|
return wrapped_lines
|
|
|
|
def compute_text_block(scale):
|
|
"""
|
|
Wrap the entire text (splitting at explicit newline characters) using the
|
|
provided scale, and then compute the overall width and height of the text block.
|
|
|
|
Returns:
|
|
wrapped_lines (List[str]): The list of wrapped lines.
|
|
block_width (int): Maximum width among the wrapped lines.
|
|
block_height (int): Total height of the text block including spacing.
|
|
sizes (List[tuple]): A list of (width, height) for each wrapped line.
|
|
spacing (int): The spacing between lines (computed from the scaled "A" height).
|
|
"""
|
|
# Split text by explicit newlines
|
|
input_lines = text.splitlines() if text else ['']
|
|
wrapped_lines = []
|
|
for line in input_lines:
|
|
wrapped = wrap_line(line, avail_width, font, thickness, scale)
|
|
wrapped_lines.extend(wrapped)
|
|
|
|
sizes = []
|
|
for line in wrapped_lines:
|
|
(text_size, _) = cv2.getTextSize(line, font, scale, thickness)
|
|
sizes.append(text_size) # (width, height)
|
|
|
|
block_width = max((w for w, h in sizes), default=0)
|
|
# Use the height of "A" (at the current scale) to compute line spacing
|
|
base_height = cv2.getTextSize("A", font, scale, thickness)[0][1]
|
|
spacing = int(line_spacing_ratio * base_height)
|
|
block_height = sum(h for w, h in sizes) + spacing * (len(sizes) - 1) if sizes else 0
|
|
|
|
return wrapped_lines, block_width, block_height, sizes, spacing
|
|
|
|
# Use binary search to find the maximum scale that allows the text block to fit
|
|
lo = 0.001
|
|
hi = max_size
|
|
eps = 0.001 # convergence threshold
|
|
best_scale = lo
|
|
best_result = None
|
|
|
|
while hi - lo > eps:
|
|
mid = (lo + hi) / 2
|
|
wrapped_lines, block_width, block_height, sizes, spacing = compute_text_block(mid)
|
|
# Ensure that both width and height constraints are met
|
|
if block_width <= avail_width and block_height <= avail_height:
|
|
best_scale = mid
|
|
best_result = (wrapped_lines, block_width, block_height, sizes, spacing)
|
|
lo = mid # try a larger scale
|
|
else:
|
|
hi = mid # reduce the scale
|
|
|
|
if best_result is None:
|
|
best_scale = 0.5
|
|
best_result = compute_text_block(best_scale)
|
|
|
|
wrapped_lines, block_width, block_height, sizes, spacing = best_result
|
|
|
|
# Compute starting y-coordinate based on vertical alignment flag
|
|
if v_align == "top":
|
|
y_top = margin
|
|
elif v_align == "center":
|
|
y_top = margin + (avail_height - block_height) // 2
|
|
elif v_align == "bottom":
|
|
y_top = margin + (avail_height - block_height)
|
|
else:
|
|
y_top = margin + (avail_height - block_height) // 2 # default to center if invalid flag
|
|
|
|
# For cv2.putText, the y coordinate represents the text baseline;
|
|
# so for the first line add its height.
|
|
y = y_top + (sizes[0][1] if sizes else 0)
|
|
|
|
# Draw each line with horizontal alignment based on the flag
|
|
for i, line in enumerate(wrapped_lines):
|
|
line_width, line_height = sizes[i]
|
|
if h_align == "left":
|
|
x = margin
|
|
elif h_align == "center":
|
|
x = margin + (avail_width - line_width) // 2
|
|
elif h_align == "right":
|
|
x = margin + (avail_width - line_width)
|
|
else:
|
|
x = margin # default to left if invalid flag
|
|
|
|
cv2.putText(img, line, (x, y), font, best_scale, (0, 0, 0), thickness, cv2.LINE_AA)
|
|
y += line_height + spacing
|
|
|
|
return img
|
|
|
|
|
|
# ==================== View index parsing ====================
|
|
|
|
def parse_view_indices(view_indices_str):
|
|
"""Parse view_indices string into a sorted deduplicated list of integers."""
|
|
if view_indices_str is None:
|
|
return None
|
|
view_indices = []
|
|
for part in view_indices_str.split(','):
|
|
if '-' in part:
|
|
start, end = map(int, part.split('-'))
|
|
view_indices.extend(range(start, end + 1))
|
|
else:
|
|
view_indices.append(int(part))
|
|
view_indices = list(set(view_indices))
|
|
view_indices.sort()
|
|
return view_indices
|
|
|
|
|
|
# ==================== Multi-view transform functions ====================
|
|
|
|
import math
|
|
|
|
def get_new_camera_matrix(radius: float, yaw: float, pitch: float, dtype=torch.float32, device='cpu'):
|
|
"""
|
|
Compute camera-to-world 4x4 transform matrix in spherical coordinates,
|
|
looking at origin with up=(0,1,0).
|
|
Uses standard LookAt formula where camera local -Z points at target and local Y is up.
|
|
An additional +90 degree rotation around local Z axis is applied to match
|
|
Blender Track-To local axis convention.
|
|
yaw, pitch are in radians.
|
|
"""
|
|
x = radius * math.cos(yaw) * math.cos(pitch)
|
|
y = radius * math.sin(yaw) * math.cos(pitch)
|
|
z = radius * math.sin(pitch)
|
|
eye = torch.tensor([x, y, z], dtype=dtype, device=device)
|
|
target = torch.zeros(3, dtype=dtype, device=device)
|
|
up_global = torch.tensor([0.0, 1.0, 0.0], dtype=dtype, device=device)
|
|
|
|
f = (target - eye)
|
|
f = f / torch.norm(f)
|
|
r = torch.cross(f, up_global)
|
|
r = r / torch.norm(r)
|
|
u = torch.cross(r, f)
|
|
|
|
z_cam = -f
|
|
x_cam = r
|
|
y_cam = u
|
|
|
|
T = torch.eye(4, dtype=dtype, device=device)
|
|
T[:3, 0] = x_cam
|
|
T[:3, 1] = y_cam
|
|
T[:3, 2] = z_cam
|
|
T[:3, 3] = eye
|
|
|
|
# +90 degree rotation matrix around local Z axis
|
|
Rz90 = torch.tensor([
|
|
[0.0, -1.0, 0.0, 0.0],
|
|
[1.0, 0.0, 0.0, 0.0],
|
|
[0.0, 0.0, 1.0, 0.0],
|
|
[0.0, 0.0, 0.0, 1.0]
|
|
], dtype=dtype, device=device)
|
|
|
|
return T @ Rz90
|
|
|
|
|
|
def transform_mesh(mesh_v, frame):
|
|
"""
|
|
Apply multi-view transform to mesh vertices based on camera transform matrix.
|
|
"""
|
|
device = mesh_v.device
|
|
c2w_orig = torch.tensor(frame['transform_matrix'], dtype=torch.float32, device=device)
|
|
|
|
# Old and new camera matrices
|
|
radius = c2w_orig[:3, 3].norm().item()
|
|
c2w_new = get_new_camera_matrix(radius=radius, yaw=-90/180.0*math.pi, pitch=0.0,
|
|
dtype=torch.float32, device=device)
|
|
w2c_orig = torch.inverse(c2w_orig)
|
|
|
|
# Initial and final axis alignment matrices
|
|
R_init = torch.tensor([
|
|
[1.0, 0.0, 0.0, 0.0],
|
|
[0.0, 0.0, -1.0, 0.0],
|
|
[0.0, 1.0, 0.0, 0.0],
|
|
[0.0, 0.0, 0.0, 1.0]
|
|
], dtype=torch.float32, device=device)
|
|
|
|
R_back = torch.tensor([
|
|
[1.0, 0.0, 0.0, 0.0],
|
|
[0.0, 0.0, 1.0, 0.0],
|
|
[0.0, -1.0, 0.0, 0.0],
|
|
[0.0, 0.0, 0.0, 1.0]
|
|
], dtype=torch.float32, device=device)
|
|
|
|
R_ply = torch.tensor([
|
|
[1.0, 0.0, 0.0, 0.0],
|
|
[0.0, 0.0, 1.0, 0.0],
|
|
[0.0, -1.0, 0.0, 0.0],
|
|
[0.0, 0.0, 0.0, 1.0]
|
|
], dtype=torch.float32, device=device)
|
|
|
|
T_cam = c2w_new @ w2c_orig @ R_ply
|
|
T_final = R_back @ T_cam @ R_init
|
|
|
|
# Apply transform
|
|
mesh_v = mesh_v.reshape(-1, 3)
|
|
verts_h = torch.cat([mesh_v, torch.ones((mesh_v.shape[0], 1), dtype=torch.float32, device=device)], dim=1)
|
|
verts_trans = (T_final @ verts_h.T).T[:, :3]
|
|
|
|
return verts_trans
|
|
|
|
|
|
def sphere_normalize_torch(vertices):
|
|
"""
|
|
Sphere normalization: normalize vertices based on sphere radius.
|
|
"""
|
|
bmin = torch.min(vertices, dim=0)[0]
|
|
bmax = torch.max(vertices, dim=0)[0]
|
|
bcenter = (bmax + bmin) / 2
|
|
assert bcenter.abs().max() < 0.25, f"bcenter is not close to origin: {bcenter}"
|
|
radius = torch.norm(vertices - bcenter, dim=-1).max()
|
|
vertices_normalized = vertices / radius
|
|
return vertices_normalized, bcenter, radius
|
|
|
|
|
|
def save_image_with_notes(img, path, notes=None):
|
|
"""
|
|
Save an image with notes.
|
|
"""
|
|
if isinstance(img, torch.Tensor):
|
|
img = img.cpu().numpy().transpose(1, 2, 0)
|
|
if img.dtype == np.float32 or img.dtype == np.float64:
|
|
img = np.clip(img * 255, 0, 255).astype(np.uint8)
|
|
img = notes_on_image(img, notes)
|
|
cv2.imwrite(path, cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
|