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 from utils_inference import infer import os # switch to "osmesa" or "egl" before loading pyrender os.environ["PYOPENGL_PLATFORM"] = "egl" hf_dataset = load_dataset("PLAID-datasets/AirfRANS_remeshed", split="all_samples") file = open('training_data.pkl', 'rb') training_data = pickle.load(file) file.close() train_ids = hf_dataset.description['split']['ML4PhySim_Challenge_train'] out_fields_names = hf_dataset.description['out_fields_names'] in_scalars_names = hf_dataset.description['in_scalars_names'] out_scalars_names = hf_dataset.description['out_scalars_names'] nb_samples = len(hf_dataset) #

MMGP demo on the AirfRANS_remeshed dataset

# MMGP paper, _HEADER_ = '''

MMGP demo on the AirfRANS_remeshed dataset

''' _HEADER_2 = ''' The model is already trained. The morphing is the same as the one used in the [MMGP paper](https://arxiv.org/abs/2305.12871), but is much less involved than the one used in the winning entry of the [ML4PhySim competition](https://www.codabench.org/competitions/1534/). The training set has 103 samples and is the one used in this competition (some evaluations are out-of-distribution). The inference takes approx 5 seconds, and is done from scratch (no precomputation is used during the inference when evaluating samples). This means that the morphing and the finite element interpolations are re-done at each evaluation. After choosing a sample id, please change the field name in the dropdown menu to update the visualization. ''' def round_num(num)->str: return '%s' % float('%.3g' % num) def compute_inference(sample_id_str): sample_id = int(sample_id_str) sample_ = hf_dataset[sample_id]["sample"] plaid_sample = Sample.model_validate(pickle.loads(sample_)) prediction = infer(hf_dataset, sample_id, training_data) reference = {fieldn:plaid_sample.get_field(fieldn) for fieldn in out_fields_names} nodes = plaid_sample.get_nodes() 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'] trimesh = Trimesh(vertices = nodes, faces = triangles) file = open('computed_inference.pkl', 'wb') pickle.dump([trimesh, reference, prediction], file) file.close() str__ = f"Training sample {sample_id_str}" if sample_id in train_ids: str__ += " (in the training set)\n\n" else: str__ += " (not in the training set)\n\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" str__ += f"\nNumber of nodes in the mesh: {nodes.shape[0]}" return str__ def show_pred(fieldn): file = open('computed_inference.pkl', 'rb') data = pickle.load(file) file.close() trimesh, reference, prediction = data[0], data[1], data[2] ref = reference[fieldn] pred = prediction[fieldn] norm = mpl.colors.Normalize(vmin=np.min(ref), vmax=np.max(ref)) cmap = cm.seismic#cm.coolwarm m = cm.ScalarMappable(norm=norm, cmap=cmap) vertex_colors = m.to_rgba(pred)[:,:3] 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, 1], [ 0, 1, 0, 0], [ 0, 0, 1, 6], [ 0, 0, 0, 1]]) # render scene r = pyrender.OffscreenRenderer(1024, 1024) color, _ = r.render(scene) return color def show_ref(fieldn): file = open('computed_inference.pkl', 'rb') data = pickle.load(file) file.close() trimesh, reference, prediction = data[0], data[1], data[2] ref = reference[fieldn] norm = mpl.colors.Normalize(vmin=np.min(ref), vmax=np.max(ref)) cmap = cm.seismic#cm.coolwarm m = cm.ScalarMappable(norm=norm, cmap=cmap) vertex_colors = m.to_rgba(ref)[:,:3] 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, 1], [ 0, 1, 0, 0], [ 0, 0, 1, 6], [ 0, 0, 0, 1]]) # render scene r = pyrender.OffscreenRenderer(1024, 1024) color, _ = r.render(scene) return color def show_err(fieldn): file = open('computed_inference.pkl', 'rb') data = pickle.load(file) file.close() trimesh, reference, prediction = data[0], data[1], data[2] ref = reference[fieldn] pred = prediction[fieldn] norm = mpl.colors.Normalize(vmin=np.min(ref), vmax=np.max(ref)) cmap = cm.seismic#cm.coolwarm m = cm.ScalarMappable(norm=norm, cmap=cmap) vertex_colors = m.to_rgba(pred-ref)[:,:3] 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, 1], [ 0, 1, 0, 0], [ 0, 0, 1, 6], [ 0, 0, 0, 1]]) # render scene r = pyrender.OffscreenRenderer(1024, 1024) color, _ = r.render(scene) return color if __name__ == "__main__": with gr.Blocks() as demo: # trimesh, reference, prediction = compute_inference(0) gr.Markdown(_HEADER_) gr.Markdown(_HEADER_2) with gr.Row(variant="panel"): with gr.Column(): d1 = gr.Slider(0, nb_samples-1, value=0, label="Training sample id", info="Choose between 0 and "+str(nb_samples-1)) # output1 = gr.Text(label="Inference status") output4 = gr.Text(label="Information on sample") output5 = gr.Image(label="Error") with gr.Column(): d2 = gr.Dropdown(out_fields_names, value=out_fields_names[0], label="Field name") output2 = gr.Image(label="Reference") output3 = gr.Image(label="MMGP prediction") # d1.input(compute_inference, [d1], [output1, output4]) d1.input(compute_inference, [d1], [output4]) d2.input(show_ref, [d2], [output2]) d2.input(show_pred, [d2], [output3]) d2.input(show_err, [d2], [output5]) demo.launch()