app.py

import gradio as gr
import pickle
from datasets import load_dataset
from plaid.containers.sample import Sample


import numpy as np
# import pyrender
# from trimesh import Trimesh
import matplotlib as mpl
import matplotlib.cm as cm

import os
# switch to "osmesa" or "egl" before loading pyrender
os.environ["PYOPENGL_PLATFORM"] = "egl"


hf_dataset = load_dataset("PLAID-datasets/2D_ElastoPlastoDynamics", split="all_samples")
sample0 = hf_dataset[0]["sample"]
plaid_sample0 = Sample.model_validate(pickle.loads(sample0))
times_steps = list(plaid_sample0._meshes.keys()) 

ids_train = hf_dataset.description["split"]["train"]
nb_samples = len(ids_train)


field_names_train = ["U_x", "U_y"]#, "Sigma_xx", "Sigma_xy", "Sigma_yy", "Epsilon_xx", "Epsilon_xy", "Epsilon_yy", "EROSION_STATUS"]


_HEADER_ = '''
<h2><b>Visualization demo of <a href='https://huggingface.co/datasets/PLAID-datasets/2D_ElastoPlastoDynamics' target='_blank'><b>2D_ElastoPlastoDynamics</b></b></h2>
'''


def round_num(num)->str:
    return '%s' % float('%.3g' % num)

def sample_info(sample_id_str, time_step, fieldn):

    time = times_steps[int(time_step)]

    sample_ = hf_dataset[ids_train[int(sample_id_str)]]["sample"]
    plaid_sample = Sample.model_validate(pickle.loads(sample_))
    # plaid_sample = Sample.load_from_dir(f"Tensile2d/dataset/samples/sample_"+str(sample_id_str).zfill(9))

    nodes = plaid_sample.get_nodes()
    if fieldn in ["U_x", "U_y"]:
        field = plaid_sample.get_field(fieldn, time=time)
    else:
        field = plaid_sample.get_field(fieldn, location="CellCenter", time=time)        
    # 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]
    
    triangles = plaid_sample.get_elements()['TRI_3']
    nb_tris = triangles.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]} {0.} {rgb_colors[i,0]} {rgb_colors[i,1]} {rgb_colors[i,2]}\n")
            # f.write(f"v {-nodes[i,0]} {0.} {nodes[i,1]} {rgb_colors[i,0]} {rgb_colors[i,1]} {rgb_colors[i,2]}\n")
       
        for i in range(nb_tris):
            f.write(f"f {triangles[i,0] + 1} {triangles[i,1] + 1} {triangles[i,2] + 1} \n")
            
    # # generate colormap
    # if np.linalg.norm(field) > 0:
    #     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]
    # 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 = 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=1000.)

    # scene.add(mesh, pose=  np.eye(4))
    # scene.add(light, pose=  np.eye(4))

    # scene.add(camera, pose=[[ 1,  0,  0,  0],
    #                         [ 0,  1,  0,  0],
    #                         [ 0,  0,  1,  3],
    #                         [ 0,  0,  0,  1]])

    # # render scene
    # r = pyrender.OffscreenRenderer(1024, 1024)
    # color, _ = r.render(scene)

    
    str__ = f"Training sample {sample_id_str}\n"
    str__ += str(plaid_sample)+"\n"
    
    if len(hf_dataset.description['in_scalars_names'])>0:        
        str__ += "\ninput scalars:\n"
        for sname in hf_dataset.description['in_scalars_names']:
            str__ += f"- {sname}: {round_num(plaid_sample.get_scalar(sname))}\n"
    if len(hf_dataset.description['out_scalars_names'])>0:        
        str__ += "\noutput scalars:\n"
        for sname in hf_dataset.description['out_scalars_names']:
            str__ += f"- {sname}: {round_num(plaid_sample.get_scalar(sname))}\n"
    str__ += f"\n\nMesh number of nodes: {nodes.shape[0]}\n"
    if len(hf_dataset.description['in_fields_names'])>0:        
        str__ += "\ninput fields:\n"
        for fname in hf_dataset.description['in_fields_names']:
            str__ += f"- {fname}\n"
    if len(hf_dataset.description['out_fields_names'])>0:        
        str__ += "\noutput fields:\n"
        for fname in hf_dataset.description['out_fields_names']:
            str__ += f"- {fname}\n"

    return str__, "./visu.obj"


if __name__ == "__main__":

    with gr.Blocks(fill_width=True) as demo:
        gr.Markdown(_HEADER_)
        with gr.Row(variant="panel"):
            with gr.Column(scale=1):
                d1 = gr.Slider(0, nb_samples-1, value=0, label="Training sample id", info="Choose between 0 and "+str(nb_samples-1))
                d2 = gr.Slider(0, len(times_steps)-1, value=0, label="Time step", info="Choose between 0 and "+str(len(times_steps)-1))
                output1 = gr.Text(label="Training sample info")
            with gr.Column(scale=2, min_width=300):
                d3 = gr.Dropdown(field_names_train, value=field_names_train[0], label="Field name")        
                # output2 = gr.Image(label="Training sample visualization")    
                output2 = gr.Model3D(label="Training sample visualization")
                
        d1.input(sample_info, [d1, d2, d3], [output1, output2])
        d2.input(sample_info, [d1, d2, d3], [output1, output2]) 
        d3.input(sample_info, [d1, d2, d3], [output1, output2]) 

    demo.launch()