from __future__ import annotations import math import tempfile import uuid import zipfile from dataclasses import dataclass from pathlib import Path from typing import Callable import numpy as np from PIL import Image, ImageDraw from shapely.geometry import GeometryCollection, MultiPolygon, Polygon import trimesh ProgressCallback = Callable[[int, int], None] | None @dataclass(slots=True) class SliceStack: output_dir: Path zip_path: Path tiff_paths: list[Path] z_values: list[float] image_size: tuple[int, int] bounds: tuple[tuple[float, float, float], tuple[float, float, float]] layer_height: float pixel_size: float def load_mesh(stl_path: str | Path) -> trimesh.Trimesh: loaded = trimesh.load(stl_path, force="scene") if isinstance(loaded, trimesh.Scene): if not loaded.geometry: raise ValueError("The STL file does not contain any mesh geometry.") mesh = trimesh.util.concatenate(tuple(loaded.geometry.values())) else: mesh = loaded if not isinstance(mesh, trimesh.Trimesh) or mesh.is_empty: raise ValueError("Unable to load a valid mesh from the STL file.") return mesh def calculate_z_levels(z_min: float, z_max: float, layer_height: float) -> list[float]: if layer_height <= 0: raise ValueError("Layer height must be greater than zero.") thickness = z_max - z_min if thickness <= 0: return [z_min] layer_count = max(1, math.ceil(thickness / layer_height)) top_guard = math.nextafter(z_max, z_min) return [ min(z_min + ((index + 0.5) * layer_height), top_guard) for index in range(layer_count) ] def _to_pixel_ring( coords: np.ndarray, x_min: float, y_min: float, pixel_size: float, image_height: int, ) -> list[tuple[int, int]]: pixels: list[tuple[int, int]] = [] for x_value, y_value in coords: x_pixel = int(round((float(x_value) - x_min) / pixel_size)) y_pixel = int(round((float(y_value) - y_min) / pixel_size)) pixels.append((x_pixel, image_height - 1 - y_pixel)) return pixels def _ring_to_world_xy(ring_coords: object, to_3d: np.ndarray) -> np.ndarray: planar = np.asarray(ring_coords, dtype=float) if planar.ndim != 2 or planar.shape[1] < 2: raise ValueError("Encountered an invalid polygon ring while slicing.") planar_3d = np.column_stack([planar[:, 0], planar[:, 1], np.zeros(len(planar))]) world = trimesh.transform_points(planar_3d, to_3d) return world[:, :2] def _compose_even_odd_polygons(polygons: list[Polygon]) -> list[Polygon]: geometry: Polygon | MultiPolygon | GeometryCollection | None = None for polygon in polygons: geometry = polygon if geometry is None else geometry.symmetric_difference(polygon) if geometry is None or geometry.is_empty: return [] if isinstance(geometry, Polygon): return [geometry] if isinstance(geometry, MultiPolygon): return list(geometry.geoms) if isinstance(geometry, GeometryCollection): return [geom for geom in geometry.geoms if isinstance(geom, Polygon) and not geom.is_empty] return [] def _extract_world_polygons(section: trimesh.path.Path3D) -> list[tuple[np.ndarray, list[np.ndarray]]]: if hasattr(section, "to_2D"): planar, to_3d = section.to_2D() else: planar, to_3d = section.to_planar() composed_polygons = _compose_even_odd_polygons(list(planar.polygons_closed)) polygons: list[tuple[np.ndarray, list[np.ndarray]]] = [] for polygon in composed_polygons: exterior = _ring_to_world_xy(polygon.exterior.coords, to_3d) holes = [_ring_to_world_xy(interior.coords, to_3d) for interior in polygon.interiors] polygons.append((exterior, holes)) return polygons def _render_slice( section: trimesh.path.Path3D | None, x_min: float, y_min: float, image_size: tuple[int, int], pixel_size: float, ) -> Image.Image: image = Image.new("L", image_size, 0) if section is None: return image polygons = _extract_world_polygons(section) if not polygons: return image draw = ImageDraw.Draw(image) for exterior, holes in polygons: draw.polygon( _to_pixel_ring(exterior, x_min, y_min, pixel_size, image.height), fill=255, ) for hole in holes: draw.polygon( _to_pixel_ring(hole, x_min, y_min, pixel_size, image.height), fill=0, ) return image def _make_output_paths(stl_path: Path, output_root: str | Path | None) -> tuple[Path, Path]: root = Path(output_root) if output_root else Path(tempfile.mkdtemp(prefix="stl_slices_")) stem = stl_path.stem or "mesh" job_dir = root / f"{stem}_{uuid.uuid4().hex[:8]}" slices_dir = job_dir / "tiff_slices" slices_dir.mkdir(parents=True, exist_ok=True) zip_path = job_dir / f"{stem}_tiff_slices.zip" return slices_dir, zip_path def _zip_tiffs(tiff_paths: list[Path], zip_path: Path) -> None: with zipfile.ZipFile(zip_path, mode="w", compression=zipfile.ZIP_DEFLATED) as archive: for tiff_path in tiff_paths: archive.write(tiff_path, arcname=tiff_path.name) def slice_stl_to_tiffs( stl_path: str | Path, layer_height: float, pixel_size: float, output_root: str | Path | None = None, progress_callback: ProgressCallback = None, ) -> SliceStack: if pixel_size <= 0: raise ValueError("Pixel size must be greater than zero.") stl_path = Path(stl_path) mesh = load_mesh(stl_path) bounds = mesh.bounds (x_min, y_min, z_min), (x_max, y_max, z_max) = bounds z_values = calculate_z_levels(float(z_min), float(z_max), layer_height) width = max(1, math.ceil((float(x_max) - float(x_min)) / pixel_size) + 1) height = max(1, math.ceil((float(y_max) - float(y_min)) / pixel_size) + 1) image_size = (width, height) output_dir, zip_path = _make_output_paths(stl_path, output_root) tiff_paths: list[Path] = [] for index, z_value in enumerate(z_values): section = mesh.section( plane_origin=np.array([0.0, 0.0, z_value], dtype=float), plane_normal=np.array([0.0, 0.0, 1.0], dtype=float), ) image = _render_slice(section, float(x_min), float(y_min), image_size, pixel_size) tiff_path = output_dir / f"slice_{index:04d}.tif" image.save(tiff_path, compression="tiff_deflate") tiff_paths.append(tiff_path) if progress_callback is not None: progress_callback(index + 1, len(z_values)) _zip_tiffs(tiff_paths, zip_path) return SliceStack( output_dir=output_dir, zip_path=zip_path, tiff_paths=tiff_paths, z_values=z_values, image_size=image_size, bounds=( (float(x_min), float(y_min), float(z_min)), (float(x_max), float(y_max), float(z_max)), ), layer_height=layer_height, pixel_size=pixel_size, )