Update app.py

app.py

@@ -40,48 +40,65 @@ def sample_info(sample_id_str, fieldn):
 
     nodes = plaid_sample.get_nodes()
     field = plaid_sample.get_field(fieldn)
-    if nodes.shape[1] == 2:
-        nodes__ = np.zeros((nodes.shape[0],nodes.shape[1]+1))
-        nodes__[:,:-1] = nodes
-        nodes = nodes__
-
+    # if nodes.shape[1] == 2:
+    #     nodes__ = np.zeros((nodes.shape[0],nodes.shape[1]+1))
+    #     nodes__[:,:-1] = nodes
+    #     nodes = nodes__
+
+    norm = (field - field.min()) / (field.max() - field.min())
+    colormap_func = mpl.pyplot.get_cmap('viridis')
+    rgb_colors = colormap_func(norm)[:, :3]
+    
+    nb_nodes = nodes.shape[0]
 
     quads = plaid_sample.get_elements()['QUAD_4']
+    nb_quads = quads.shape[0]
+
+    assert field.shape[0] == nb_nodes
+
+    with open("visu.obj", 'w') as f:
+        for i in range(nb_nodes):
+            f.write(f"v {nodes[i,0]} {nodes[i,1]} {nodes[i,2]} {rgb_colors[i,0]} {rgb_colors[i,1]} {rgb_colors[i,2]}\n")
+       
+        for i in range(nb_quads):
+            f.write(f"f {quads[i,0] + 1} {quads[i,1] + 1} {quads[i,2] + 1} {quads[i,3] + 1}\n")
+
+    # quads = plaid_sample.get_elements()['QUAD_4']
 
-    # generate colormap
-    if np.linalg.norm(field) > 0:
-        norm = mpl.colors.Normalize(vmin=np.min(field), vmax=np.max(field))
-        cmap = cm.nipy_spectral#cm.coolwarm
-        m = cm.ScalarMappable(norm=norm, cmap=cmap)
+    # # generate colormap
+    # if np.linalg.norm(field) > 0:
+    #     norm = mpl.colors.Normalize(vmin=np.min(field), vmax=np.max(field))
+    #     cmap = cm.nipy_spectral#cm.coolwarm
+    #     m = cm.ScalarMappable(norm=norm, cmap=cmap)
     
-        vertex_colors = m.to_rgba(field)[:,:3]
-    else:
-        vertex_colors = 1+np.zeros((field.shape[0], 3))
-        vertex_colors[:,0] = 0.2298057
-        vertex_colors[:,1] = 0.01555616
-        vertex_colors[:,2] = 0.15023281
-
-    # generate mesh
-    trimesh = Trimesh(vertices = nodes, faces = quads)
-    trimesh.visual.vertex_colors = vertex_colors
-    mesh = pyrender.Mesh.from_trimesh(trimesh, smooth=False)
-
-    # compose scene
-    scene = pyrender.Scene(ambient_light=[.1, .1, .3], bg_color=[0, 0, 0])
-    camera = pyrender.PerspectiveCamera( yfov=np.pi / 6.0)
-    light = pyrender.DirectionalLight(color=[1,1,1], intensity=1000.)
-
-    scene.add(mesh, pose=  np.eye(4))
-    scene.add(light, pose=  np.eye(4))
-
-    scene.add(camera, pose=[[ 1,  0,  0,  0.02],
-                            [ 0,  1,  0,  0.21],
-                            [ 0,  0,  1,  0.19],
-                            [ 0,  0,  0,  1]])
-
-    # render scene
-    r = pyrender.OffscreenRenderer(1024, 1024)
-    color, _ = r.render(scene)
+    #     vertex_colors = m.to_rgba(field)[:,:3]
+    # else:
+    #     vertex_colors = 1+np.zeros((field.shape[0], 3))
+    #     vertex_colors[:,0] = 0.2298057
+    #     vertex_colors[:,1] = 0.01555616
+    #     vertex_colors[:,2] = 0.15023281
+
+    # # generate mesh
+    # trimesh = Trimesh(vertices = nodes, faces = quads)
+    # trimesh.visual.vertex_colors = vertex_colors
+    # mesh = pyrender.Mesh.from_trimesh(trimesh, smooth=False)
+
+    # # compose scene
+    # scene = pyrender.Scene(ambient_light=[.1, .1, .3], bg_color=[0, 0, 0])
+    # camera = pyrender.PerspectiveCamera( yfov=np.pi / 6.0)
+    # light = pyrender.DirectionalLight(color=[1,1,1], intensity=1000.)
+
+    # scene.add(mesh, pose=  np.eye(4))
+    # scene.add(light, pose=  np.eye(4))
+
+    # scene.add(camera, pose=[[ 1,  0,  0,  0.02],
+    #                         [ 0,  1,  0,  0.21],
+    #                         [ 0,  0,  1,  0.19],
+    #                         [ 0,  0,  0,  1]])
+
+    # # render scene
+    # r = pyrender.OffscreenRenderer(1024, 1024)
+    # color, _ = r.render(scene)
     
 
     str__ = f"Training sample {sample_id_str}\n"
@@ -105,7 +122,7 @@ def sample_info(sample_id_str, fieldn):
         for fname in hf_dataset.description['out_fields_names']:
             str__ += f"- {fname}\n"
 
-    return str__, color
+    return str__, "./visu.obj"
 
 
 if __name__ == "__main__":
@@ -118,7 +135,8 @@ if __name__ == "__main__":
                 output1 = gr.Text(label="Training sample info")
             with gr.Column():
                 d2 = gr.Dropdown(field_names_train, value=field_names_train[0], label="Field name")        
-                output2 = gr.Image(label="Training sample visualization")
+                # output2 = gr.Image(label="Training sample visualization")    
+                output2 = gr.Model3D(label="Training sample visualization")
                 
         d1.input(sample_info, [d1, d2], [output1, output2])
         d2.input(sample_info, [d1, d2], [output1, output2])