Update app.py

app.py

@@ -11,6 +11,7 @@ from torch import nn
 from transformers import (
     AutoTokenizer, Qwen2ForCausalLM, Qwen2Model, PreTrainedModel)
 from transformers.modeling_outputs import CausalLMOutputWithPast
+from pytorch3d.ops import sample_farthest_points
 
 
 class FourierPointEncoder(nn.Module):
@@ -18,13 +19,13 @@ class FourierPointEncoder(nn.Module):
         super().__init__()
         frequencies = 2.0 ** torch.arange(8, dtype=torch.float32)
         self.register_buffer('frequencies', frequencies, persistent=False)
-        self.projection = nn.Linear(54, hidden_size)
+        self.projection = nn.Linear(51, hidden_size)
 
     def forward(self, points):
-        x = points[..., :3]
+        x = points
         x = (x.unsqueeze(-1) * self.frequencies).view(*x.shape[:-1], -1)
-        x = torch.cat((points[..., :3], x.sin(), x.cos()), dim=-1)
-        x = self.projection(torch.cat((x, points[..., 3:]), dim=-1))
+        x = torch.cat((points, x.sin(), x.cos()), dim=-1)
+        x = self.projection(x)
         return x
 
 
@@ -113,15 +114,11 @@ class CADRecode(Qwen2ForCausalLM):
         return model_inputs
 
 
-def mesh_to_point_cloud(mesh, n_points=256):
-    vertices, faces = trimesh.sample.sample_surface(mesh, n_points)
-    point_cloud = np.concatenate((
-        np.asarray(vertices),
-        mesh.face_normals[faces]
-    ), axis=1)
-    ids = np.lexsort((point_cloud[:, 0], point_cloud[:, 1], point_cloud[:, 2]))
-    point_cloud = point_cloud[ids]
-    return point_cloud
+def mesh_to_point_cloud(mesh, n_points=256, n_pre_points=8192):
+    vertices, _ = trimesh.sample.sample_surface(mesh, n_pre_points)
+    _, ids = sample_farthest_points(torch.tensor(vertices).unsqueeze(0), K=n_points)
+    ids = ids[0].numpy()
+    return np.asarray(vertices[ids])
 
 
 def py_string_to_mesh_file(py_string, mesh_path, queue):
@@ -219,9 +216,7 @@ def run():
         with gr.Row(equal_height=True):
             in_model = gr.Model3D(label='1. Input Mesh', interactive=True)
             point_model = gr.Model3D(label='2. Sampled Point Cloud', display_mode='point_cloud', interactive=False)
-            out_model = gr.Model3D(
-                label='4. Result CAD Model', interactive=False
-            )
+            out_model = gr.Model3D(label='4. Result CAD Model', interactive=False)
         
         with gr.Row():
             with gr.Column():
@@ -230,27 +225,30 @@ def run():
                 with gr.Row():
                     gr.Examples(
                         examples=[
-                            ['./data/49215_5368e45e_0000.stl', 42],
-                            ['./data/00882236.stl', 6],
-                            ['./data/User Library-engrenage.stl', 18],
+                            ['./data/49215_5368e45e_0000.stl', 46],
+                            ['./data/User Library-engrenage.stl', 46],
                             ['./data/00010900.stl', 42],
-                            ['./data/21492_8bd34fc1_0008.stl', 42],
-                            ['./data/00375556.stl', 96],
-                            ['./data/49121_adb01620_0000.stl', 42],
-                            ['./data/41473_c2137170_0023.stl', 42],
-                            ['./data/22447_4062c6cb_0011.stl', 42],
+                            ['./data/00375556.stl', 44],
+                            ['./data/41473_c2137170_0023.stl', 47],
+                            ['./data/22447_4062c6cb_0011.stl', 43],
                             ['./data/27694_7801dc67_0017.stl', 45],
-                            ['./data/49562_6df35938_0005.stl', 44],
-                            ['./data/131709_8b86dfb6_0000.stl', 45],
+                            ['./data/49562_6df35938_0005.stl', 48],
                             ['./data/00081523.stl', 44],
-                            ['./data/00614972.stl', 44],
-                            ['./data/User Library-Cople Ventosa V2.stl', 50],
-                            ['./data/User Library-Caster, Superior Brand, 600lb 2 EDIT.stl', 52]],
+                            ['./data/User Library-Caster, Superior Brand, 600lb 2 EDIT.stl', 57],
+                            ['./data/00810270.stl', 44],
+                            ['./data/00295262.stl', 42],
+                            ['./data/132156_cd3f1428_0000.stl', 59],
+                            ['./data/User Library-All Weather Padlock_Padlock Key.stl', 66],
+                            ['./data/00047559.stl', 42],
+                            ['./data/00052220.stl', 42],
+                            ['./data/141665_0564e852_0003.stl', 42],
+                            ['./data/User Library-_50 Cal Round.stl', 42]],
                         example_labels=[
-                            'fusion360_table1', 'deepcad_star', 'cc3d_gear', 'deepcad_barrels',
-                            'fusion360_gear', 'deepcad_house', 'fusion360_table2', 'fusion360_omega',
-                            'fusion360_hat', 'fusion360_table3', 'fusion360_bolt', 'fusion360_clamp',
-                            'deepcad_alpha', 'deepcad_pulley', 'cc3d_flange', 'cc3d_caster'],
+                            'fusion360_table1', 'cc3d_gear', 'deepcad_barrels', 'deepcad_house',
+                            'fusion360_omega', 'fusion360_hat', 'fusion360_table3', 'fusion360_bolt',
+                            'deepcad_alpha', 'cc3d_caster', 'deepcad_spinner', 'deepcad_5bricks',
+                            'fusion360_star', 'cc3d_key', 'deepcad_3cubes', 'deepcad_arrow',
+                            'fusion360_frame', 'cc3d_bullet'],
                         inputs=[in_model, seed_slider],
                         cache_examples=False,
                         examples_per_page=20)
@@ -286,7 +284,7 @@ tokenizer = AutoTokenizer.from_pretrained(
     pad_token='<|im_end|>',
     padding_side='left')
 cad_recode = CADRecode.from_pretrained(
-    'filapro/cad-recode',
+    'filapro/cad-recode-v1.5',
     torch_dtype='auto',
     attn_implementation='flash_attention_2').eval()