assembly_handbook/ahb_raster_view.py

import FreeCAD as App
import FreeCADGui as Gui

class RasterView:
	def __init__(self, view):
		self.source_view = view
		doc = view.Document
		self.image_file_name = doc.FileName.replace('.FCStd', '') + '_raster/' + view.Name + '.png'
		
	def init_image_projection(self):
		doc = self.source_view.Document
		image_name = self.source_view.Label + "_raster"
		image = doc.getObject(image_name)
		if image is None:
			return False
		
		self.image_view = image
		
		if image.Assembly_handbook_ViewVolumeWidth > 0:
			self._precompute_image_projection()
			return True
		
		return False
		
	def init_image(self):
		workbench = Gui.getWorkbench("AssemblyHandbookWorkbench") #: :type workbench: AssemblyHandbookWorkbench
		
		doc = self.source_view.Document
		page = workbench.techDrawExtensions.getViewPage(self.source_view)
		
		image_name = self.source_view.Label + "_raster"
		image = doc.getObject(image_name)
		if image is None:
			image = doc.addObject('TechDraw::DrawViewImage', image_name)
			image.addProperty("App::PropertyFloat", "Assembly_handbook_ViewVolumeWidth", "Assembly_handbook")
			image.addProperty("App::PropertyFloat", "Assembly_handbook_ViewVolumeHeight", "Assembly_handbook")
			image.addProperty("App::PropertyFloat", "Assembly_handbook_ViewVolumeDepth", "Assembly_handbook")
			image.addProperty("App::PropertyVector", "Assembly_handbook_ViewVolumeOffset", "Assembly_handbook")
		
		if not image in page.Views:
			page.addView(image)
			
			new_views_list = page.Views
			new_views_list.remove(image)
			view_idx = new_views_list.index(self.source_view)
			new_views_list.insert(view_idx, image)
			page.Views = new_views_list
		
		self.image_view = image
		
		if image.Assembly_handbook_ViewVolumeWidth > 0:
			self._precompute_image_projection()
			
	def _precompute_image_projection(self):
		YDirection = self.source_view.Direction.cross(self.source_view.XDirection)
		
		self.image_x_dir = self.source_view.XDirection / self.image_view.Assembly_handbook_ViewVolumeWidth
		self.image_y_dir = YDirection / self.image_view.Assembly_handbook_ViewVolumeHeight
		self.image_z_dir = self.source_view.Direction / self.image_view.Assembly_handbook_ViewVolumeDepth
		
		self.image_x_dir_inv = self.source_view.XDirection * self.image_view.Assembly_handbook_ViewVolumeWidth
		self.image_y_dir_inv = YDirection * self.image_view.Assembly_handbook_ViewVolumeHeight
		self.image_z_dir_inv = self.source_view.Direction * self.image_view.Assembly_handbook_ViewVolumeDepth
		
	def project3DPointToImageView(self, point3d):
		offset = self.image_view.Assembly_handbook_ViewVolumeOffset
		return App.Vector(self.image_x_dir.dot(point3d) + offset.x, self.image_y_dir.dot(point3d) + offset.y, self.image_z_dir.dot(point3d) + offset.z)
	
	def project3DPointToSourceView(self, point3d):
		offset = self.image_view.Assembly_handbook_ViewVolumeOffset
		offset = App.Vector((offset.x-0.5) * self.image_view.Assembly_handbook_ViewVolumeWidth, (offset.y-0.5) * self.image_view.Assembly_handbook_ViewVolumeHeight, (offset.z-0.5) * self.image_view.Assembly_handbook_ViewVolumeDepth)
		#image_view_point = App.Vector(self.image_x_dir.dot(point3d), self.image_y_dir.dot(point3d), self.image_z_dir.dot(point3d))
		#return App.Vector(image_view_point.x * self.image_view.Assembly_handbook_ViewVolumeWidth, image_view_point.y * self.image_view.Assembly_handbook_ViewVolumeHeight, 0)
		
		YDirection = self.source_view.Direction.cross(self.source_view.XDirection)
		return App.Vector(self.source_view.XDirection.dot(point3d) + offset.x, YDirection.dot(point3d) + offset.y, self.image_z_dir.dot(point3d) + offset.z)
	
	def projectImageViewPointTo3D(self, point2d):
		offset = self.image_view.Assembly_handbook_ViewVolumeOffset
		p = point2d - offset
		return self.image_x_dir_inv * p.x + self.image_y_dir_inv * p.y + self.image_z_dir_inv * p.z
	
	def _flatten_objects_tree(self, obj_list):
		result = []
		for obj in obj_list:
			if obj.TypeId == 'Part::FeaturePython' and hasattr(obj, 'LinkedObject'): # variant link
				result.extend(self._flatten_objects_tree(obj.Group))
			elif obj.TypeId in ['App::Link']:
				result.extend(self._flatten_objects_tree([obj.LinkedObject]))
			elif obj.TypeId in ['App::Part', 'App::DocumentObjectGroup']:
				result.extend(self._flatten_objects_tree(obj.Group))
			elif obj.TypeId in ['Part::Feature', 'Part::FeaturePython', 'PartDesign::Body', 'PartDesign::CoordinateSystem', 'PartDesign::Line', 'Part::Mirroring', 'Part::Cut']:
				result.append(obj)
				if hasattr(obj, 'Group'):
					result.extend(self._flatten_objects_tree(obj.Group))
			
		return result
	
	def _should_render(self, obj):
		return obj.TypeId in ['Part::Feature', 'Part::FeaturePython', 'PartDesign::Body', 'Part::Mirroring', 'Part::Cut']
		
	def render(self):
		from pivy import coin
		import os
		from PIL import Image, ImageDraw, ImageChops
		
		workbench = Gui.getWorkbench("AssemblyHandbookWorkbench") #: :type workbench: AssemblyHandbookWorkbench
		
		view = self.source_view
		
		self.init_image()
		
		print('Rasterizing ' + view.Label + " to " + self.image_file_name + "...")
		
		dir = os.path.dirname(self.image_file_name)
		if not os.path.exists(dir):
			os.makedirs(dir)
			
		doc = App.newDocument('tmp_raster', hidden=False, temp=False)
			
		objects_to_reset = {}
		try:
			for part in view.XSource:
				link = doc.addObject('App::Link', part.Name)
				link.Label = part.Label
				if part.TypeId == 'App::Link':
					link.LinkedObject = part.LinkedObject
					link.Placement = part.Placement
				elif part.TypeId == 'Part::FeaturePython' and hasattr(part, 'LinkedObject'): # variant link
					link.LinkedObject = part.LinkedObject
					link.Placement = part.Placement
				else:
					link.LinkedObject = part
					
				new_part = workbench.techDrawExtensions.isNewPartInView(view, part)
				
				link.ViewObject.OverrideMaterial = True
				link.ViewObject.ShapeMaterial.DiffuseColor = (0.0, 0.0, 0.0, 0.0) if new_part else (0.5, 0.5, 0.5, 0.0) # this actually changes the line color
				
				# in current version of freecad, link override material does not allow to override all material properties, for example emissive color, so we have to change material of the linked object
				for obj in self._flatten_objects_tree([link]):
					if obj in objects_to_reset.keys():
						continue
					
					if self._should_render(obj):
						objects_to_reset[obj] = (
							obj.ViewObject.LineColor,
							obj.ViewObject.ShapeMaterial.AmbientColor,
							obj.ViewObject.ShapeMaterial.DiffuseColor,
							obj.ViewObject.ShapeMaterial.SpecularColor,
							obj.ViewObject.ShapeMaterial.EmissiveColor,
							obj.ViewObject.LineWidth
						)
						
						obj.ViewObject.LineColor = (0.0, 0.0, 0.0, 0.0)
						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 = (1.0, 1.0, 1.0, 0.0)
						
						# We need to set two different values otherwise freecad does not always update LineWidth of sub-elements
						obj.ViewObject.LineWidth = 1.0						
						obj.ViewObject.LineWidth = 2.0
					else:
						objects_to_reset[obj] = (
							obj.ViewObject.Visibility,
						)
						
						obj.ViewObject.Visibility = False
					
			docView = Gui.getDocument(doc.Name).mdiViewsOfType('Gui::View3DInventor')[0]
			
			cam = docView.getCameraNode()
			
			rot = coin.SbRotation(coin.SbVec3f(1,0,0), coin.SbVec3f(view.XDirection.x,view.XDirection.y,view.XDirection.z))
			rot *= coin.SbRotation(coin.SbVec3f(0,0,1), coin.SbVec3f(view.Direction.x,view.Direction.y,view.Direction.z))
			cam.orientation.setValue(rot)
			
			docView.fitAll()
			
			viewVolume = cam.getViewVolume(0.0)
			self.image_view.Assembly_handbook_ViewVolumeWidth = viewVolume.getWidth()
			self.image_view.Assembly_handbook_ViewVolumeHeight = viewVolume.getHeight()
			self.image_view.Assembly_handbook_ViewVolumeDepth = viewVolume.getDepth()
			
			max_res = 4096
			docView.saveImage(self.image_file_name, int(min(max_res, viewVolume.getWidth() * view.Scale * 10)), int(min(max_res, viewVolume.getHeight() * view.Scale * 10)), "#ffffff")
			
		finally:
			for obj, props in objects_to_reset.items():
				if self._should_render(obj):
					obj.ViewObject.LineColor = props[0]
					obj.ViewObject.ShapeMaterial.AmbientColor = props[1]
					obj.ViewObject.ShapeMaterial.DiffuseColor = props[2]
					obj.ViewObject.ShapeMaterial.SpecularColor = props[3]
					obj.ViewObject.ShapeMaterial.EmissiveColor = props[4]
					obj.ViewObject.LineWidth = props[5]
				else:
					obj.ViewObject.Visibility = props[0]
		
		with Image.open(self.image_file_name) as img:
			original_size = img.size
			
			bg = Image.new(img.mode, img.size, '#ffffff') # fills an image with the background color
			diff = ImageChops.difference(img, bg) # diff between the actual image and the background color
			bbox = diff.getbbox() # finds border size (non-black portion of the image)
			#print(bbox)
			#image_center = (bbox[0] + (bbox[2] - bbox[0])/2 - img.size[0]/2, bbox[1] + (bbox[3] - bbox[1])/2 - img.size[1]/2)
			#print(image_center)
			img = img.crop(bbox)
			
			draw = ImageDraw.Draw(img)
			
			def debugPoint(p3d):
				p2d = self.project3DPointToImageView(p3d)
				pp = App.Vector(p2d.x * original_size[0] - bbox[0], (1.0-p2d.y) * original_size[1] - bbox[1])
				#print('pp', pp)
				
				len = 100
				draw.line([(pp.x, pp.y-len), (pp.x, pp.y+len)], fill=128, width = 7)
				draw.line([(pp.x-len, pp.y), (pp.x+len, pp.y)], fill=128, width = 7)
			
			#debugPoint(App.Vector(-12.5, 37.5, 25.0))
			#debugPoint(App.Vector(-12.5, -1387.5, 25.0))
			#debugPoint(App.Vector(131.23702882966705, -655.0000021095163, 145.21130178331268))
			
			img.save(self.image_file_name)
			
			sb_offset = viewVolume.projectToScreen(coin.SbVec3f(0,0,0))
			crop_offset = App.Vector(((bbox[0] + bbox[2])/2 - original_size[0]/2)/original_size[0], ((bbox[1] + bbox[3])/2 - original_size[1]/2)/original_size[1], 0)
			self.image_view.Assembly_handbook_ViewVolumeOffset = App.Vector(sb_offset[0] - crop_offset.x, sb_offset[1] + crop_offset.y, sb_offset[2])
			
			self._precompute_image_projection()
			
			p2dA = self.project3DPointToImageView(App.Vector(0,0,0))
			p2dB = self.project3DPointToImageView(view.XDirection)
			imageScale = view.Scale / (p2dB.x - p2dA.x) / original_size[0] * 10
			#print('imageScale', imageScale)
		
		App.closeDocument(doc.Name)
		
		image = self.image_view
		image.ImageFile = ""
		image.Scale = imageScale
		image.X = view.X
		image.Y = view.Y
		image.ImageFile = self.image_file_name # TODO: see if it's possible to set a relative path
		image.recompute()