forked from youen/assembly_handbook
You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
192 lines
9.0 KiB
192 lines
9.0 KiB
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.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) |
|
|
|
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" |
|
|
|
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") |
|
|
|
# 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,0.0)) |
|
Gui.ActiveDocument.ActiveView.saveImage(temp_lines_file_name, resolution[0], resolution[1], "#ff0000") |
|
|
|
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)) |
|
|
|
# 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: 0 if p == 255 else 255)) |
|
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: 0 if p == 255 else 255)) |
|
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())
|
|
|