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 render(self):
		from pivy import coin
		import os
		from PIL import Image, ImageDraw, ImageChops
		
		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)
		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
			else:
				link.LinkedObject = part
				
		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()
		
		docView.saveImage(self.image_file_name, 4096, 4096, "#ffffff")
		
		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()