ecopus/mesh-refinement-agent
0
1## MVP 1 FILE!2import gradio as gr3import os4import json5import tempfile6import gmsh7 8from mesh_service.run_interface import (9 get_result,10 extract_faces_for_gui # <-- NEW helper from updated run_interface11)12 13# Absolute path to example files inside Docker container14os.chdir(os.path.dirname(__file__))15EXAMPLE_DIR = os.path.join(os.path.dirname(__file__), "example_steps")16 17 18# -----------------------------------------19# Gmsh initialize only once20# -----------------------------------------21try:22 gmsh.initialize()23except:24 pass25 26 27# -----------------------------------------28# STEP Upload → Extract Faces for GUI29# -----------------------------------------30def on_step_upload(step_file):31 if step_file is None:32 return gr.update(choices=[], value=None), "[]"33 34 # Work in temp copy if needed35 if hasattr(step_file, "read"):36 with tempfile.NamedTemporaryFile(delete=False, suffix=".step") as f:37 f.write(step_file.read())38 step_path = f.name39 else:40 step_path = step_file.name41 42 faces = extract_faces_for_gui(step_path)43 labels = [f["label"] for f in faces]44 45 faces_json = json.dumps(faces)46 #default = labels[0] if labels else None47 default = None48 49 return gr.update(choices=labels, value=default), faces_json50 51 52# -----------------------------------------53# Main pipeline call54# -----------------------------------------55def run_mesh(step_file, thickness, load_type, direction, scale,56 selected_face_label, faces_json, view_mode):57 58 # 1. Handle STEP file path59 if hasattr(step_file, "read"):60 # Cloud / HuggingFace mode61 with tempfile.NamedTemporaryFile(delete=False, suffix=".step") as f:62 f.write(step_file.read())63 step_path = f.name64 else:65 step_path = step_file.name66 67 # 2. Convert direction text → number68 direction_val = 1 if direction == "positive" else -169 70 # 3. Run model pipeline71 sentence, global_png, refined_png, fig, msh_path = get_result(72 step_path=step_path,73 thickness=float(thickness),74 load_face=load_type, # still used as categorical feature75 load_direction=direction_val,76 load_scale=scale,77 selected_face_label=selected_face_label, # NEW78 faces_json=faces_json, # NEW79 view_mode=view_mode80 )81 82 # 4. Convert PNG paths to absolute (for Gradio)83 refined_img = os.path.abspath(refined_png)84 global_img = os.path.abspath(global_png)85 msh_path = os.path.abspath(msh_path)86 87 return sentence, refined_img, global_img, fig, msh_path88 89 90# -----------------------------------------91# Gradio Interface92# -----------------------------------------93with gr.Blocks(title="Mesh Refinement Agent") as iface:94 95 gr.Markdown("## Mesh Refinement Prediction Demo")96 97 gr.Markdown(98'''### 📘 How to Use This Mesh Refinement Model99 1001. **Upload a STEP file** (`.step` / `.stp`).101OR: choose from one of the example STEP files at the bottom of the GUI and click *Run Mesh Prediction*1022. **Set basic parameters:**103 - *Thickness* (mm) 104 - *Load Type* (bend from below or tension) 105 - *Direction* (positive/negative) 106 - *Load Scale* (low/medium/high)1073. After uploading, the system automatically detects all faces. 108 Select the **Load Face** from the dropdown.1094. Choose a **3D visualization mode**:110 - **Wireframe** – full mesh structure 111 - **Highlight** – locally refined regions in red 112 - **Heatmap** – continuous refinement magnitude1135. Click **Run Mesh Prediction**.114 115The model will generate:116- A natural-language summary 117- PNGs of the global and refined meshes 118- An interactive 3D mesh viewer 119- A downloadable `.msh` file120 121This workflow reproduces the full ML-driven refinement pipeline used during training.'''122 )123 124 gr.HTML("""125<details style="margin-bottom: 18px;">126 <summary style="font-size: 18px; cursor: pointer; font-weight: bold;">127 How to Choose the Correct Load Face128 </summary>129 <div style="padding-left: 12px; margin-top: 10px;">130 131 <p>Different load types correspond to different physical loading scenarios.<br>132 After uploading your STEP file, the face dropdown will list <b>all detected faces</b>.<br>133 To ensure the ML model receives inputs aligned with its training distribution,134 please choose the face following these guidelines:</p>135 136 <h3><b>Load Type: bend_bottom</b></h3>137 <p>Use this when the bracket or part is <b>loaded from below</b>, such as:</p>138 <ul>139 <li>A shelf bracket being pressed upward</li>140 <li>A support arm being bent from underneath</li>141 </ul>142 143 <p><b>Choose the face that is physically on the bottom of the part</b>, i.e.:</p>144 <ul>145 <li>The largest downward-facing surface</li>146 <li>The surface that would contact a support or wall in real usage</li>147 </ul>148 149 <p>💡 <i>Tip:</i> In most STEP files this is a face with a downward normal 150 (<code>normal ≈ [0, 0, -1]</code>).</p>151 152 <hr>153 154 <h3><b>Load Type: tension</b></h3>155 <p>Use this when the bracket or part is pulled <b>backwards</b> or <b>outwards</b>, such as:</p>156 <ul>157 <li>A wall bracket being pulled away from a mounting plane</li>158 <li>A hook experiencing outward tension</li>159 </ul>160 161 <p><b>Choose the face that is physically the rear/back face of the part</b>, i.e.:</p>162 <ul>163 <li>The flat face that would mount to a wall</li>164 <li>A large face with outward-pointing orientation</li>165 </ul>166 167 <p>💡 <i>Tip:</i> Often the normal points along ±X 168 (<code>normal ≈ [±1, 0, 0]</code>).</p>169 170 <hr>171 172 <h3><b>If You Are Unsure</b></h3>173 <ul>174 <li>Choose the face that best represents the real-life load direction.</li>175 <li>Picking an incorrect face has <b>minimal effect</b> on the mesh output.</li>176 </ul>177 178 </div>179</details>180""")181 182 gr.HTML("""183<details style="margin-bottom: 18px;">184 <summary style="font-size: 18px; cursor: pointer; font-weight: bold;">185 3D Visualization Mode Guidance186 </summary>187 <div style="padding-left: 12px; margin-top: 10px;">188 189 <p>In the 3D visualization mode dropdown, there are three modes to choose from:</p>190 191 <ul>192 <li><b>wireframe</b>: 3D interactive wireframe of the final optimized mesh</li>193 <li><b>highlight</b>: wireframe with refined regions shown in red</li>194 <li><b>heatmap</b>: continuous heatmap quantifying refinement magnitude</li>195 </ul>196 197 <p><b>Heatmap color scale:</b></p>198 <ul>199 <li><b>Red</b>: very small predicted mesh size → highest refinement</li>200 <li><b>Blue</b>: very large predicted mesh size → coarsest regions</li>201 </ul>202 203 <p>The heatmap provides a smooth, global visualization of refinement intensity across the mesh.</p>204 205 </div>206</details>207""")208 209 210 with gr.Row():211 with gr.Column(scale=1):212 213 # STEP file upload214 step_file = gr.File(215 label="Upload STEP File",216 file_types=[".step", ".stp"]217 )218 219 thickness = gr.Textbox(220 label="Thickness (mm)",221 value="3"222 )223 224 load_type = gr.Dropdown(225 ["bend_bottom", "tension"],226 value="bend_bottom",227 label="Load Type (categorical feature)"228 )229 230 direction = gr.Dropdown(231 ["positive", "negative"],232 value="positive",233 label="Direction"234 )235 236 scale = gr.Dropdown(237 ["low", "medium", "high"],238 value="high",239 label="Load Scale"240 )241 242 # toggle visualization mode243 view_mode = gr.Dropdown(244 ["wireframe", "highlight", "heatmap"],245 value="wireframe",246 label="3D Visualization Mode"247 )248 249 # NEW: face selection & json hidden box250 selected_face_label = gr.Dropdown(251 label="Select Load Face (detected from STEP)",252 choices=[],253 value=None,254 interactive=True255 )256 257 faces_json_box = gr.Textbox(258 visible=False259 )260 261 examples_data = [262 [f"{EXAMPLE_DIR}/example1.step", "3", "bend_bottom", "positive", "high", None, "[]", "wireframe"],263 [f"{EXAMPLE_DIR}/example2.step", "2", "tension", "negative", "medium", None, "[]", "highlight"],264 [f"{EXAMPLE_DIR}/example3.step", "3", "bend_bottom", "positive", "high", None, "[]", "heatmap"],265 [f"{EXAMPLE_DIR}/example4.step", "3.5", "tension", "negative", "medium", None, "[]", "wireframe"],266 ]267 268 269 # Populate face dropdown upon STEP upload270 step_file.change(271 on_step_upload,272 inputs=step_file,273 outputs=[selected_face_label, faces_json_box]274 )275 276 run_button = gr.Button("Run Mesh Prediction")277 278 with gr.Column(scale=1):279 280 result_text = gr.Textbox(281 label="Result",282 lines=12,283 max_lines=20,284 interactive=False285 )286 287 refined_img = gr.Image(288 label="Refined Mesh (PNG)"289 )290 291 global_img = gr.Image(292 label="Global Mesh (PNG)"293 )294 295 # NEW: 3D interactive refined mesh296 mesh_plot_3d = gr.Plot(297 label="Refined Mesh (3D Interactive)"298 )299 300 mesh_download = gr.File(301 label="Download Generated Mesh (.msh)",302 interactive=False303 )304 305 gr.Examples(306 examples=examples_data,307 inputs=[308 step_file, thickness, load_type, direction, scale, 309 selected_face_label, faces_json_box, view_mode310 ],311 # The `run_on_click` parameter automatically runs the main function (run_mesh)312 # when an example is selected, giving immediate results.313 # If you prefer users click the 'Run' button manually after selection, set this to False.314 fn=run_mesh,315 outputs=[316 result_text,317 refined_img,318 global_img,319 mesh_plot_3d,320 mesh_download321 ],322 run_on_click=False 323 )324 325 # Connect compute button326 run_button.click(327 run_mesh,328 inputs=[329 step_file,330 thickness,331 load_type,332 direction,333 scale,334 selected_face_label,335 faces_json_box,336 view_mode,337 ],338 outputs=[339 result_text,340 refined_img,341 global_img,342 mesh_plot_3d,343 mesh_download344 ]345 )346 347 348iface.launch(server_name="0.0.0.0", server_port=7860)349