server/geometry.py

import logging
import time
from collections import Counter
from typing import Any, Dict, List, Optional, Tuple

import numpy as np

logger = logging.getLogger(__name__)


def extract_properties(shape) -> Dict[str, Any]:
    t0 = time.time()
    try:
        props = {}

        bb = shape.BoundingBox()
        dims = [bb.xlen, bb.ylen, bb.zlen]
        sorted_axes = sorted(
            zip(["X", "Y", "Z"], dims), key=lambda x: x[1], reverse=True
        )
        longest_axis = sorted_axes[0][0]

        props["is_valid"] = shape.isValid()

        volume = shape.Volume()
        surface_area = shape.Area()
        props["volume_mm3"] = round(volume, 4)
        props["surface_area_mm2"] = round(surface_area, 4)

        props["bbox_x_mm"] = round(bb.xlen, 4)
        props["bbox_y_mm"] = round(bb.ylen, 4)
        props["bbox_z_mm"] = round(bb.zlen, 4)
        props["bbox_longest_axis"] = longest_axis

        max_dim = max(dims)
        props["bbox_ratio_yx"] = round(sorted_axes[1][1] / max_dim, 4) if max_dim > 0 else 0.0
        props["bbox_ratio_zx"] = round(sorted_axes[2][1] / max_dim, 4) if max_dim > 0 else 0.0

        faces = shape.Faces()
        face_types = [f.geomType() for f in faces]
        type_counts = dict(Counter(face_types))
        total_faces = len(face_types)

        props["face_count"] = total_faces
        props["face_type_counts"] = type_counts
        props["dominant_face_type"] = max(type_counts, key=type_counts.get) if type_counts else "UNKNOWN"
        props["face_type_distribution"] = {
            k: round(v / total_faces, 4) for k, v in type_counts.items()
        } if total_faces > 0 else {}

        props["edge_count"] = len(shape.Edges())
        props["vertex_count"] = len(shape.Vertices())

        V = props["vertex_count"]
        E = props["edge_count"]
        F = props["face_count"]
        props["euler_characteristic"] = V - E + F

        props["is_watertight"] = _check_watertight(shape)

        xy_sym, xz_sym, yz_sym = _check_symmetry(shape, bb)
        props["has_xy_symmetry"] = xy_sym
        props["has_xz_symmetry"] = xz_sym
        props["has_yz_symmetry"] = yz_sym

        props["shape_class"] = _classify_shape(props)

        elapsed = time.time() - t0
        logger.info(f"extract_properties took {elapsed:.3f}s")
        return props

    except Exception as e:
        elapsed = time.time() - t0
        logger.error(f"extract_properties failed after {elapsed:.3f}s: {e}")
        raise


def _check_watertight(shape) -> bool:
    try:
        shells = shape.Shells()
        if not shells:
            return False
        for shell in shells:
            if not shell.Closed():
                return False
        return True
    except Exception:
        return False


def _check_symmetry(shape, bb, n_sample: int = 200, tol_ratio: float = 0.05) -> Tuple[bool, bool, bool]:
    try:
        from OCP.BRepMesh import BRepMesh_IncrementalMesh
        from OCP.BRep import BRep_Tool
        from OCP.TopExp import TopExp_Explorer
        from OCP.TopAbs import TopAbs_FACE
        from OCP.TopLoc import TopLoc_Location
        from OCP.TopoDS import TopoDS

        mesh = BRepMesh_IncrementalMesh(shape.wrapped, 0.5, False, 0.5, True)
        mesh.Perform()

        points = []
        explorer = TopExp_Explorer(shape.wrapped, TopAbs_FACE)
        while explorer.More():
            face = TopoDS.Face_s(explorer.Current())
            loc = TopLoc_Location()
            tri = BRep_Tool.Triangulation_s(face, loc)
            if tri is not None:
                for i in range(1, tri.NbNodes() + 1):
                    p = tri.Node(i)
                    trsf = loc.Transformation()
                    p.Transform(trsf)
                    points.append([p.X(), p.Y(), p.Z()])
            explorer.Next()

        if len(points) < 10:
            return False, False, False

        pts = np.array(points)
        if len(pts) > n_sample:
            idx = np.random.choice(len(pts), n_sample, replace=False)
            pts = pts[idx]

        diag = np.sqrt(bb.xlen**2 + bb.ylen**2 + bb.zlen**2)
        tol = diag * tol_ratio

        from scipy.spatial import cKDTree

        tree = cKDTree(pts)

        def check_plane_symmetry(pts_arr, axis_idx):
            reflected = pts_arr.copy()
            reflected[:, axis_idx] = -reflected[:, axis_idx]
            dists, _ = tree.query(reflected)
            return float(np.mean(dists)) < tol

        cx = (bb.xmin + bb.xmax) / 2
        cy = (bb.ymin + bb.ymax) / 2
        cz = (bb.zmin + bb.zmax) / 2
        centered = pts - np.array([cx, cy, cz])
        tree_c = cKDTree(centered)

        def check_sym(axis_idx):
            reflected = centered.copy()
            reflected[:, axis_idx] = -reflected[:, axis_idx]
            dists, _ = tree_c.query(reflected)
            return float(np.mean(dists)) < tol

        has_yz = check_sym(0)
        has_xz = check_sym(1)
        has_xy = check_sym(2)

        return has_xy, has_xz, has_yz

    except Exception as e:
        logger.warning(f"Symmetry check failed: {e}")
        return False, False, False


def _classify_shape(props: Dict[str, Any]) -> str:
    ratio_yx = props.get("bbox_ratio_yx", 0)
    ratio_zx = props.get("bbox_ratio_zx", 0)
    dominant = props.get("dominant_face_type", "")
    euler = props.get("euler_characteristic", 2)
    is_valid = props.get("is_valid", False)

    if not is_valid or props.get("volume_mm3", 0) <= 0:
        return "DEGENERATE"

    is_flat = ratio_zx < 0.25
    is_cubic = ratio_yx > 0.7 and ratio_zx > 0.7
    is_round = dominant in ("CYLINDER", "CONE", "SPHERE", "TORUS")
    is_tall = ratio_zx > 0.5 and not is_cubic

    if is_round:
        if is_flat:
            return "FLAT_ROUND_SOLID"
        if euler != 2:
            return "HOLLOW_ROUND_SOLID"
        if is_tall:
            return "TALL_ROUND_SOLID"
        return "FLAT_ROUND_SOLID"

    if is_flat:
        if euler != 2:
            return "FLAT_SOLID_WITH_HOLES"
        return "FLAT_SOLID"

    if is_cubic:
        if euler != 2:
            return "CUBIC_SOLID_WITH_HOLES"
        return "CUBIC_SOLID"

    face_count = props.get("face_count", 0)

    if face_count > 8 and ratio_yx < 0.5:
        return "L_SHAPED_SOLID"

    if face_count > 8 and ratio_yx > 0.5:
        return "T_SHAPED_SOLID"

    if face_count > 10:
        return "STEPPED_SOLID"

    return "COMPLEX_SOLID"