| | import numpy |
| | import trimesh |
| | import trimesh.sample |
| | import trimesh.visual |
| | import trimesh.proximity |
| | import objaverse |
| | import streamlit as st |
| | import plotly.graph_objects as go |
| | import matplotlib.pyplot as plotlib |
| |
|
| |
|
| | def get_bytes(x: str): |
| | import io, requests |
| | return io.BytesIO(requests.get(x).content) |
| |
|
| |
|
| | def get_image(x: str): |
| | try: |
| | return plotlib.imread(get_bytes(x), 'auto') |
| | except Exception: |
| | raise ValueError("Invalid image", x) |
| |
|
| |
|
| | def model_to_pc(mesh: trimesh.Trimesh, n_sample_points=10000): |
| | f32 = numpy.float32 |
| | rad = numpy.sqrt(mesh.area / (3 * n_sample_points)) |
| | for _ in range(24): |
| | pcd, face_idx = trimesh.sample.sample_surface_even(mesh, n_sample_points, rad) |
| | rad *= 0.85 |
| | if len(pcd) == n_sample_points: |
| | break |
| | else: |
| | raise ValueError("Bad geometry, cannot finish sampling.", mesh.area) |
| | if isinstance(mesh.visual, trimesh.visual.ColorVisuals): |
| | rgba = mesh.visual.face_colors[face_idx] |
| | elif isinstance(mesh.visual, trimesh.visual.TextureVisuals): |
| | bc = trimesh.proximity.points_to_barycentric(mesh.triangles[face_idx], pcd) |
| | if mesh.visual.uv is None or len(mesh.visual.uv) < mesh.faces[face_idx].max(): |
| | uv = numpy.zeros([len(bc), 2]) |
| | st.warning("Invalid UV, filling with zeroes") |
| | else: |
| | uv = numpy.einsum('ntc,nt->nc', mesh.visual.uv[mesh.faces[face_idx]], bc) |
| | material = mesh.visual.material |
| | if hasattr(material, 'materials'): |
| | if len(material.materials) == 0: |
| | rgba = numpy.ones_like(pcd) * 0.8 |
| | texture = None |
| | st.warning("Empty MultiMaterial found, falling back to light grey") |
| | else: |
| | material = material.materials[0] |
| | if hasattr(material, 'image'): |
| | texture = material.image |
| | if texture is None: |
| | rgba = numpy.zeros([len(uv), len(material.main_color)]) + material.main_color |
| | elif hasattr(material, 'baseColorTexture'): |
| | texture = material.baseColorTexture |
| | if texture is None: |
| | rgba = numpy.zeros([len(uv), len(material.main_color)]) + material.main_color |
| | else: |
| | texture = None |
| | rgba = numpy.ones_like(pcd) * 0.8 |
| | st.warning("Unknown material, falling back to light grey") |
| | if texture is not None: |
| | rgba = trimesh.visual.uv_to_interpolated_color(uv, texture) |
| | if rgba.max() > 1: |
| | if rgba.max() > 255: |
| | rgba = rgba.astype(f32) / rgba.max() |
| | else: |
| | rgba = rgba.astype(f32) / 255.0 |
| | return numpy.concatenate([numpy.array(pcd, f32), numpy.array(rgba, f32)[:, :3]], axis=-1) |
| |
|
| |
|
| | def trimesh_to_pc(scene_or_mesh): |
| | if isinstance(scene_or_mesh, trimesh.Scene): |
| | meshes = [] |
| | for node_name in scene_or_mesh.graph.nodes_geometry: |
| | |
| | transform, geometry_name = scene_or_mesh.graph[node_name] |
| |
|
| | |
| | geometry = scene_or_mesh.geometry[geometry_name].copy() |
| | if not hasattr(geometry, 'triangles'): |
| | continue |
| | geometry: trimesh.Trimesh |
| | geometry = geometry.apply_transform(transform) |
| | meshes.append(geometry) |
| | total_area = sum(geometry.area for geometry in meshes) |
| | if total_area < 1e-6: |
| | raise ValueError("Bad geometry: total area too small (< 1e-6)") |
| | pcs = [] |
| | for geometry in meshes: |
| | pcs.append(model_to_pc(geometry, max(1, round(geometry.area / total_area * 10000)))) |
| | if not len(pcs): |
| | raise ValueError("Unsupported mesh object: no triangles found") |
| | return numpy.concatenate(pcs) |
| | else: |
| | assert isinstance(scene_or_mesh, trimesh.Trimesh) |
| | return model_to_pc(scene_or_mesh, 10000) |
| |
|
| |
|
| | def input_3d_shape(key=None): |
| | if key is None: |
| | objaid_key = model_key = npy_key = swap_key = None |
| | else: |
| | objaid_key = key + "_objaid" |
| | model_key = key + "_model" |
| | npy_key = key + "_npy" |
| | swap_key = key + "_swap" |
| | objaid = st.text_input("Enter an Objaverse ID", key=objaid_key) |
| | model = st.file_uploader("Or upload a model (.glb/.obj/.ply)", key=model_key) |
| | npy = st.file_uploader("Or upload a point cloud numpy array (.npy of Nx3 XYZ or Nx6 XYZRGB)", key=npy_key) |
| | swap_yz_axes = st.radio("Gravity", ["Y is up (for most Objaverse shapes)", "Z is up"], key=swap_key) == "Z is up" |
| | f32 = numpy.float32 |
| |
|
| | def load_data(prog): |
| | |
| | prog.progress(0.05, "Preparing Point Cloud") |
| | if npy is not None: |
| | pc: numpy.ndarray = numpy.load(npy) |
| | elif model is not None: |
| | pc = trimesh_to_pc(trimesh.load(model, model.name.split(".")[-1])) |
| | elif objaid: |
| | prog.progress(0.1, "Downloading Objaverse Object") |
| | objamodel = objaverse.load_objects([objaid])[objaid] |
| | prog.progress(0.2, "Preparing Point Cloud") |
| | pc = trimesh_to_pc(trimesh.load(objamodel)) |
| | else: |
| | raise ValueError("You have to supply 3D input!") |
| | prog.progress(0.25, "Preprocessing Point Cloud") |
| | assert pc.ndim == 2, "invalid pc shape: ndim = %d != 2" % pc.ndim |
| | assert pc.shape[1] in [3, 6], "invalid pc shape: should have 3/6 channels, got %d" % pc.shape[1] |
| | pc = pc.astype(f32) |
| | if swap_yz_axes: |
| | pc[:, [1, 2]] = pc[:, [2, 1]] |
| | pc[:, :3] = pc[:, :3] - numpy.mean(pc[:, :3], axis=0) |
| | pc[:, :3] = pc[:, :3] / numpy.linalg.norm(pc[:, :3], axis=-1).max() |
| | if pc.shape[1] == 3: |
| | pc = numpy.concatenate([pc, numpy.ones_like(pc) * 0.4], axis=-1) |
| | prog.progress(0.27, "Normalized Point Cloud") |
| | if pc.shape[0] >= 10000: |
| | pc = pc[numpy.random.permutation(len(pc))[:10000]] |
| | elif pc.shape[0] == 0: |
| | raise ValueError("Got empty point cloud!") |
| | elif pc.shape[0] < 10000: |
| | pc = numpy.concatenate([pc, pc[numpy.random.randint(len(pc), size=[10000 - len(pc)])]]) |
| | prog.progress(0.3, "Preprocessed Point Cloud") |
| | return pc.astype(f32) |
| |
|
| | return load_data |
| |
|
| |
|
| | def render_pc(pc): |
| | rand = numpy.random.permutation(len(pc))[:2048] |
| | pc = pc[rand] |
| | rgb = (pc[:, 3:] * 255).astype(numpy.uint8) |
| | g = go.Scatter3d( |
| | x=pc[:, 0], y=pc[:, 1], z=pc[:, 2], |
| | mode='markers', |
| | marker=dict(size=2, color=[f'rgb({rgb[i, 0]}, {rgb[i, 1]}, {rgb[i, 2]})' for i in range(len(pc))]), |
| | ) |
| | fig = go.Figure(data=[g]) |
| | fig.update_layout(scene_camera=dict(up=dict(x=0, y=1, z=0))) |
| | fig.update_scenes(aspectmode="data") |
| | col1, col2 = st.columns(2) |
| | with col1: |
| | st.plotly_chart(fig, use_container_width=True) |
| | |
| | return col2 |
| |
|
