working proto

app.py

@@ -1,8 +1,7 @@
-from pytorch3d.io import load_obj
-from pytorch3d.structures import Meshes
 import torch
 import gradio as gr
 import plotly.graph_objects as go
+import trimesh
 
 device = torch.device("cpu")
 model = torch.jit.load('model_scripted.pt').to(device)
@@ -14,11 +13,10 @@ def normalize_vertices(verts):
     scale = max(verts.abs().max(0)[0])
     return verts / scale
 
-
 def plot_3d_results(verts, faces, uv_seam_edge_indices):
     # Convert vertices to NumPy for easier manipulation
     verts_np = verts.cpu().numpy()
-    faces_np = faces.verts_idx.cpu().numpy()
+    faces_np = faces.cpu().numpy()
 
     # Prepare the vertex coordinates for the Mesh3d plot
     x, y, z = verts_np[:, 0], verts_np[:, 1], verts_np[:, 2]
@@ -56,21 +54,56 @@ def plot_3d_results(verts, faces, uv_seam_edge_indices):
 
 
 def generate_prediction(file_input, treshold_value=0.5):
-    verts, faces, aux = load_obj(file_input)
-    verts = normalize_vertices(verts)
-    mesh = Meshes(verts=[verts.to(device)], faces=[faces.verts_idx.to(device)])
+    # Load the triangle mesh
+    mesh = trimesh.load_mesh(file_input)
+
+    # For production, we should use a faster method to preprocess the mesh!
+
+    # Convert vertices to a PyTorch tensor
+    vertices = torch.tensor(mesh.vertices, dtype=torch.float32)
+
+    # Initialize containers for unique vertices and mapping
+    unique_vertices = []
+    vertex_mapping = {}
+    new_faces = []
+
+    # Populate unique vertices and create new faces with updated indices
+    for face in mesh.faces:
+        new_face = []
+        for orig_index in face:
+            vertex = tuple(vertices[orig_index].tolist())  # Convert to tuple (hashable)
+            if vertex not in vertex_mapping:
+                vertex_mapping[vertex] = len(unique_vertices)
+                unique_vertices.append(vertices[orig_index])
+            new_face.append(vertex_mapping[vertex])
+        new_faces.append(new_face)
+
+    # Create edge set to ensure uniqueness
+    edge_set = set()
+    for face in new_faces:
+        # Unpack the vertex indices
+        v1, v2, v3 = face
+        # Create undirected edges (use tuple sorting to ensure uniqueness)
+        edge_set.add(tuple(sorted((v1, v2))))
+        edge_set.add(tuple(sorted((v2, v3))))
+        edge_set.add(tuple(sorted((v1, v3))))
+
+    # Convert edges back to tensor
+    edges = torch.tensor(list(edge_set), dtype=torch.long)
+
+    # Convert unique vertices and new faces back to tensors
+    verts = torch.stack(unique_vertices)
+    faces = torch.tensor(new_faces, dtype=torch.long)
 
     model.eval()
 
-    test_verts = mesh.verts_packed().to(device)
-    test_edges = mesh.edges_packed().to(device)
-
     with torch.no_grad():
-        test_outputs_logits = model(test_verts, test_edges).to(device)
+        test_outputs_logits = model(verts, edges).to(device)
         test_outputs = torch.sigmoid(test_outputs_logits).to(device)
         test_predictions = (test_outputs > treshold_value).int().cpu()
-        uv_seam_edges_mask = test_predictions.cpu().squeeze() == 1
-        uv_seam_edges = test_edges[uv_seam_edges_mask].cpu().tolist()
+
+    uv_seam_edges_mask = test_predictions.cpu().squeeze() == 1
+    uv_seam_edges = edges[uv_seam_edges_mask].cpu().tolist()
 
     # Return the HTML content generated by plot_3d_results
     return plot_3d_results(verts, faces, uv_seam_edges)

model_scripted.pt

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5475053a74a596b223abc768bf8d926247bf94f826c1ed41b469f593a11828c8
+oid sha256:6f0d5dcd806540ba8061e3d70ed17eac539e560ae0932f04679f706898126588
 size 255324

requirements.txt

@@ -1,4 +1,3 @@
-git+https://github.com/facebookresearch/pytorch3d.git
 torch
-pytorch3d
-plotly
+plotly
+trimesh