Merge pull request #7 from RomainGuimbal/fixes

fixing cache, rebuild and hierarchy with single shape

Author: RomainGuimbal

__init__.py

@@ -151,7 +151,7 @@ class STEP_OT_ImportStepCADOperator(bpy.types.Operator, ImportHelper):
             ),
         ],
         name="Tree hierarchy",
-        default=str(int(HierarchyType.EMPTIES_TREE)),
+        default=str(int(HierarchyType.COLLECTION_TREE)),
         description="Organization style of objects",
     )
 
@@ -300,7 +300,6 @@ class STEP_OT_ClearCache(bpy.types.Operator):
     bl_description = "Clear STEP cache, enabling the reload of a file"
 
     def execute(self, context):
-        # utils.memorytrace_print()
         # global GLOBAL_FILE_CACHE
         # items = list(GLOBAL_FILE_CACHE.values())
         # for entry in items:
@@ -437,7 +436,6 @@ def execute(self, context):
                     other_obj.select_set(True)
 
         # Reload files if not in cache
-        reload_needed = False
         for o in selected_objects:
             if o["STEP_file"] not in GLOBAL_FILE_CACHE:
                 bpy.ops.object.occ_reload_step()

importer.py

@@ -56,7 +56,7 @@
     XCAFDoc_DocumentTool,
     XCAFDoc_ColorGen,
     XCAFDoc_ColorSurf,
-    # XCAFDoc_ColorCurv,
+    XCAFDoc_ColorCurv,
 )
 from OCP.XSControl import XSControl_WorkSession
 
@@ -251,23 +251,24 @@ def query_color(self, lab, overwrite=False):
             lab: shape label
         """
         # default color = pink
-        c = Quantity_Color(1.0, 0.0, 1.0, Quantity_TOC_RGB)
+        color = Quantity_Color(1.0, 0.0, 1.0, Quantity_TOC_RGB)
         iscolorset = False
         colortype = None
 
         shape = self.shape_tool.GetShape_s(lab)
-        print(shape.ShapeType())
+        # print(shape.ShapeType())
 
-        c_gen = self.color_tool.GetColor(shape, XCAFDoc_ColorGen, c)
-        c_surf = self.color_tool.GetColor(shape, XCAFDoc_ColorSurf, c)
-        c_curv = False  # self.color_tool.GetColor(shape, XCAFDoc_ColorCurv, c) # TODO uncomment once priority are working
+        # Overwrites each type
+        c_gen_exists = self.color_tool.GetColor(shape, XCAFDoc_ColorGen, color)
+        c_surf_exists = self.color_tool.GetColor(shape, XCAFDoc_ColorSurf, color)
+        c_curv_exists = False  # self.color_tool.GetColor(shape, XCAFDoc_ColorCurv, color) Supposed to be a fallback but overwrite't
 
-        if c_gen or c_surf or c_curv:
+        if c_gen_exists or c_surf_exists or c_curv_exists:
             iscolorset = True
             # Color priority (1/type) is the same as CAD assistant material tree display
-            colortype = c_gen * 1 + c_surf * 2 + c_curv * 3
+            colortype = c_gen_exists * 1 + c_surf_exists * 2 + c_curv_exists * 3
 
-        return c, colortype, iscolorset
+        return color, colortype, iscolorset
 
     def print_all_colors(self):
         tcol = Quantity_Color(1.0, 0.0, 1.0, Quantity_TOC_RGB)
@@ -424,27 +425,27 @@ def transfer_with_units(self, filename):
 
         self.doc = doc
 
-    def transfer_simple(self, fname):
-        # see stepanalyzer.py for license details
-        print("Init simple transfer")
-
-        # Create the application, empty document and shape_tool
-        doc = TDocStd_Document(TCollection_ExtendedString("STEP"))
-        app = XCAFApp_Application.GetApplication()
-        app.NewDocument("MDTV-XCAF", doc)
-
-        # Read file and return populated doc
-        step_reader = STEPCAFControl_Reader()
-        step_reader.SetColorMode(True)
-        step_reader.SetLayerMode(True)
-        step_reader.SetNameMode(True)
-        step_reader.SetMatMode(True)
-        status = step_reader.ReadFile(fname)
-        if status == IFSelect_RetDone:
-            step_reader.Transfer(doc)
-        self.scale = 0.001
-
-        self.doc = doc
+    # def transfer_simple(self, fname):
+    #     # see stepanalyzer.py for license details
+    #     print("Init simple transfer")
+
+    #     # Create the application, empty document and shape_tool
+    #     doc = TDocStd_Document(TCollection_ExtendedString("STEP"))
+    #     app = XCAFApp_Application.GetApplication()
+    #     app.NewDocument("MDTV-XCAF", doc)
+
+    #     # Read file and return populated doc
+    #     step_reader = STEPCAFControl_Reader()
+    #     step_reader.SetColorMode(True)
+    #     step_reader.SetLayerMode(True)
+    #     step_reader.SetNameMode(True)
+    #     step_reader.SetMatMode(True)
+    #     status = step_reader.ReadFile(fname)
+    #     if status == IFSelect_RetDone:
+    #         step_reader.Transfer(doc)
+    #     self.scale = 0.001
+
+    #     self.doc = doc
 
     def init_reader(self, filename):
         if not os.path.isfile(filename):

object_generator.py

@@ -20,6 +20,7 @@
 from .importer import ShapeTreeNode
 
 GLOBAL_FILE_CACHE = {}
+ROOT_PARENT = -1
 
 
 class HierarchyType(IntEnum):
@@ -262,6 +263,10 @@ def build_mesh(
     trimesh.filter_zero_area()
     trimesh.filter_same_face()
     trimesh.fill_empty_color()
+
+    # Temporary fix. Original addon was apparently smarter at replacing only the rebuilt faces
+    obj.data.clear_geometry()
+
     trimesh.add_to_mesh(obj.data)
     obj.data.update()
 
@@ -310,7 +315,9 @@ def node_parse(
         # link objects to tree
         if len(hierarchy_collections.items()) > 0:
             for obj in created_objs:
-                hierarchy_collections[obj["STEP_parent"]].objects.link(obj)
+                parent_col_uuid = obj.get("STEP_parent", ROOT_PARENT)
+                parent_col = hierarchy_collections.get(parent_col_uuid, hierarchy_collections[-1])
+                parent_col.objects.link(obj)
                 global_t = tree.nodes[obj["STEP_tree_location"]].global_transform
                 set_obj_matrix_world(obj, global_t)
 
@@ -389,7 +396,7 @@ def build_collection_instances(instanced_objects, hierarchy_collections, scale):
         component_col.objects.link(source_obj)
 
         # instance original
-        empty = bpy.data.objects.new(source_obj.name + "_instance", None)
+        empty = bpy.data.objects.new(source_obj.name, None)
         empty.instance_type = "COLLECTION"
         empty.instance_collection = component_col
         empty.matrix_world = trsf
@@ -400,13 +407,13 @@ def build_collection_instances(instanced_objects, hierarchy_collections, scale):
         # instance copies
         for shape_name, local_t, global_t, parent_uuid in instances_info:
             # create instance
-            empty = bpy.data.objects.new(shape_name + "_instance", None)
+            empty = bpy.data.objects.new(source_obj.name + "_instance", None)
             empty.instance_type = "COLLECTION"
             empty.instance_collection = component_col
             empty.empty_display_size = scale
             scale_translation(global_t, scale)
             set_obj_matrix_world(empty, global_t)
-            parent_col = hierarchy_collections[parent_uuid]
+            parent_col = hierarchy_collections.get(parent_uuid, hierarchy_collections[-1])
             parent_col.objects.link(empty)
 
 
@@ -422,8 +429,9 @@ def load_step(
 ):
     # Find file
     from . import importer
-
     filename = "".join(ntpath.basename(filepath).split(".")[:-1])
+    
+    # Try retrieve step_reader from cache
     if filepath not in GLOBAL_FILE_CACHE:
         try:
             step_reader = importer.ReadSTEP(filepath)
@@ -458,7 +466,7 @@ def load_step(
     all_shapes = tree.get_shapes()
     total = len(all_shapes)
 
-    # Build shapes
+    # Build shapes objects
     wm.progress_begin(0, total)
     for i, (shp, node_index) in enumerate(all_shapes):
         parent_uuid, self_uuid, tag, obj_name, _, local_t, global_t = tree.nodes[
@@ -488,12 +496,14 @@ def load_step(
                 source_obj = created_names[shape_name]
                 if htypes == HierarchyType.COLLECTION_INSTANCES:
                     if source_obj in instanced_objects:
+                        parent_for_instance = parent_uuid if parent_uuid != 0 else ROOT_PARENT
                         instanced_objects[source_obj].append(
-                            (shape_name, local_t, global_t, parent_uuid)
+                            (shape_name, local_t, global_t, parent_for_instance)
                         )
                     else:
+                        parent_for_instance = parent_uuid if parent_uuid != 0 else ROOT_PARENT
                         instanced_objects[source_obj] = [
-                            (shape_name, local_t, global_t, parent_uuid)
+                            (shape_name, local_t, global_t, parent_for_instance)
                         ]
                 else:
                     obj = source_obj.copy()
@@ -526,7 +536,8 @@ def load_step(
         if obj:
             # assign property to obj
             obj["STEP_tag"] = tag
-            obj["STEP_parent"] = parent_uuid
+            obj_parent = parent_uuid if parent_uuid != 0 else ROOT_PARENT
+            obj["STEP_parent"] = obj_parent
             obj["STEP_uuid"] = self_uuid
             obj["STEP_file"] = filepath
             obj["STEP_name"] = obj_name

trimesh.py

@@ -1,6 +1,7 @@
 import numpy as np
 import bpy
 from dataclasses import dataclass
+from mathutils import Vector
 
 
 def make_tri_hash(f):
@@ -210,16 +211,16 @@ def set_material_name(self, name):
         self.tris.material_name = [name] * len(self.tris)
 
     def fill_empty_color(self):
-        "Fill color==None with undef_color (currently pink)"
-        # TODO : make undef black and add attributes to tell if face is colored so the shader can know
-        undef_color = (1.0, 0.0, 1.0)  # pink
+        """
+        Fill tris empty colors
+        """
         if len(self.tris) == 0:
-            # Empty mesh
             return
+        undef_color = (0.8, 0.8, 0.8)
         self.tris.colors = [c if c else undef_color for c in self.tris.colors]
 
     def add_to_mesh(self, mesh: bpy.types.Mesh):
-        mesh.from_pydata(self.verts, [], self.tris.indices)
+        mesh.from_pydata([Vector(v) for v in self.verts], [], self.tris.indices)
 
     def get_loop_colors(self):
         "Return colors in triangle loop creation order"