assembly_handbook/ahb_cmd_render.py

import os
import tempfile

import FreeCADGui as Gui
import FreeCAD as App

class AHB_Render:
    def GetResources(self):
        return {"MenuText": "Render",
                "ToolTip": "Render to images",
                "Pixmap": ""
                }

    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)):
        doc = App.activeDocument()
        for obj in doc.Objects:
            if obj.TypeId == 'Part::Feature':
                if 'AssemblyHandbook_Stage' in obj.PropertiesList:
                    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.ShapeMaterial.AmbientColor = (0.0, 0.0, 0.0, 0.0)
                    obj.ViewObject.ShapeMaterial.EmissiveColor = background_color

    def set_render_outlines(self):
        doc = App.activeDocument()
        step = 8
        r = step
        g = step
        b = step
        for obj in doc.Objects:
            if obj.TypeId == 'Part::Feature':
                if 'AssemblyHandbook_Stage' in obj.PropertiesList:
                    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)

                    r = r + step
                    if r >= 256 - step:
                        r = step
                        g = g + step
                        if g >= 256 - step:
                            g = step
                            b = b + step
                            if b >= 256 - step:
                                b = step

    def reset_display(self):
        doc = App.activeDocument()
        for obj in doc.Objects:
            if obj.TypeId == 'Part::Feature':
                if 'AssemblyHandbook_Stage' in obj.PropertiesList:
                    obj.ViewObject.LineColor = (0.0, 0.0, 0.0, 0.0)
                    obj.ViewObject.DisplayMode = 'Flat Lines'
                    obj.ViewObject.ShapeMaterial.AmbientColor = (0.3, 0.3, 0.3, 0.0)
                    obj.ViewObject.ShapeMaterial.DiffuseColor = (1.0, 1.0, 1.0, 0.0)
                    obj.ViewObject.ShapeMaterial.SpecularColor = (0.5, 0.5, 0.5, 0.0)
                    obj.ViewObject.ShapeMaterial.EmissiveColor = (0.0, 0.0, 0.0, 0.0)

        workbench = Gui.getWorkbench("AssemblyHandbookWorkbench")
        workbench.context.onPartStageChanged(None)

    def Activated(self):
        resolution = (2000, 2000)
        filename = "/home/youen/dev_linux/vhelio-render/vhelio.png"

        temp_lines_file_name = tempfile.gettempdir() + "/ahb_temp_lines.png"
        temp_shapes_file_name = tempfile.gettempdir() + "/ahb_temp_shapes.png"

        self.set_render_outlines()
        Gui.ActiveDocument.ActiveView.saveImage(temp_shapes_file_name, resolution[0] * 2, resolution[1] * 2, "#ffffff")

        self.set_render_lines()
        Gui.ActiveDocument.ActiveView.saveImage(temp_lines_file_name, resolution[0], resolution[1], "#ffffff")

        self.reset_display()

        from PIL import Image, ImageFilter
        import time
        lines = Image.open(temp_lines_file_name)
        lines = lines.convert("L")
        shapes = Image.open(temp_shapes_file_name)

        startTime = 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
                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")
                partial_outlines = partial_outlines.point(lambda p: 255 if p == 255 else 0)
                if outlines is None:
                    outlines = partial_outlines
                else:
                    outlines.paste(partial_outlines, None, partial_outlines.point(lambda p: 255 - p))

        #outlines.save("/home/youen/dev_linux/vhelio-render/vhelio-outlines.png")

        #outlines = outlines.resize(lines.size, Image.BILINEAR)
        #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 = lines_fullres.resize(lines.size, Image.BILINEAR)

        endTime = time.perf_counter()
        print("Outline detection: " + str(round((endTime - startTime)*1000)/1000) + "s")
        lines.save(filename)

from ahb_command import AHB_CommandWrapper
AHB_CommandWrapper.addGuiCommand('AHB_render', AHB_Render())