added show function to igeshandler + plot now with matlplotlib and sa… (#56)

* added show function to igeshandler + plot now with matlplotlib and save stl of iges in png format

* added show function to igeshandler + plot now with matlplotlib and save stl of iges in png format

* added show function to igeshandler + plot now with matlplotlib and save stl of iges in png format

* added show function to igeshandler + plot now with matlplotlib and save stl of iges in png format

Author: fsalmoir

pygem/igeshandler.py

@@ -2,7 +2,7 @@
 Utilities for reading and writing different CAD files.
 """
 import numpy as np
-import sys
+import os
 import pygem.filehandler as fh
 from OCC.IGESControl import (IGESControl_Reader, IGESControl_Writer)
 from OCC.BRep import (BRep_Tool, BRep_Builder)
@@ -14,6 +14,10 @@
 from OCC.gp import (gp_Pnt, gp_XYZ)
 from OCC.Display.SimpleGui import init_display
 from OCC.ShapeFix import ShapeFix_ShapeTolerance
+from OCC.StlAPI import StlAPI_Writer
+from mpl_toolkits import mplot3d
+from matplotlib import pyplot
+from stl import mesh
 
 
 class IgesHandler(fh.FileHandler):
@@ -198,9 +202,6 @@ def plot(self, plot_file=None, save_fig=False):
 		:param string plot_file: the iges filename you want to plot.
 		:param bool save_fig: a flag to save the figure in png or not. If True the
 			plot is not shown.
-			
-		.. warning::
-			It does not work well up to now
 		"""
 		if plot_file is None:
 			plot_file = self.infile
@@ -213,40 +214,53 @@ def plot(self, plot_file=None, save_fig=False):
 		reader.TransferRoots()
 		shape = reader.Shape()		
 
-		display, start_display, add_menu, add_function_to_menu = init_display()
-		display.FitAll()
-		display.DisplayShape(shape, update=True)
+		stl_writer = StlAPI_Writer()
+		# Do not switch SetASCIIMode() from False to True.
+		stl_writer.SetASCIIMode(False)
+		stl_writer.Write(shape,'aux_figure.stl')
 
-		def export_to_BMP(event=None):
-			display.View.Dump('./capture_bmp.bmp')
-
-
-		def export_to_PNG(event=None):
-			display.View.Dump('./capture_png.png')
-
+		# Create a new plot
+		figure = pyplot.figure()
+		axes = mplot3d.Axes3D(figure)
 
-		def export_to_JPEG(event=None):
-			display.View.Dump('./capture_jpeg.jpeg')
+		# Load the STL files and add the vectors to the plot
+		stl_mesh = mesh.Mesh.from_file('aux_figure.stl')
+		axes.add_collection3d(mplot3d.art3d.Poly3DCollection(stl_mesh.vectors))
 
+		# Auto scale to the mesh size
+		scale = stl_mesh.points.flatten(-1)
+		axes.auto_scale_xyz(scale, scale, scale)
+		
+		# Show the plot to the screen
+		if not save_fig:
+			pyplot.show()
+		else:
+			figure.savefig(plot_file.split('.')[0] + '.png')
+			
+		os.remove('aux_figure.stl')
+			
+			
+	def show(self, show_file=None):
+		"""
+		Method to show an iges file. If `show_file` is not given it plots `self.infile`.
 
-		def export_to_TIFF(event=None):
-			display.View.Dump('./capture_tiff.tiff')
-
+		:param string show_file: the iges filename you want to plot.
+		"""
+		if show_file is None:
+			show_file = self.infile
+		else:
+			self._check_filename_type(show_file)
 
-		def exit(event=None):
-			sys.exit()
+		## read in the IGES file
+		reader = IGESControl_Reader()
+		reader.ReadFile(show_file)
+		reader.TransferRoots()
+		shape = reader.Shape()		
 
-		add_menu('screencapture')
-		add_function_to_menu('screencapture', export_to_BMP)
-		add_function_to_menu('screencapture', export_to_PNG)
-		add_function_to_menu('screencapture', export_to_JPEG)
-		add_function_to_menu('screencapture', export_to_TIFF)
-		add_function_to_menu('screencapture', exit)
+		display, start_display, add_menu, add_function_to_menu = init_display()
+		display.FitAll()
+		display.DisplayShape(shape, update=True)
 
 		# Show the plot to the screen
-		if not save_fig:
-			start_display()
-		else:
-			f = display.View.View().GetObject()
-			display.View.Dump(plot_file.split('.')[0] + '.ppm')
+		start_display()
 

tests/test_igeshandler.py

@@ -138,6 +138,14 @@ def test_iges_write_comparison(self):
 
 		iges_handler.write(mesh_points, outfilename)
 		os.remove(outfilename)
+		
+		
+	def test_iges_plot_save_fig(self):
+		iges_handler = ih.IgesHandler()
+		mesh_points = iges_handler.parse('tests/test_datasets/test_pipe.iges')
+		iges_handler.plot(save_fig=True)
+		self.assertTrue(os.path.isfile('tests/test_datasets/test_pipe.png'))
+		os.remove('tests/test_datasets/test_pipe.png')
 
 
 	def test_iges_plot_failing_outfile_type(self):