Update app.py

app.py

@@ -1,25 +1,28 @@
 import gradio as gr
-import pickle
+# import pickle
 # from datasets import load_from_disk
 from plaid.containers.sample import Sample
 # import pyvista as pv
 
+import numpy as np
 import pyrender
 import trimesh
-import matplotlib.pyplot as plt
+import matplotlib as mpl
+import matplotlib.cm as cm
 
 import os
 # switch to "osmesa" or "egl" before loading pyrender
 os.environ["PYOPENGL_PLATFORM"] = "egl"
 
-import numpy as np
+
+os.system("wget https://zenodo.org/records/10124594/files/Tensile2d.tar.gz")
+os.system("tar -xvf Tensile2d.tar.gz")
 
 # FOLDER = "plot"
 
 # dataset = load_from_disk("Rotor37")
 
-field_names_train = ["Temperature", "Pressure", "Density"]#pickle.loads(dataset[0]["sample"]).get_field_names()
-field_names_test  = []
+field_names_train = ["sig11", "sig22", "sig12", "U1", "U2", "evrcum"]
 
 
 
@@ -29,45 +32,51 @@ field_names_test  = []
 
 def sample_info(sample_id_str, fieldn):
 
-    # sample_id = int(sample_id_str)
-    # plaid_sample = Sample.load_from_dir(f"Rotor37/dataset/samples/sample_"+str(sample_id_str).zfill(9))
+    plaid_sample = Sample.load_from_dir(f"Tensile2d/dataset/samples/sample_"+str(sample_id_str).zfill(9))
+    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__
+
+
+    triangles = plaid_sample.get_elements()['TRI_3']
+
+    # generate colormap
+    norm = mpl.colors.Normalize(vmin=np.min(field), vmax=np.max(field))
+    cmap = cm.coolwarm
+
+    m = cm.ScalarMappable(norm=norm, cmap=cmap)
+    vertex_colors = m.to_rgba(field)[:,:3]
 
-    str__ = f"loading sample {sample_id_str}"
-    
-    
     # generate mesh
-    sphere = trimesh.creation.icosphere(subdivisions=4, radius=0.8)
-    sphere.vertices+=1e-2*np.random.randn(*sphere.vertices.shape)
-    mesh = pyrender.Mesh.from_trimesh(sphere, smooth=False)
-    
+    trimesh = Trimesh(vertices = nodes, faces = triangles)
+    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 / 3.0)
-    light = pyrender.DirectionalLight(color=[1,1,1], intensity=2e3)
-    
+    light = pyrender.DirectionalLight(color=[1,1,1], intensity=1000.)
+
     scene.add(mesh, pose=  np.eye(4))
     scene.add(light, pose=  np.eye(4))
-    
-    c = 2**-0.5
+
+    c = 3**-0.5
     scene.add(camera, pose=[[ 1,  0,  0,  0],
                             [ 0,  c, -c, -2],
-                            [ 0,  c,  c,  2],
+                            [ 0,  c,  c,  1.2],
                             [ 0,  0,  0,  1]])
-    
+
     # render scene
-    r = pyrender.OffscreenRenderer(512, 512)
+    r = pyrender.OffscreenRenderer(1024, 1024)
     color, _ = r.render(scene)
-    # color = np.random.rand(512, 512)
-    
-    plt.figure(figsize=(8,8))
-    plt.imshow(color)
     
-    plt.savefig("test.png")
-
 
+    str__ = f"loading sample {sample_id_str}"
 
-    return str__, "test.png"
-    # return str__, str__
+    return str__, color
 
 
 if __name__ == "__main__":
@@ -78,10 +87,8 @@ if __name__ == "__main__":
 
 
         output1 = gr.Text(label="Training sample info")
-        # output2 = gr.Text(label="Training sample visualization")
         output2 = gr.Image(label="Training sample visualization")
 
-        # d1.input(update_second, d1, d2)
 
         d1.input(sample_info, [d1, d2], [output1, output2])
         d2.input(sample_info, [d1, d2], [output1, output2])