Upload utils.py

demo_support/utils.py

@@ -0,0 +1,168 @@
+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:
+            # which geometry does this node refer to
+            transform, geometry_name = scene_or_mesh.graph[node_name]
+
+            # get the actual potential mesh instance
+            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):
+        # load the model
+        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)
+        # st.caption("Point Cloud Preview")
+    return col2