Spaces:
Sleeping
Sleeping
| from __future__ import annotations | |
| from pathlib import Path | |
| from typing import Any | |
| import gradio as gr | |
| from PIL import Image | |
| from stl_slicer import SliceStack, load_mesh, slice_stl_to_tiffs | |
| ViewerState = dict[str, Any] | |
| SAMPLE_STL_FILENAMES = ("Hollow_Pyramid.stl", "balanced_die.stl", "halfsphere.stl") | |
| SAMPLE_STL_DIR = Path(__file__).resolve().parent / "sample_stls" | |
| APP_CSS = """ | |
| .gradio-container { | |
| font-size: 90%; | |
| padding-top: 0.5rem !important; | |
| padding-bottom: 0.5rem !important; | |
| } | |
| .gradio-container .gr-row { | |
| gap: 0.5rem !important; | |
| } | |
| .gradio-container .gr-form, | |
| .gradio-container .gr-box, | |
| .gradio-container .block { | |
| padding: 0.4rem !important; | |
| } | |
| .gradio-container .prose { | |
| margin-bottom: 0.4rem !important; | |
| } | |
| """ | |
| def _read_slice_preview(path: str) -> Image.Image: | |
| with Image.open(path) as image: | |
| return image.copy() | |
| def _empty_state() -> ViewerState: | |
| return {"tiff_paths": [], "z_values": []} | |
| def _reset_slider() -> dict[str, Any]: | |
| return gr.update(minimum=0, maximum=0, value=0, step=1, interactive=False) | |
| def _stack_to_state(stack: SliceStack) -> ViewerState: | |
| return { | |
| "tiff_paths": [str(path) for path in stack.tiff_paths], | |
| "z_values": stack.z_values, | |
| } | |
| def _format_model_details(source_name: str, mesh) -> str: | |
| extents = mesh.extents | |
| return "\n".join( | |
| [ | |
| "### Model Details", | |
| f"- Source: `{source_name}`", | |
| f"- Extents: `{extents[0]:.3f} x {extents[1]:.3f} x {extents[2]:.3f}`", | |
| f"- Faces: `{len(mesh.faces)}`", | |
| f"- Vertices: `{len(mesh.vertices)}`", | |
| f"- Watertight: `{'yes' if mesh.is_watertight else 'no'}`", | |
| ] | |
| ) | |
| def _slice_label(state: ViewerState, index: int) -> str: | |
| path = Path(state["tiff_paths"][index]).name | |
| z_value = state["z_values"][index] | |
| total = len(state["tiff_paths"]) | |
| return f"Slice {index + 1} / {total} | z = {z_value:.4f} | {path}" | |
| def _render_selected_slice(state: ViewerState, index: int) -> tuple[str, Image.Image | None]: | |
| tiff_paths = state.get("tiff_paths", []) | |
| if not tiff_paths: | |
| return "No slice stack loaded yet.", None | |
| bounded_index = max(0, min(int(index), len(tiff_paths) - 1)) | |
| selected_path = tiff_paths[bounded_index] | |
| return ( | |
| _slice_label(state, bounded_index), | |
| _read_slice_preview(selected_path), | |
| ) | |
| def load_single_model(stl_file: str | None) -> tuple[str | None, str]: | |
| if not stl_file: | |
| return None, "No model loaded." | |
| mesh = load_mesh(stl_file) | |
| return stl_file, _format_model_details(Path(stl_file).name, mesh) | |
| def preload_sample_models() -> tuple: | |
| outputs: list[Any] = [] | |
| for filename in SAMPLE_STL_FILENAMES: | |
| stl_path = SAMPLE_STL_DIR / filename | |
| if not stl_path.exists(): | |
| outputs.extend([ | |
| None, | |
| None, | |
| f"Sample file not found: {stl_path}", | |
| ]) | |
| continue | |
| try: | |
| mesh = load_mesh(stl_path) | |
| except Exception as exc: | |
| outputs.extend([ | |
| str(stl_path), | |
| None, | |
| f"Failed to load sample model: {stl_path.name} ({exc})", | |
| ]) | |
| continue | |
| outputs.extend([ | |
| str(stl_path), | |
| str(stl_path), | |
| _format_model_details(stl_path.name, mesh), | |
| ]) | |
| return tuple(outputs) | |
| def generate_all_stacks( | |
| stl1: str | None, | |
| stl2: str | None, | |
| stl3: str | None, | |
| layer_height: float, | |
| pixel_size: float, | |
| progress: gr.Progress = gr.Progress(), | |
| ): | |
| files = [stl1, stl2, stl3] | |
| valid_count = max(1, sum(1 for f in files if f)) | |
| results: list = [] | |
| completed = 0 | |
| for stl_file in files: | |
| if not stl_file: | |
| results.extend([ | |
| _empty_state(), | |
| _reset_slider(), | |
| "No slice stack loaded yet.", | |
| None, | |
| None, | |
| ]) | |
| continue | |
| slot_offset = completed | |
| def report_progress(cur: int, tot: int, offset: int = slot_offset) -> None: | |
| progress( | |
| (offset + cur / tot) / valid_count, | |
| desc=f"Slicing object {offset + 1} of {valid_count}\u2026", | |
| ) | |
| stack = slice_stl_to_tiffs( | |
| stl_file, | |
| layer_height=layer_height, | |
| pixel_size=pixel_size, | |
| progress_callback=report_progress, | |
| ) | |
| state = _stack_to_state(stack) | |
| label, preview = _render_selected_slice(state, 0) | |
| slider = gr.update( | |
| minimum=0, | |
| maximum=max(0, len(stack.tiff_paths) - 1), | |
| value=0, | |
| step=1, | |
| interactive=len(stack.tiff_paths) > 1, | |
| ) | |
| results.extend([ | |
| state, | |
| slider, | |
| label, | |
| preview, | |
| str(stack.zip_path), | |
| ]) | |
| completed += 1 | |
| return tuple(results) | |
| def jump_to_slice(state: ViewerState, index: float) -> tuple[str, Image.Image | None]: | |
| return _render_selected_slice(state, int(index)) | |
| def shift_slice(state: ViewerState, index: float, delta: int) -> tuple[int, str, Image.Image | None]: | |
| tiff_paths = state.get("tiff_paths", []) | |
| if not tiff_paths: | |
| return 0, "No slice stack loaded yet.", None | |
| new_index = max(0, min(int(index) + delta, len(tiff_paths) - 1)) | |
| label, preview = _render_selected_slice(state, new_index) | |
| return new_index, label, preview | |
| def build_demo() -> gr.Blocks: | |
| with gr.Blocks(title="STL TIFF Slicer", css=APP_CSS) as demo: | |
| gr.Markdown( | |
| """ | |
| # STL to TIFF Slicer | |
| Upload up to three STL files, choose a shared layer height and XY pixel size, then generate TIFF stacks for all uploaded models. | |
| """ | |
| ) | |
| with gr.Row(): | |
| load_samples_button = gr.Button( | |
| "Load Sample STLs", | |
| variant="secondary", | |
| size="sm", | |
| min_width=140, | |
| scale=0, | |
| ) | |
| # --- Upload + 3D viewer row --- | |
| stl_files: list[gr.File] = [] | |
| model_viewers: list[gr.Model3D] = [] | |
| model_details_list: list[gr.Markdown] = [] | |
| with gr.Row(): | |
| for i in range(3): | |
| with gr.Column(): | |
| stl_file = gr.File( | |
| label=f"STL File {i + 1}", | |
| file_types=[".stl"], | |
| type="filepath", | |
| ) | |
| model_viewer = gr.Model3D( | |
| label=f"3D Viewer {i + 1}", | |
| display_mode="solid", | |
| clear_color=(0.94, 0.95, 0.97, 1.0), | |
| height=270, | |
| ) | |
| model_details = gr.Markdown(f"No model {i + 1} loaded.") | |
| stl_files.append(stl_file) | |
| model_viewers.append(model_viewer) | |
| model_details_list.append(model_details) | |
| # --- Shared slicing controls --- | |
| with gr.Row(): | |
| layer_height = gr.Number(label="Layer Height", value=0.1, minimum=0.0001, step=0.01) | |
| pixel_size = gr.Number(label="Pixel Size", value=0.05, minimum=0.0001, step=0.01) | |
| generate_button = gr.Button("Generate TIFF Stacks", variant="primary") | |
| # --- Per-object slice browsers --- | |
| states: list[gr.State] = [] | |
| sliders: list[gr.Slider] = [] | |
| slice_labels: list[gr.Markdown] = [] | |
| slice_previews: list[gr.Image] = [] | |
| download_zips: list[gr.File] = [] | |
| with gr.Row(): | |
| for i in range(3): | |
| with gr.Column(): | |
| slice_label = gr.Markdown("No slice stack loaded yet.") | |
| slice_preview = gr.Image( | |
| label=f"Slice Preview {i + 1}", | |
| type="pil", | |
| image_mode="L", | |
| height=324, | |
| ) | |
| with gr.Row(): | |
| prev_button = gr.Button("\u25c4 Prev") | |
| next_button = gr.Button("Next \u25ba") | |
| slice_slider = gr.Slider( | |
| label="Slice Index", | |
| minimum=0, | |
| maximum=0, | |
| value=0, | |
| step=1, | |
| interactive=False, | |
| ) | |
| download_zip = gr.File(label=f"Download TIFF ZIP {i + 1}") | |
| state = gr.State(_empty_state()) | |
| slice_labels.append(slice_label) | |
| slice_previews.append(slice_preview) | |
| sliders.append(slice_slider) | |
| download_zips.append(download_zip) | |
| states.append(state) | |
| slice_slider.release( | |
| fn=jump_to_slice, | |
| inputs=[state, slice_slider], | |
| outputs=[slice_label, slice_preview], | |
| queue=False, | |
| ) | |
| prev_button.click( | |
| fn=lambda sv, idx: shift_slice(sv, idx, -1), | |
| inputs=[state, slice_slider], | |
| outputs=[slice_slider, slice_label, slice_preview], | |
| queue=False, | |
| ) | |
| next_button.click( | |
| fn=lambda sv, idx: shift_slice(sv, idx, 1), | |
| inputs=[state, slice_slider], | |
| outputs=[slice_slider, slice_label, slice_preview], | |
| queue=False, | |
| ) | |
| # --- File upload handlers --- | |
| for i in range(3): | |
| stl_files[i].change( | |
| fn=load_single_model, | |
| inputs=stl_files[i], | |
| outputs=[model_viewers[i], model_details_list[i]], | |
| ) | |
| # --- Generate button --- | |
| generate_outputs: list = [] | |
| for i in range(3): | |
| generate_outputs.extend([ | |
| states[i], | |
| sliders[i], | |
| slice_labels[i], | |
| slice_previews[i], | |
| download_zips[i], | |
| ]) | |
| preload_outputs: list = [] | |
| for i in range(3): | |
| preload_outputs.extend([ | |
| stl_files[i], | |
| model_viewers[i], | |
| model_details_list[i], | |
| ]) | |
| load_samples_button.click( | |
| fn=preload_sample_models, | |
| outputs=preload_outputs, | |
| ) | |
| generate_button.click( | |
| fn=generate_all_stacks, | |
| inputs=[stl_files[0], stl_files[1], stl_files[2], layer_height, pixel_size], | |
| outputs=generate_outputs, | |
| ) | |
| return demo | |
| demo = build_demo() | |
| if __name__ == "__main__": | |
| demo.launch(ssr_mode=False) | |