forked from youen/assembly_handbook
333 lines
16 KiB
Python
333 lines
16 KiB
Python
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]+2, resolution[1]+2, "#ff0000")
|
|
|
|
self.set_render_outlines(mask_stages_below=mask_stages_below)
|
|
Gui.ActiveDocument.ActiveView.saveImage(temp_shapes_file_name, (resolution[0]+2) * 2, (resolution[1]+2) * 2, "#ffffff")
|
|
|
|
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
|
|
])
|
|
|
|
# crop 1px borders
|
|
result = result.crop((1, 1, result.size[0] - 1, result.size[1] - 1))
|
|
|
|
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
|
|
import re
|
|
|
|
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 = (3000,3000)
|
|
workbench.context.setAllStagesVisible(True)
|
|
self.render(resolution, filename)
|
|
self.reset_display()
|
|
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 = (3000, 3000)
|
|
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
|
|
self.reset_display()
|
|
|
|
if render_parts:
|
|
rendered_references = []
|
|
max_length = 1500
|
|
max_resolution = 800
|
|
min_resolution = 350
|
|
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.pow(min(1.0, main_length / max_length), 0.3) * 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/" + re.sub('[^A-Z0-9_]', '-', part_to_render.Base_Reference.upper()) + ".png"
|
|
self.render(resolution, part_filename)
|
|
|
|
image = Image.new('RGB', resolution, (255, 255, 255))
|
|
fg = Image.open(part_filename)
|
|
image.paste(fg, None, fg.getchannel('A'))
|
|
#image.save(part_filename)
|
|
|
|
final_apsect_ratio = 2.0 # width/height
|
|
final_resolution = (max_resolution, int(max_resolution / final_apsect_ratio))
|
|
final_size = image.size
|
|
if final_size[0] > final_resolution[0]:
|
|
final_size = (final_size[0] * final_resolution[0] / final_size[0], final_size[1] * final_resolution[0] / final_size[0])
|
|
if final_size[1] > final_resolution[1]:
|
|
final_size = (final_size[0] * final_resolution[1] / final_size[1], final_size[1] * final_resolution[1] / final_size[1])
|
|
final_size = (int(final_size[0]), int(final_size[1]))
|
|
|
|
if final_size[0] < image.size[0] or final_size[1] < image.size[1]:
|
|
image.thumbnail(final_size, Image.ANTIALIAS)
|
|
|
|
final_image = Image.new('RGB', final_resolution, (255, 255, 255))
|
|
final_image.paste(image, (final_resolution[0] - final_size[0], int(final_resolution[1]/2 - final_size[1]/2)))
|
|
final_image.save(part_filename)
|
|
|
|
count = count + 1
|
|
if count == 10:
|
|
pass
|
|
|
|
self.reset_display()
|
|
|
|
workbench.context.setAllStagesVisible(True)
|
|
|
|
from ahb_command import AHB_CommandWrapper
|
|
AHB_CommandWrapper.addGuiCommand('AHB_render', AHB_Render())
|