possibility to specify different rendering colors for previous stages and current stage

ahb_cmd_render.py

@@ -1,5 +1,6 @@
 import os
 import tempfile
+from typing import Optional
 
 import FreeCADGui as Gui
 import FreeCAD as App
@@ -14,20 +15,21 @@ class AHB_Render:
     def IsActive(self):
         return True
 
-    def set_render_lines(self, line_color = (0.0,0.0,0.0,0.0), background_color = (1.0,1.0,1.0,0.0)):
+    def set_render_lines(self, line_color = (0.0,0.0,0.0,0.0), background_color = (1.0,1.0,1.0,0.0), mask_stages_below: Optional[int] = None, mask_color=(1.0,1.0,1.0)):
         doc = App.activeDocument()
         for obj in doc.Objects:
             if obj.TypeId == 'Part::Feature':
                 if 'AssemblyHandbook_Stage' in obj.PropertiesList:
+                    masked = mask_stages_below is not None and obj.AssemblyHandbook_Stage < mask_stages_below
                     if 'AssemblyHandbook_RenderLines' in obj.PropertiesList and not obj.AssemblyHandbook_RenderLines:
                         obj.ViewObject.LineColor = background_color
                     else:
                         obj.ViewObject.LineColor = line_color
-                    obj.ViewObject.DisplayMode = 'Flat Lines'
+                    obj.ViewObject.DisplayMode = 'Flat Lines' if not masked else 'Shaded'
                     obj.ViewObject.ShapeMaterial.AmbientColor = (0.0, 0.0, 0.0, 0.0)
-                    obj.ViewObject.ShapeMaterial.EmissiveColor = background_color
+                    obj.ViewObject.ShapeMaterial.EmissiveColor = background_color if not masked else mask_color
 
-    def set_render_outlines(self):
+    def set_render_outlines(self, mask_stages_below: Optional[int] = None):
         doc = App.activeDocument()
         step = 8
         r = step
@@ -36,11 +38,12 @@ class AHB_Render:
         for obj in doc.Objects:
             if obj.TypeId == 'Part::Feature':
                 if 'AssemblyHandbook_Stage' in obj.PropertiesList:
+                    masked = mask_stages_below is not None and obj.AssemblyHandbook_Stage < mask_stages_below
                     obj.ViewObject.DisplayMode = 'Shaded'
                     obj.ViewObject.ShapeMaterial.AmbientColor = (0.0, 0.0, 0.0, 0.0)
                     obj.ViewObject.ShapeMaterial.DiffuseColor = (0.0, 0.0, 0.0, 0.0)
                     obj.ViewObject.ShapeMaterial.SpecularColor = (0.0, 0.0, 0.0, 0.0)
-                    obj.ViewObject.ShapeMaterial.EmissiveColor = (r/255.0, g/255.0, b/255.0, 0.0)
+                    obj.ViewObject.ShapeMaterial.EmissiveColor = (r/255.0, g/255.0, b/255.0, 0.0) if not masked else (1.0,1.0,1.0,0.0)
 
                     r = r + step
                     if r >= 256 - step:
@@ -67,35 +70,38 @@ class AHB_Render:
         workbench = Gui.getWorkbench("AssemblyHandbookWorkbench")
         workbench.context.onPartStageChanged(None)
 
-    def Activated(self):
+    def render(self, resolution, filename: str, line_color=(0.0, 0.0, 0.0, 0.0), fill_color=(1.0, 1.0, 1.0, 0.0), mask_stages_below: Optional[int] = None):
-        resolution = (2000, 2000)
+        import time
-        filename = "/home/youen/dev_linux/vhelio-render/vhelio.png"
+        from PIL import Image, ImageFilter
+
+        render_start_time = time.perf_counter()
 
         temp_lines_file_name = tempfile.gettempdir() + "/ahb_temp_lines.png"
         temp_shapes_file_name = tempfile.gettempdir() + "/ahb_temp_shapes.png"
 
-        self.set_render_outlines()
+        self.set_render_outlines(mask_stages_below=mask_stages_below)
         Gui.ActiveDocument.ActiveView.saveImage(temp_shapes_file_name, resolution[0] * 2, resolution[1] * 2, "#ffffff")
 
-        self.set_render_lines()
+        # render lines in black, background in red, fill shapes in green
-        Gui.ActiveDocument.ActiveView.saveImage(temp_lines_file_name, resolution[0], resolution[1], "#ffffff")
+        # the green band contains the lines images, the red band contains the inverted alpha layer
+        self.set_render_lines((0.0,0.0,0.0), (0.0,1.0,0.0), mask_stages_below=mask_stages_below, mask_color=(1.0,0.0,0.0))
+        Gui.ActiveDocument.ActiveView.saveImage(temp_lines_file_name, resolution[0], resolution[1], "#ff0000")
 
         self.reset_display()
 
-        from PIL import Image, ImageFilter
+        lines_bands = Image.open(temp_lines_file_name).split()
-        import time
+        lines = lines_bands[1]
-        lines = Image.open(temp_lines_file_name)
+        alpha_band = lines_bands[0].point(lambda p: 255 - p)
-        lines = lines.convert("L")
         shapes = Image.open(temp_shapes_file_name)
 
-        startTime = time.perf_counter()
+        outlines_start_time = time.perf_counter()
 
         outlines = None
         for x in range(0, 3):
             for y in range(0, 3):
                 if x == 1 and y == 1: continue
                 kernel = [0, 0, 0, 0, 1, 0, 0, 0, 0]
-                kernel[y*3 + x] = -1
+                kernel[y * 3 + x] = -1
                 partial_outlines = shapes.filter(ImageFilter.Kernel((3, 3), kernel, 1, 127))
                 partial_outlines = partial_outlines.point(lambda p: 255 if p == 127 else 0)
                 partial_outlines = partial_outlines.convert("L")
@@ -103,20 +109,84 @@ class AHB_Render:
                 if outlines is None:
                     outlines = partial_outlines
                 else:
-                    outlines.paste(partial_outlines, None, partial_outlines.point(lambda p: 255 - p))
+                    outlines.paste(partial_outlines, None, partial_outlines.point(lambda p: 0 if p == 255 else 255))
 
-        #outlines.save("/home/youen/dev_linux/vhelio-render/vhelio-outlines.png")
+        # outlines.save("/home/youen/dev_linux/vhelio-render/vhelio-outlines.png")
 
-        #outlines = outlines.resize(lines.size, Image.BILINEAR)
+        # outlines = outlines.resize(lines.size, Image.BILINEAR)
-        #lines.paste(outlines, None, outlines.point(lambda p: 255 - p))
+        # lines.paste(outlines, None, outlines.point(lambda p: 255 - p))
 
         lines_fullres = lines.resize(outlines.size, Image.NEAREST)
-        lines_fullres.paste(outlines, None, outlines.point(lambda p: 255 - p))
+        lines_fullres.paste(outlines, None, outlines.point(lambda p: 0 if p == 255 else 255))
         lines = lines_fullres.resize(lines.size, Image.BILINEAR)
 
-        endTime = time.perf_counter()
+        alpha_band_fullres = alpha_band.resize(outlines.size, Image.NEAREST)
-        print("Outline detection: " + str(round((endTime - startTime)*1000)/1000) + "s")
+        alpha_band_fullres.paste(outlines.point(lambda p: 255), None, outlines.point(lambda p: 0 if p == 255 else 255))
-        lines.save(filename)
+        alpha_band = alpha_band_fullres.resize(lines.size, Image.BILINEAR)
+
+        outlines_end_time = time.perf_counter()
+
+        # colorize
+
+        result = Image.merge("RGBA", [
+            lines.point(lambda p: int(fill_color[0] * p + line_color[0]*(255.0-p))),
+            lines.point(lambda p: int(fill_color[1] * p + line_color[1] * (255.0 - p))),
+            lines.point(lambda p: int(fill_color[2] * p + line_color[2] * (255.0 - p))),
+            alpha_band
+        ])
+
+        result.save(filename)
+
+        print("Rendered " + filename + " in " + str(
+            round((outlines_end_time - render_start_time) * 1000) / 1000) + "s (outlines detection in " + str(
+            round((outlines_end_time - outlines_start_time) * 1000) / 1000) + "s)")
+
+    def Activated(self):
+        import shutil
+        from PIL import Image, ImageFilter
+
+        workbench = Gui.getWorkbench("AssemblyHandbookWorkbench")
+        doc = App.activeDocument()
+        doc_file_name: str = doc.FileName
+        if doc_file_name is None:
+            raise BaseException("You must save your FreeCAD document before rendering images")
+
+        filename = os.path.splitext(doc_file_name)[0] + ".png"
+        dir = os.path.dirname(filename)
+
+        workbench.context.setAllStagesVisible(True)
+        self.render((2000, 2000), filename)
+
+        render_stages = True
+        if render_stages:
+            shutil.rmtree(dir + "/stages", ignore_errors=True)
+            os.makedirs(dir + "/stages", exist_ok=True)
+            all_stages = workbench.context.getAllStages()
+            prev_stage_id: Optional[int] = None
+            for stage_id in all_stages:
+                stage_name = str(stage_id)
+                while len(stage_name) < 6:
+                    stage_name = "0" + stage_name
+                resolution = (2000, 2000)
+                if prev_stage_id is not None:
+                    workbench.context.setActiveStage(prev_stage_id)
+                    workbench.context.setAllStagesVisible(False)
+                    self.render(resolution, dir + "/stages/" + stage_name + "-bg.png", (0.5, 0.5, 0.5), (0.75,0.75,0.75))
+                workbench.context.setActiveStage(stage_id)
+                workbench.context.setAllStagesVisible(False)
+                self.render(resolution, dir + "/stages/" + stage_name + ".png", mask_stages_below=stage_id)
+                if prev_stage_id is not None:
+                    # merge previous stages background with new stage
+                    bg = Image.open(dir + "/stages/" + stage_name + "-bg.png")
+                    os.remove(dir + "/stages/" + stage_name + "-bg.png")
+                    new = Image.open(dir + "/stages/" + stage_name + ".png")
+                    bg.paste(new, None, new.getchannel('A'))
+                    bg.save(dir + "/stages/" + stage_name + ".png")
+                if prev_stage_id is not None:
+                    pass
+                prev_stage_id = stage_id
+
+        workbench.context.setAllStagesVisible(True)
 
 from ahb_command import AHB_CommandWrapper
 AHB_CommandWrapper.addGuiCommand('AHB_render', AHB_Render())