dawood/mesh-test
1
1import spaces2import subprocess3# Install flash attention, skipping CUDA build if necessary4subprocess.run(5 "pip install flash-attn --no-build-isolation",6 env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},7 shell=True,8)9import os10import torch11import trimesh12from accelerate.utils import set_seed13from accelerate import Accelerator14import numpy as np15import gradio as gr16from main import get_args, load_model17from mesh_to_pc import process_mesh_to_pc18import time19import matplotlib.pyplot as plt20from mpl_toolkits.mplot3d.art3d import Poly3DCollection21from PIL import Image22import io23 24args = get_args()25model = load_model(args)26 27device = torch.device('cuda')28accelerator = Accelerator(29 mixed_precision="fp16",30)31model = accelerator.prepare(model)32model.eval()33print("Model loaded to device")34 35def wireframe_render(mesh):36 views = [37 (90, 20), (270, 20)38 ]39 mesh.vertices = mesh.vertices[:, [0, 2, 1]]40 41 bounding_box = mesh.bounds42 center = mesh.centroid43 scale = np.ptp(bounding_box, axis=0).max()44 45 fig = plt.figure(figsize=(10, 10))46 47 # Function to render and return each view as an image48 def render_view(mesh, azimuth, elevation):49 ax = fig.add_subplot(111, projection='3d')50 ax.set_axis_off()51 52 # Extract vertices and faces for plotting53 vertices = mesh.vertices54 faces = mesh.faces55 56 # Plot faces57 ax.add_collection3d(Poly3DCollection(58 vertices[faces],59 facecolors=(0.8, 0.5, 0.2, 1.0), # Brownish yellow60 edgecolors='k',61 linewidths=0.5,62 ))63 64 # Set limits and center the view on the object65 ax.set_xlim(center[0] - scale / 2, center[0] + scale / 2)66 ax.set_ylim(center[1] - scale / 2, center[1] + scale / 2)67 ax.set_zlim(center[2] - scale / 2, center[2] + scale / 2)68 69 # Set view angle70 ax.view_init(elev=elevation, azim=azimuth)71 72 # Save the figure to a buffer73 buf = io.BytesIO()74 plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0, dpi=300)75 plt.clf()76 buf.seek(0)77 78 return Image.open(buf)79 80 # Render each view and store in a list81 images = [render_view(mesh, az, el) for az, el in views]82 83 # Combine images horizontally84 widths, heights = zip(*(i.size for i in images))85 total_width = sum(widths)86 max_height = max(heights)87 88 combined_image = Image.new('RGBA', (total_width, max_height))89 90 x_offset = 091 for img in images:92 combined_image.paste(img, (x_offset, 0))93 x_offset += img.width94 95 # Save the combined image96 save_path = f"combined_mesh_view_{int(time.time())}.png"97 combined_image.save(save_path)98 99 plt.close(fig)100 return save_path101 102@spaces.GPU(duration=300)103def do_inference(input_3d, sample_seed=0, do_sampling=False, do_marching_cubes=False):104 set_seed(sample_seed)105 print("Seed value:", sample_seed)106 107 input_mesh = trimesh.load(input_3d)108 pc_list, mesh_list = process_mesh_to_pc([input_mesh], marching_cubes = do_marching_cubes)109 pc_normal = pc_list[0] # 4096, 6110 mesh = mesh_list[0]111 vertices = mesh.vertices112 113 pc_coor = pc_normal[:, :3]114 normals = pc_normal[:, 3:]115 116 bounds = np.array([vertices.min(axis=0), vertices.max(axis=0)])117 # scale mesh and pc118 vertices = vertices - (bounds[0] + bounds[1])[None, :] / 2119 vertices = vertices / (bounds[1] - bounds[0]).max()120 mesh.vertices = vertices121 pc_coor = pc_coor - (bounds[0] + bounds[1])[None, :] / 2122 pc_coor = pc_coor / (bounds[1] - bounds[0]).max()123 124 mesh.merge_vertices()125 mesh.update_faces(mesh.unique_faces())126 mesh.fix_normals()127 if mesh.visual.vertex_colors is not None:128 orange_color = np.array([255, 165, 0, 255], dtype=np.uint8)129 130 mesh.visual.vertex_colors = np.tile(orange_color, (mesh.vertices.shape[0], 1))131 else:132 orange_color = np.array([255, 165, 0, 255], dtype=np.uint8)133 mesh.visual.vertex_colors = np.tile(orange_color, (mesh.vertices.shape[0], 1))134 input_save_name = f"processed_input_{int(time.time())}.obj"135 mesh.export(input_save_name)136 137 pc_coor = pc_coor / np.abs(pc_coor).max() * 0.9995 # input should be from -1 to 1138 assert (np.linalg.norm(normals, axis=-1) > 0.99).all(), "normals should be unit vectors, something wrong"139 normalized_pc_normal = np.concatenate([pc_coor, normals], axis=-1, dtype=np.float16)140 141 input = torch.tensor(normalized_pc_normal, dtype=torch.float16, device=device)[None]142 print("Data loaded")143 144 # with accelerator.autocast():145 with accelerator.autocast():146 outputs = model(input, do_sampling)147 print("Model inference done")148 recon_mesh = outputs[0]149 150 recon_mesh = recon_mesh[~torch.isnan(recon_mesh[:, 0, 0])] # nvalid_face x 3 x 3151 vertices = recon_mesh.reshape(-1, 3).cpu()152 vertices_index = np.arange(len(vertices)) # 0, 1, ..., 3 x face153 triangles = vertices_index.reshape(-1, 3)154 155 artist_mesh = trimesh.Trimesh(vertices=vertices, faces=triangles, force="mesh",156 merge_primitives=True)157 artist_mesh.merge_vertices()158 artist_mesh.update_faces(artist_mesh.unique_faces())159 artist_mesh.fix_normals()160 161 if artist_mesh.visual.vertex_colors is not None:162 orange_color = np.array([255, 165, 0, 255], dtype=np.uint8)163 164 artist_mesh.visual.vertex_colors = np.tile(orange_color, (artist_mesh.vertices.shape[0], 1))165 else:166 orange_color = np.array([255, 165, 0, 255], dtype=np.uint8)167 artist_mesh.visual.vertex_colors = np.tile(orange_color, (artist_mesh.vertices.shape[0], 1))168 169 num_faces = len(artist_mesh.faces)170 171 brown_color = np.array([165, 42, 42, 255], dtype=np.uint8)172 face_colors = np.tile(brown_color, (num_faces, 1))173 174 artist_mesh.visual.face_colors = face_colors175 # add time stamp to avoid cache176 save_name = f"output_{int(time.time())}.obj"177 artist_mesh.export(save_name)178 return input_save_name, input_save_name, save_name, save_name179 180 181_HEADER_ = '''182<h2><b>Official 🤗 Gradio Demo</b></h2><h2><a href='https://github.com/buaacyw/MeshAnything' target='_blank'><b>MeshAnything: Artist-Created Mesh Generation with Autoregressive Transformers</b></a></h2>183 184**MeshAnything** converts any 3D representation into meshes created by human artists, i.e., Artist-Created Meshes (AMs).185 186Code: <a href='https://github.com/buaacyw/MeshAnything' target='_blank'>GitHub</a>. Arxiv Paper: <a href='https://arxiv.org/abs/2406.10163' target='_blank'>ArXiv</a>.187 188❗️❗️❗️**Important Notes:**189- Gradio doesn't support interactive wireframe rendering currently. For interactive mesh visualization, please use download the obj file and open it with MeshLab or https://3dviewer.net/.190- The input mesh will be normalized to a unit bounding box. The up vector of the input mesh should be +Y for better results. Click **Preprocess with Marching Cubes** if the input mesh is a manually created mesh.191- Limited by computational resources, MeshAnything is trained on meshes with fewer than 800 faces and cannot generate meshes with more than 800 faces. The shape of the input mesh should be sharp enough; otherwise, it will be challenging to represent it with only 800 faces. Thus, feed-forward image-to-3D methods may often produce bad results due to insufficient shape quality.192- For point cloud input, please refer to our github repo <a href='https://github.com/buaacyw/MeshAnything' target='_blank'>GitHub</a>.193'''194 195 196_CITE_ = r"""197If MeshAnything is helpful, please help to ⭐ the <a href='https://github.com/buaacyw/MeshAnything' target='_blank'>Github Repo</a>. Thanks!198---199📋 **License**200 201S-Lab-1.0 LICENSE. Please refer to the [LICENSE file](https://github.com/buaacyw/GaussianEditor/blob/master/LICENSE.txt) for details.202 203📧 **Contact**204 205If you have any questions, feel free to open a discussion or contact us at <b>yiwen002@e.ntu.edu.sg</b>.206 207"""208output_model_obj = gr.Model3D(209 label="Generated Mesh (OBJ Format)",210 clear_color=[1, 1, 1, 1],211)212preprocess_model_obj = gr.Model3D(213 label="Processed Input Mesh (OBJ Format)",214 clear_color=[1, 1, 1, 1],215)216input_wireframe_render = gr.Model3D(217 label="Wireframe Render of Processed Input Mesh",218 clear_color=[1,1,1,1],219 display_mode="wireframe"220 )221output_wireframe_render = gr.Model3D(222 label="Wireframe Render of Generated Mesh",223 clear_color=[1, 1, 1, 1],224 display_mode="wireframe"225)226with (gr.Blocks() as demo):227 gr.Markdown(_HEADER_)228 with gr.Row(variant="panel"):229 with gr.Column():230 with gr.Row():231 input_3d = gr.Model3D(232 label="Input Mesh",233 clear_color=[1,1,1,1],234 )235 236 with gr.Row():237 with gr.Group():238 do_marching_cubes = gr.Checkbox(label="Preprocess with Marching Cubes", value=False)239 do_sampling = gr.Checkbox(label="Random Sampling", value=False)240 sample_seed = gr.Number(value=0, label="Seed Value", precision=0)241 242 with gr.Row():243 submit = gr.Button("Generate", elem_id="generate", variant="primary")244 245 with gr.Row(variant="panel"):246 mesh_examples = gr.Examples(247 examples=[248 os.path.join("examples", img_name) for img_name in sorted(os.listdir("examples"))249 ],250 inputs=input_3d,251 outputs=[preprocess_model_obj, input_wireframe_render, output_model_obj, output_wireframe_render],252 fn=do_inference,253 cache_examples = "lazy",254 examples_per_page=10255 )256 with gr.Column():257 with gr.Row():258 input_wireframe_render.render()259 with gr.Row():260 with gr.Tab("OBJ"):261 preprocess_model_obj.render()262 with gr.Row():263 output_wireframe_render.render()264 with gr.Row():265 with gr.Tab("OBJ"):266 output_model_obj.render()267 with gr.Row():268 gr.Markdown('''Try click random sampling and different <b>Seed Value</b> if the result is unsatisfying''')269 270 gr.Markdown(_CITE_)271 272 mv_images = gr.State()273 274 submit.click(275 fn=do_inference,276 inputs=[input_3d, sample_seed, do_sampling, do_marching_cubes],277 outputs=[preprocess_model_obj, input_wireframe_render, output_model_obj, output_wireframe_render],278 )279 280demo.launch(share=True)