gradio updated

app.py

@@ -0,0 +1,258 @@
+import jax
+import numpy as np
+import jax.numpy as jnp
+import matplotlib.pyplot as plt
+from functools import partial
+import torch
+from torch_geometric.nn import knn
+import gradio as gr
+import os
+import glob
+import pickle
+import time
+import cv2
+import numpy as np
+import matplotlib.pyplot as plt
+import plotly.graph_objects as go
+import open3d as o3d
+import trimesh
+import gradio
+
+
+
+def createPlane(normal, point_on_plane):
+    normal = normal / np.linalg.norm(normal)
+    
+    # Find a vector in the plane
+    if np.allclose(normal, [1, 0, 0]):
+        v1 = np.cross(normal, [0, 1, 0])
+    else:
+        v1 = np.cross(normal, [1, 0, 0])
+    
+    v1 = v1 / np.linalg.norm(v1)
+    v2 = np.cross(normal, v1)
+    v2 = v2 / np.linalg.norm(v2)
+    
+    half_width = 1
+    half_height = 1
+    
+    # Calculate the corners
+    corner1 = point_on_plane + half_width * v1 + half_height * v2
+    corner2 = point_on_plane - half_width * v1 + half_height * v2
+    corner3 = point_on_plane - half_width * v1 - half_height * v2
+    corner4 = point_on_plane + half_width * v1 - half_height * v2
+
+    vertices = np.array([corner1, corner2, corner3, corner4])
+
+    faces = np.array([
+        [0, 1, 2],
+        [0, 2, 3],
+        [2, 1, 0],
+        [3, 2, 0]
+    ])
+    # Define the color (sky blue) with opacity (alpha)
+    # Define the color (sky blue) with transparency
+    sky_blue_with_alpha = [255, 255, 255, 128]  # RGBA format, with 128 alpha for half opacity
+
+    # Set the vertex colors with transparency
+    vertex_colors = np.tile(sky_blue_with_alpha, (vertices.shape[0], 1))
+
+    # Create a mesh for the rectangle
+    plane_mesh = trimesh.Trimesh(vertices=vertices, faces=faces, vertex_colors=vertex_colors)
+    return plane_mesh
+
+def reflect_points_multiple_3d(d, n, p):
+    points_expanded = p.unsqueeze(0).expand(n.size(0), -1, -1)
+    normals_expanded = n.unsqueeze(1).expand(-1, p.size(0), -1)
+    distances_expanded = d.unsqueeze(1).expand(-1, p.size(0))
+    dot_products = torch.sum(points_expanded * normals_expanded, dim=2)
+    # reflections: (m, n, 3)
+    reflections = points_expanded - 2 * (dot_products - distances_expanded).unsqueeze(2) * normals_expanded
+
+    return reflections
+
+def reflection_point_association_3d(d, n, q, threshold):
+    # d in shape (m, ), n in shape (m, 3), q in shape (n, 3)
+    reflections = reflect_points_multiple_3d(d, n, q)  # shape: (m, n, 3)
+    reflections = reflections.view(-1, 3)  # Flatten to (m*n, 3)
+
+    # Using knn to find the closest points in q for each point in reflections
+    # knn finds indices of the nearest neighbors
+    _, indices = knn(q, reflections, k=1, batch_x=None, batch_y=None) # indices shape: (m*n, 1)
+
+    # Gather nearest points based on indices from q
+    nearest_points = q[indices.squeeze()] # shape: (m*n, 3)
+
+    # Calculate distances for the nearest neighbors
+    distances = (nearest_points - reflections).norm(dim=1) # shape: (m*n,)
+
+    # Reshape distances back to (m, n) and check threshold
+    distances = distances.view(d.size(0), -1) # shape: (m, n)
+    within_threshold = distances <= threshold # shape: (m, n), sum this along axis 1 to get the number of associated points
+    return within_threshold
+
+def get_patches(points, centroids):
+    norm = np.linalg.norm(centroids, axis=1)
+    n = centroids / norm[:, None]
+    d = norm - 1
+
+    association = reflection_point_association_3d(torch.tensor(np.array(d)), torch.tensor(np.array(n)),
+                                                  torch.tensor(np.array(points)), 0.03)
+
+    return np.array(association)
+
+def left_right(allpoints, patchbool, planepoints):
+    """
+    inputs:     patchpoints:    (n,3) 
+                planepoints:    (4,3) 
+    outputs:    leftpoints, rightpoints: (k,3), (k',3)
+    """
+    def signed_distance(point, plane_point, normal):
+        return np.dot(normal, point - plane_point) / (np.linalg.norm(normal)+1e-6)
+    patchpoints = allpoints[patchbool]
+    p1 = planepoints[0]
+    p2 = planepoints[1]
+    p3 = planepoints[2]
+    v1 = p2 - p1
+    v2 = p3 - p1
+    normal = np.cross(v1, v2)
+
+    distances = np.array([signed_distance(point, p1, normal) for point in patchpoints])
+    contains_nan = np.isnan(distances).any()
+    l_idx = distances<0
+    allidcs = np.arange(len(allpoints))[patchbool]
+    left_idx = allidcs[l_idx]
+    right_idx = allidcs[~l_idx]
+      
+    #left_points = patchpoints[l_idx]
+    #right_points = patchpoints[~l_idx]
+
+    
+    return left_idx, right_idx#left_points, right_points
+
+def dbscan(D, eps, MinPts):
+    labels = [0]*len(D)
+    C = 0
+    for P in range(0, len(D)):
+        if not (labels[P] == 0):
+           continue
+        NeighborPts = region_query(D, P, eps)
+        if len(NeighborPts) < MinPts:
+            labels[P] = -1
+        else: 
+           C += 1
+           grow_cluster(D, labels, P, NeighborPts, C, eps, MinPts)
+    return labels
+
+def grow_cluster(D, labels, P, NeighborPts, C, eps, MinPts):
+    labels[P] = C
+    i = 0
+    while i < len(NeighborPts):    
+        Pn = NeighborPts[i]
+        if labels[Pn] == -1:
+           labels[Pn] = C
+        elif labels[Pn] == 0:
+            labels[Pn] = C
+            PnNeighborPts = region_query(D, Pn, eps)
+            if len(PnNeighborPts) >= MinPts:
+                NeighborPts = NeighborPts + PnNeighborPts
+
+        i += 1        
+    
+def region_query(D, P, eps):
+    neighbors = []
+    for Pn in range(0, len(D)):
+        #if geodesic_dist(D[P], D[Pn])<eps:
+        if np.linalg.norm(D[P] - D[Pn]) < eps:
+           neighbors.append(Pn)
+            
+    return neighbors
+
+def compute_centroids(data, labels):
+    unique = np.unique(labels)
+    unique_labels = unique[unique!=-1]
+    centroids = []
+    for label in unique_labels:
+        mask = labels == label
+        points = data[mask]
+        centroid = jnp.mean(points, axis=0)
+        centroids.append(centroid)
+    return np.stack(centroids)
+
+
+def proc_all(mesh_path, mode_path):
+  with open(mode_path, 'rb') as f:
+    fin = pickle.load(f)
+  name = mesh_path.split('/')[-1].split('.')[0]
+  mesh = trimesh.load(mesh_path)
+  verts = np.array(mesh.vertices)
+
+  my_labels = dbscan(fin, eps=0.1, MinPts=1) 
+  centroid = np.array(compute_centroids(fin, my_labels))
+  print("total number of modes found: " + str(len(centroid)))
+
+  norm = torch.norm(torch.tensor(np.array(centroid)), dim=-1)
+  n = centroid / norm[:,None]
+  d = norm - 1
+  point = n * d[...,None]
+
+  #pts = create_points_3d(fin, alpha = 1, markersize = 4, label = 'final timestep')
+  #pts2 = create_points_3d(centroid, alpha = 1, markersize = 10, label = 'final timestep')
+  #plot_all_3d([pts, pts2])
+
+  pats = get_patches(torch.tensor(np.array(verts)), torch.tensor(centroid))
+  alldicts = []
+  pntdict = {}
+  for i in range(len(n)):
+      plane = createPlane(n[i], point[i])
+      l,r = left_right(verts, pats[i], plane.vertices)
+      single = {'plane': plane, 'left': l, 'right': r}
+      alldicts.append(single)
+  return alldicts
+
+def create_scene(mesh_path, dicts, plane_idx):
+    mesh = trimesh.load(mesh_path)
+    verts = np.array(mesh.vertices)
+    colors = []
+    for i in range(len(verts)):
+        if i in dicts[plane_idx]['left']:
+            color = [8, 136, 255]
+        elif i in dicts[plane_idx]['right']:
+            color = [204, 20, 245]
+        else:
+            color = [225, 225, 225]
+        colors.append(color)
+
+    mesh.visual.vertex_colors = colors
+    newplane = dicts[plane_idx]['plane']
+    newplane.visual.vertex_colors = np.tile([100, 100, 100, 100], (4, 1))
+    scene = trimesh.Scene([mesh, newplane])
+    temp_file = f"/tmp/scene_{plane_idx}.obj"
+    scene.export(temp_file)
+    return temp_file
+
+def load_mesh(mesh_file_name, mode_file_name):
+    
+    dicts = proc_all(mesh_file_name, mode_file_name)  # Assuming proc_all is defined and returns a list of dictionaries
+    allmesh = []
+    for i in range(min(5, len(dicts))):  # Limiting to maximum 5 outputs
+        temp_file = create_scene(mesh_file_name, dicts, i)
+        allmesh.append(temp_file)
+    return allmesh + [None] * (5 - len(allmesh))  # Fill the rest with None if less than 5
+
+examples = [["mesh_rescaled/" + mesh, "out_3d/" + mesh.split('.')[0] + "_50000_0.1_modes.pickle"] for mesh in os.listdir("mesh_rescaled/")]
+
+outputs = [gr.Model3D(label=f"3D Model {i+1}") for i in range(10)]
+
+demo = gr.Interface(
+    fn=load_mesh,
+    inputs=[gr.File(label="Upload Mesh"), gr.File(label="Upload Mode File")],
+    outputs=outputs,
+    examples=examples,
+    cache_examples=False  # Set to False to ensure it doesn't cache examples, requires actual file uploads
+)
+
+if __name__ == "__main__":
+    demo.launch()
+    
+    

requirements.txt

@@ -0,0 +1,4 @@
+trimesh
+open3d
+torch_geometric
+pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.3.0+cu121.html