Upload face_labeler.py

face_labeler.py

@@ -1,5 +1,7 @@
-import sys
+import sys
 import os
+import argparse
+import pathlib
 from dataclasses import dataclass
 from typing import List, Optional, Tuple
 
@@ -22,31 +24,83 @@ from OCC.Core.TopAbs import TopAbs_FACE
 from OCC.Core.TopoDS import topods
 from OCC.Core.Quantity import Quantity_Color, Quantity_TOC_RGB
 from OCC.Display.qtDisplay import qtViewer3d
+try:
+    from OCC.Core.Aspect import Aspect_TOL_SOLID
+    from OCC.Core.Prs3d import Prs3d_LineAspect
+except Exception:  # pragma: no cover - optional OCC build
+    Aspect_TOL_SOLID = None
+    Prs3d_LineAspect = None
+try:
+    from OCC.Core.Graphic3d import Graphic3d_NOM_MATTE, Graphic3d_NOM_NEON
+except Exception:  # pragma: no cover - optional OCC build
+    Graphic3d_NOM_MATTE = None
+    Graphic3d_NOM_NEON = None
+try:
+    from OCC.Core.Graphic3d import Graphic3d_TOSM_UNLIT
+except Exception:  # pragma: no cover - optional OCC build
+    Graphic3d_TOSM_UNLIT = None
 
 CLASS_NAMES = [
-    "ExtrudeSide",
-    "ExtrudeEnd",
-    "CutSide",
-    "CutEnd",
-    "Fillet",
-    "Chamfer",
-    "RevolveSide",
-    "RevolveEnd",
+    "SOL",
+    "EOS",
+    "rect_slot",
+    "tri_slot",
+    "cir_slot",
+    "rect_psg",
+    "tri_psg",
+    "hexa_psg",
+    "hole",
+    "rect_step",
+    "tside_step",
+    "slant_step",
+    "rect_b_step",
+    "tri_step",
+    "cir_step",
+    "rect_b_slot",
+    "cir_b_slot",
+    "u_b_slot",
+    "rect_pkt",
+    "key_pkt",
+    "tri_pkt",
+    "hexa_pkt",
+    "o_ring",
+    "b_hole",
+    "chamfer",
+    "fillet",
 ]
 
 CLASS_COLORS_HEX = [
-    "#e41a1c",  # red
-    "#377eb8",  # blue
-    "#4daf4a",  # green
-    "#984ea3",  # purple
-    "#0e2579",  # teal-blue (fillet)
-    "#a65628",  # brown
-    "#f781bf",  # pink
-    "#00a7a7",  # teal
+    "#1f77b4",  # blue
+    "#ff7f0e",  # orange
+    "#2ca02c",  # green
+    "#d62728",  # red
+    "#9467bd",  # purple
+    "#8c564b",  # brown
+    "#e377c2",  # pink
+    "#7f7f7f",  # gray
+    "#bcbd22",  # olive
+    "#17becf",  # cyan
+    "#393b79",  # dark blue
+    "#637939",  # dark green
+    "#8c6d31",  # dark mustard
+    "#843c39",  # dark red
+    "#7b4173",  # dark purple
+    "#3182bd",  # blue alt
+    "#e6550d",  # orange alt
+    "#31a354",  # green alt
+    "#756bb1",  # purple alt
+    "#636363",  # dark gray
+    "#6baed6",  # light blue
+    "#fd8d3c",  # light orange
+    "#74c476",  # light green
+    "#9e9ac8",  # light purple
+    "#a1d99b",  # pale green
+    "#fdd0a2",  # pale orange
 ]
 
 UNLABELED_COLOR_HEX = "#d0d0d0"
 HIGHLIGHT_COLOR_HEX = "#FFD400"  # fixed for current selection
+EDGE_COLOR_HEX = "#2b2b2b"
 
 
 def hex_to_rgb01(color_hex: str) -> Tuple[float, float, float]:
@@ -74,9 +128,9 @@ class FaceItem:
 
 
 class FaceLabeler(QtWidgets.QMainWindow):
-    def __init__(self, step_path: Optional[str] = None):
+    def __init__(self, step_path: Optional[str] = None, output_dir: Optional[str] = None):
         super().__init__()
-        self.setWindowTitle("BRep Face Labeler")
+        self.setWindowTitle("BRepMFR Face Labeler")
         self.resize(1600, 1000)
         self.setMinimumSize(1200, 800)
 
@@ -89,10 +143,13 @@ class FaceLabeler(QtWidgets.QMainWindow):
         self.labels: List[Optional[int]] = []
         self.current_index: Optional[int] = None
         self.highlight_enabled = True
+        self.step_path: Optional[str] = None
+        self.output_dir: Optional[str] = output_dir
 
         self.highlight_color = rgb01_to_quantity(hex_to_rgb01(HIGHLIGHT_COLOR_HEX))
 
         self._build_ui()
+        self.update_step_label()
 
         if step_path:
             self.load_step(step_path)
@@ -110,6 +167,7 @@ class FaceLabeler(QtWidgets.QMainWindow):
             self.display.Context.SetAutomaticHilight(False)
         except Exception:
             pass
+        self._configure_viewer_visuals()
         root_layout.addWidget(self.viewer, 1)
 
         panel = QtWidgets.QWidget(central)
@@ -141,6 +199,10 @@ class FaceLabeler(QtWidgets.QMainWindow):
         nav_layout.addWidget(self.btn_next)
         panel_layout.addLayout(nav_layout)
 
+        self.step_label = QtWidgets.QLabel("STEP: (none)")
+        self.step_label.setWordWrap(True)
+        panel_layout.addWidget(self.step_label)
+
         self.info_label = QtWidgets.QLabel("No STEP loaded")
         self.info_label.setWordWrap(True)
         panel_layout.addWidget(self.info_label)
@@ -166,6 +228,15 @@ class FaceLabeler(QtWidgets.QMainWindow):
 
         self.setCentralWidget(central)
 
+    def update_step_label(self) -> None:
+        if self.step_path:
+            name = os.path.basename(self.step_path)
+            self.step_label.setText(f"STEP: {name}")
+            self.setWindowTitle(f"BRepMFR Face Labeler - {name}")
+        else:
+            self.step_label.setText("STEP: (none)")
+            self.setWindowTitle("BRepMFR Face Labeler")
+
     def keyPressEvent(self, event) -> None:  # pragma: no cover - UI only
         if event.key() in (QtCore.Qt.Key_Right, QtCore.Qt.Key_D):
             self.on_next()
@@ -193,6 +264,17 @@ class FaceLabeler(QtWidgets.QMainWindow):
                 "Unlabeled faces remain. Label all faces before exporting.",
             )
             return
+        if self.output_dir:
+            if not self.step_path:
+                QtWidgets.QMessageBox.warning(self, "Export", "No STEP loaded.")
+                return
+            output_dir = pathlib.Path(self.output_dir)
+            output_dir.mkdir(parents=True, exist_ok=True)
+            filename = pathlib.Path(self.step_path).with_suffix(".seg").name
+            path = output_dir / filename
+            self.save_seg(str(path))
+            return
+
         path, _ = QtWidgets.QFileDialog.getSaveFileName(
             self, "Export .seg", "", "SEG Files (*.seg)"
         )
@@ -267,6 +349,8 @@ class FaceLabeler(QtWidgets.QMainWindow):
             return
         reader.TransferRoots()
         shape = reader.OneShape()
+        self.step_path = path
+        self.update_step_label()
 
         self.display.EraseAll()
         self.face_items.clear()
@@ -280,10 +364,15 @@ class FaceLabeler(QtWidgets.QMainWindow):
             ais = self.display.DisplayShape(face, update=False, color=self.unlabeled_color)
             if isinstance(ais, list):
                 ais = ais[0]
+            self._apply_face_material(ais)
             try:
                 self.display.Context.SetDisplayMode(ais, 1, False)
             except Exception:
                 pass
+            try:
+                self.display.Context.Redisplay(ais, False)
+            except Exception:
+                pass
             self.face_items.append(FaceItem(face=face, ais=ais))
             self.labels.append(None)
             explorer.Next()
@@ -297,6 +386,7 @@ class FaceLabeler(QtWidgets.QMainWindow):
 
         self.display.FitAll()
         self.set_current_index(0)
+        self.display.Repaint()
 
     def save_seg(self, path: str) -> None:
         with open(path, "w", encoding="utf-8") as handle:
@@ -359,17 +449,73 @@ class FaceLabeler(QtWidgets.QMainWindow):
             ais.SetColor(color)
         except Exception:
             self.display.Context.SetColor(ais, color, False)
+        self._apply_face_material(ais)
         self.display.Context.Redisplay(ais, False)
 
+    def _apply_face_material(self, ais) -> None:
+        # Prefer emissive material to keep colors stable regardless of lighting.
+        applied = False
+        if Graphic3d_NOM_NEON is not None:
+            try:
+                ais.SetMaterial(Graphic3d_NOM_NEON)
+                applied = True
+            except Exception:
+                pass
+        if not applied and Graphic3d_NOM_MATTE is not None:
+            try:
+                ais.SetMaterial(Graphic3d_NOM_MATTE)
+            except Exception:
+                pass
+        self._apply_face_edges(ais)
+
+    def _configure_viewer_visuals(self) -> None:
+        # Try unlit shading to avoid view-dependent brightening.
+        if Graphic3d_TOSM_UNLIT is None:
+            return
+        try:
+            self.display.View.SetShadingModel(Graphic3d_TOSM_UNLIT)
+        except Exception:
+            pass
+
+    def _apply_face_edges(self, ais) -> None:
+        if Prs3d_LineAspect is None or Aspect_TOL_SOLID is None:
+            return
+        try:
+            drawer = ais.Attributes()
+            drawer.SetFaceBoundaryDraw(True)
+            line_aspect = Prs3d_LineAspect(
+                rgb01_to_quantity(hex_to_rgb01(EDGE_COLOR_HEX)),
+                Aspect_TOL_SOLID,
+                1.0,
+            )
+            drawer.SetFaceBoundaryAspect(line_aspect)
+        except Exception:
+            pass
+
 
 def main() -> int:
+    parser = argparse.ArgumentParser(description="BRepMFR Face Labeler")
+    parser.add_argument("step_path", nargs="?", help="Optional STEP file to open")
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default=None,
+        help="Output directory for .seg exports (auto-save with input filename)",
+    )
+    parser.add_argument(
+        "--output_folder",
+        type=str,
+        default=None,
+        help="Alias of --output_dir",
+    )
+    args = parser.parse_args()
+    if args.output_dir and args.output_folder and args.output_dir != args.output_folder:
+        raise SystemExit("--output_dir and --output_folder must match when both are provided")
+
     app = QtWidgets.QApplication(sys.argv)
-    step_path = None
-    if len(sys.argv) > 1:
-        candidate = sys.argv[1]
-        if os.path.exists(candidate):
-            step_path = candidate
-    window = FaceLabeler(step_path=step_path)
+    step_path = args.step_path if args.step_path and os.path.exists(args.step_path) else None
+    output_dir = args.output_folder or args.output_dir
+    window = FaceLabeler(step_path=step_path, output_dir=output_dir)
     window.show()
     return app.exec_()