import os import tempfile from typing import Optional 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), 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' if not masked else '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 = background_color if not masked else mask_color def set_render_outlines(self, mask_stages_below: Optional[int] = None): 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: 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) if not masked else (1.0,1.0,1.0,0.0) obj.ViewObject.ShapeMaterial.EmissiveColor = (r/255.0, g/255.0, b/255.0, 0.0) if masked: obj.ViewObject.Visibility = False 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 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): import time 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" #temp_lines_file_name = filename + "-lines.png" #temp_shapes_file_name = filename + "-shapes.png" # render lines in black, background in red, fill shapes in green # 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, 1.0)) Gui.ActiveDocument.ActiveView.saveImage(temp_lines_file_name, resolution[0], resolution[1], "#ff0000") 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.reset_display() lines_bands = Image.open(temp_lines_file_name).split() lines = lines_bands[1] alpha_band = lines_bands[0].point(lambda p: 255 - p) shapes = Image.open(temp_shapes_file_name) 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 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: 0 if p == 255 else 255)) # erase masked outlines outlines.paste(outlines.point(lambda p: 255), None, lines_bands[2].resize(outlines.size).point(lambda p: 255 if p == 255 else 0)) # 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 if p == 0 else 0)) lines = lines_fullres.resize(lines.size, Image.BILINEAR) alpha_band_fullres = alpha_band.resize(outlines.size, Image.NEAREST) alpha_band_fullres.paste(outlines.point(lambda p: 255), None, outlines.point(lambda p: 255 if p == 0 else 0)) 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 import math from pivy import coin render_main = False render_stages = False render_parts = True Gui.Selection.clearSelection() Gui.activeDocument().activeView().setCameraType("Orthographic") 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) if render_main: resolution = (2000,2000) workbench.context.setAllStagesVisible(True) self.render(resolution, filename) img_full = Image.new('RGB', resolution, (255, 255, 255)) img = Image.open(filename) img_full.paste(img, None, img.getchannel('A')) img_full.save(filename) 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.7, 0.7, 0.7), (1,1,1)) workbench.context.setActiveStage(stage_id) workbench.context.setAllStagesVisible(False) self.render(resolution, dir + "/stages/" + stage_name + ".png", mask_stages_below=stage_id) # merge images bg = Image.new('RGB', resolution, (255,255,255)) if prev_stage_id is not None: prev_stage = Image.open(dir + "/stages/" + stage_name + "-bg.png") prev_stage.load() os.remove(dir + "/stages/" + stage_name + "-bg.png") bg.paste(prev_stage, None, prev_stage.getchannel('A')) fg = Image.open(dir + "/stages/" + stage_name + ".png") bg.paste(fg, None, fg.getchannel('A')) bg.save(dir + "/stages/" + stage_name + ".png") if prev_stage_id is not None: pass prev_stage_id = stage_id if render_parts: rendered_references = [] max_resolution = 1500 max_length = 2000 min_resolution = 250 shutil.rmtree(dir + "/parts", ignore_errors=True) os.makedirs(dir + "/parts", exist_ok=True) count = 0 for part_to_render in doc.Objects: if part_to_render.TypeId in ['Part::Feature', 'App::Part']: if 'Base_Reference' not in part_to_render.PropertiesList or part_to_render.Base_Reference == "" or part_to_render.Base_Reference in rendered_references: continue parent = part_to_render.Parents[0][0] if len(part_to_render.Parents) > 0 else None if parent is not None and parent.TypeId != 'App::Part': parent = None if parent is not None and 'Base_Reference' in parent.PropertiesList and parent.Base_Reference != "": continue #if not part_to_render.Label.startswith("L") and not part_to_render.Label.startswith("M") and not part_to_render.Label.startswith("T") and not part_to_render.Label.startswith("R"): # continue rendered_references.append(part_to_render.Base_Reference) for part_to_hide in doc.Objects: if part_to_hide.TypeId == 'Part::Feature': part_to_hide.ViewObject.Visibility = False part_to_render.ViewObject.Visibility = True if part_to_render.TypeId == 'App::Part': for feature in part_to_render.Group: feature.ViewObject.Visibility = True dimensions = [part_to_render.Shape.BoundBox.XLength, part_to_render.Shape.BoundBox.YLength, part_to_render.Shape.BoundBox.ZLength] main_axis: int = 0; main_length: float = part_to_render.Shape.BoundBox.XLength if part_to_render.Shape.BoundBox.YLength > main_length: main_axis = 1 main_length = part_to_render.Shape.BoundBox.YLength if part_to_render.Shape.BoundBox.ZLength > main_length: main_axis = 2 main_length = part_to_render.Shape.BoundBox.ZLength center = coin.SbVec3f(part_to_render.Shape.BoundBox.Center) offset = coin.SbVec3f(800, -2000, 800) up = coin.SbVec3f(0, 0, 1) if main_axis == 1: offset[0], offset[1] = -offset[1], offset[0] elif main_axis == 2: offset[0], offset[2] = -offset[2], offset[0] up = coin.SbVec3f(-1, 0, 0) cam = Gui.ActiveDocument.ActiveView.getCameraNode() cam.position = center + offset cam.pointAt(center, up) Gui.ActiveDocument.ActiveView.fitAll() cross_length = 0 if main_axis == 0: cross_length = max(dimensions[1] * 0.5, dimensions[2]) elif main_axis == 1: cross_length = max(dimensions[0] * 0.5, dimensions[2]) elif main_axis == 2: cross_length = max(dimensions[0], dimensions[1] * 0.5) cross_length = max(cross_length, main_length * 0.25 + cross_length * 0.8) main_length_pixels = max(min_resolution, int(math.sqrt(min(1.0, main_length / max_length)) * max_resolution)) cross_length_pixels = int(cross_length * main_length_pixels / main_length) cam.scaleHeight(0.05 + 0.95 * cross_length_pixels / main_length_pixels) resolution = (main_length_pixels, cross_length_pixels) part_filename = dir + "/parts/" + part_to_render.Base_Reference.upper().replace(' ', '-') + ".png" self.render(resolution, part_filename) bg = Image.new('RGB', resolution, (255, 255, 255)) fg = Image.open(part_filename) bg.paste(fg, None, fg.getchannel('A')) bg.save(part_filename) count = count + 1 if count == 10: pass workbench.context.setAllStagesVisible(True) from ahb_command import AHB_CommandWrapper AHB_CommandWrapper.addGuiCommand('AHB_render', AHB_Render())