Newer
Older
CNCTools / ReferenceSurfaceGenerator / backend / app / processing.py
import trimesh
import numpy as np
import ezdxf
from scipy.spatial import ConvexHull
from scipy.interpolate import splprep, splev
import os
import logging

# Configure logging
logging.basicConfig(level=logging.INFO, format='[%(levelname)s] %(message)s')

def create_layered_curves_dxf(input_file, output_file, num_layers=20, num_points_per_layer=30):
    layer_errors = [] 
    try:
        yield {"status": "processing", "progress": 5, "message": "Loading and repairing mesh..."}
        loaded = trimesh.load(input_file)

        if isinstance(loaded, trimesh.Scene):
            meshes = [g for g in loaded.geometry.values() if isinstance(g, trimesh.Trimesh)]
            if not meshes:
                raise ValueError("No valid mesh geometry found in scene.")
            mesh = trimesh.util.concatenate(meshes)
        else:
            mesh = loaded

        mesh.update_faces(mesh.nondegenerate_faces())
        mesh.fill_holes()

        if len(mesh.vertices) == 0:
            raise ValueError("Mesh has no vertices after processing.")

        yield {"status": "processing", "progress": 15, "message": "Optimizing orientation..."}
        extents = mesh.extents
        longest_axis_index = np.argmax(extents)
        if longest_axis_index != 2:
            source_vec = np.eye(3)[longest_axis_index]
            transform = trimesh.geometry.align_vectors(source_vec, [0, 0, 1])
            mesh.apply_transform(transform)

        mesh.vertices -= mesh.center_mass

        bounds = mesh.bounds
        z_min, z_max = bounds[:, 2]
        z_height = z_max - z_min
        safe_buffer = min(z_height * 0.01, 0.05)
        z_levels = np.linspace(z_min + safe_buffer, z_max - safe_buffer, num_layers)

        all_profiles_3d = []
        yield {"status": "processing", "progress": 30, "message": "Starting slice generation..."}

        for i, z in enumerate(z_levels):
            section = mesh.section(plane_origin=[0, 0, z], plane_normal=[0, 0, 1])

            if section is None:
                logging.warning(f"Layer {i+1}/{num_layers}: No section found at z={z:.2f}.")
                layer_errors.append(f"Layer {i+1} (z={z:.2f}): No section.")
                continue

            try:
                # Use to_planar() to get 2D representation
                # Note: to_2D() is the newer recommended method in trimesh
                paths = section.to_planar()[0].discrete
                if not paths:
                    layer_errors.append(f"Layer {i+1}: No discrete paths.")
                    continue

                slice_points = np.vstack(paths)
                
                # Splprep requires more points than the degree of the spline (k=3 by default)
                if len(slice_points) < 5:
                    logging.warning(f"Layer {i+1}: Insufficient points ({len(slice_points)}).")
                    layer_errors.append(f"Layer {i+1}: Insufficient points.")
                    continue

                hull = ConvexHull(slice_points)
                hull_pts = slice_points[hull.vertices]

                # splprep needs unique points, so we add a tiny bit of noise if needed
                tck, u = splprep(
                    [hull_pts[:, 0], hull_pts[:, 1]], s=0, per=True)
                # Generate N points over the interval [0, 1), excluding the endpoint because it's a closed loop
                u_new = np.linspace(0, 1, num_points_per_layer, endpoint=False)
                x_new, y_new = splev(u_new, tck)

                pts_2d = np.column_stack((x_new, y_new))
                start_idx = np.argmax(pts_2d[:, 0])
                pts_2d = np.roll(pts_2d, -start_idx, axis=0)

                profile_3d = np.column_stack((pts_2d, np.full(len(pts_2d), z)))
                all_profiles_3d.append(profile_3d)

            except Exception as e:
                logging.warning(f"Layer {i+1}: Error: {e}")
                layer_errors.append(f"Layer {i+1}: {str(e)[:50]}")
                continue

            yield {"status": "processing", "progress": 30 + int((i/num_layers)*50), "message": f"Layer {i+1}/{num_layers}"}

        if not all_profiles_3d:
            raise ValueError("Could not generate any valid profiles. Result would be empty.")

        yield {"status": "processing", "progress": 90, "message": "Writing DXF..."}
        doc = ezdxf.new()
        msp = doc.modelspace()
        for poly in all_profiles_3d:
            msp.add_polyline3d(poly, close=True)

        doc.saveas(output_file)
        
        if os.path.exists(output_file):
            msg = "DXF created successfully."
            if layer_errors:
                msg += f" (with {len(layer_errors)} skipped layers)"
            yield {"status": "complete", "progress": 100, "message": msg}
        else:
            raise IOError(f"File not found after save: {output_file}")

    except Exception as e:
        logging.error(f"Job failed: {e}")
        yield {"status": "failed", "progress": 0, "message": str(e)}