vtk show method and some cleaning with pylint (#57)

* vtk show method and some cleaning with pylint

* show iges test

Author: mtezzele

pygem/igeshandler.py

@@ -1,12 +1,14 @@
 """
 Utilities for reading and writing different CAD files.
 """
-import numpy as np
 import os
-import pygem.filehandler as fh
+import numpy as np
+from mpl_toolkits import mplot3d
+from matplotlib import pyplot
 from OCC.IGESControl import (IGESControl_Reader, IGESControl_Writer)
 from OCC.BRep import (BRep_Tool, BRep_Builder)
-from OCC.BRepBuilderAPI import (BRepBuilderAPI_NurbsConvert, BRepBuilderAPI_MakeWire, BRepBuilderAPI_MakeEdge, BRepBuilderAPI_MakeFace)
+from OCC.BRepBuilderAPI import (BRepBuilderAPI_NurbsConvert, BRepBuilderAPI_MakeWire)
+from OCC.BRepBuilderAPI import (BRepBuilderAPI_MakeEdge, BRepBuilderAPI_MakeFace)
 from OCC.GeomConvert import geomconvert_SurfaceToBSplineSurface
 import OCC.TopoDS
 from OCC.TopAbs import (TopAbs_FACE, TopAbs_EDGE)
@@ -15,9 +17,8 @@
 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
+import pygem.filehandler as fh
 
 
 class IgesHandler(fh.FileHandler):
@@ -27,7 +28,8 @@ class IgesHandler(fh.FileHandler):
 	:cvar string infile: name of the input file to be processed.
 	:cvar string outfile: name of the output file where to write in.
 	:cvar string extension: extension of the input/output files. It is equal to '.iges'.
-	:cvar list control_point_position: index of the first NURBS control point (or pole) of each face of the iges file.
+	:cvar list control_point_position: index of the first NURBS control point (or pole)
+		of each face of the iges file.
 	"""
 	def __init__(self):
 		super(IgesHandler, self).__init__()
@@ -51,22 +53,21 @@ def parse(self, filename):
 		self._check_extension(filename)
 
 		self.infile = filename
-		
-		## read in the IGES file
+
+		# read in the IGES file
 		reader = IGESControl_Reader()
 		reader.ReadFile(self.infile)
 		reader.TransferRoots()
 		shape = reader.Shape()
 
-		## cycle on the faces to get the control points
+		# cycle on the faces to get the control points
 		# init some quantities
 		n_faces = 0
 		control_point_position = [0]
 		faces_explorer = TopExp_Explorer(shape, TopAbs_FACE)
-		mesh_points = np.zeros(shape=(0,3))
-		
-		while faces_explorer.More():
+		mesh_points = np.zeros(shape=(0, 3))
 
+		while faces_explorer.More():
 			# performing some conversions to get the right format (BSplineSurface)
 			iges_face = OCC.TopoDS.topods_Face(faces_explorer.Current())
 			iges_nurbs_converter = BRepBuilderAPI_NurbsConvert(iges_face)
@@ -81,27 +82,28 @@ def parse(self, filename):
 			# extract the Control Points of each face
 			n_poles_u = occ_face.NbUPoles()
 			n_poles_v = occ_face.NbVPoles()
-			control_polygon_coordinates = np.zeros(shape=(n_poles_u*n_poles_v,3))
+			control_polygon_coordinates = np.zeros(shape=(n_poles_u * n_poles_v, 3))
 
 			# cycle over the poles to get their coordinates
 			i = 0
 			for pole_u_direction in xrange(n_poles_u):
 				for pole_v_direction in xrange(n_poles_v):
-					control_point_coordinates = occ_face.Pole(pole_u_direction+1,pole_v_direction+1)
-					control_polygon_coordinates[i,:] = [control_point_coordinates.X(), control_point_coordinates.Y(), control_point_coordinates.Z()]
-					i += 1	
+					control_point_coordinates = occ_face.Pole(pole_u_direction+1, pole_v_direction+1)
+					control_polygon_coordinates[i, :] = [control_point_coordinates.X(), \
+						control_point_coordinates.Y(), control_point_coordinates.Z()]
+					i += 1
 
 			# pushing the control points coordinates to the mesh_points array (used for FFD)
 			mesh_points = np.append(mesh_points, control_polygon_coordinates, axis=0)
-			control_point_position.append(control_point_position[-1] + n_poles_u*n_poles_v)
+			control_point_position.append(control_point_position[-1] + n_poles_u * n_poles_v)
 
 			n_faces += 1
 			faces_explorer.Next()
 
 		self._control_point_position = control_point_position
 
 		return mesh_points
-		
+
 
 	def write(self, mesh_points, filename):
 		"""
@@ -117,18 +119,18 @@ def write(self, mesh_points, filename):
 		self._check_extension(filename)
 		self._check_infile_instantiation(self.infile)
 
-		self.outfile = filename	
+		self.outfile = filename
 
-		## init the ouput file writer
+		# init the ouput file writer
 		writer = IGESControl_Writer()
 
-		## read in the IGES file
+		# read in the IGES file
 		reader = IGESControl_Reader()
 		reader.ReadFile(self.infile)
 		reader.TransferRoots()
 		shape_read = reader.Shape()
 
-		## cycle on the faces to update the control points position
+		# cycle on the faces to update the control points position
 		# init some quantities
 		faces_explorer = TopExp_Explorer(shape_read, TopAbs_FACE)
 		n_faces = 0
@@ -137,9 +139,8 @@ def write(self, mesh_points, filename):
 		compound_builder = BRep_Builder()
 		compound = OCC.TopoDS.TopoDS_Compound()
 		compound_builder.MakeCompound(compound)
-		
+
 		while faces_explorer.More():
-	
 			# similar to the parser method
 			iges_face = OCC.TopoDS.topods_Face(faces_explorer.Current())
 			iges_nurbs_converter = BRepBuilderAPI_NurbsConvert(iges_face)
@@ -156,18 +157,19 @@ def write(self, mesh_points, filename):
 			i = 0
 			for pole_u_direction in xrange(n_poles_u):
 				for pole_v_direction in xrange(n_poles_v):
-					control_point_coordinates = mesh_points[i+control_point_position[n_faces],:]
-					point_xyz = gp_XYZ(control_point_coordinates[0], control_point_coordinates[1], control_point_coordinates[2])
+					control_point_coordinates = mesh_points[i+control_point_position[n_faces], :]
+					point_xyz = gp_XYZ(control_point_coordinates[0], control_point_coordinates[1], \
+						control_point_coordinates[2])
 					gp_point = gp_Pnt(point_xyz)
-					occ_face.SetPole(pole_u_direction+1,pole_v_direction+1,gp_point)
+					occ_face.SetPole(pole_u_direction+1, pole_v_direction+1, gp_point)
 					i += 1
 
-			## construct the deformed wire for the trimmed surfaces
+			# construct the deformed wire for the trimmed surfaces
 			wire_maker = BRepBuilderAPI_MakeWire()
 			tol = ShapeFix_ShapeTolerance()
 			brep = BRepBuilderAPI_MakeFace(occ_face.GetHandle(), 1e-4).Face()
 			brep_face = BRep_Tool.Surface(brep)
-		
+
 			# cycle on the edges
 			edge_explorer = TopExp_Explorer(nurbs_face, TopAbs_EDGE)
 			while edge_explorer.More():
@@ -181,10 +183,10 @@ def write(self, mesh_points, filename):
 				wire_maker.Add(edge_phis_coordinates)
 				edge_explorer.Next()
 
-			#grouping the edges in a wire
+			# grouping the edges in a wire
 			wire = wire_maker.Wire()
 
-			## trimming the surfaces
+			# trimming the surfaces
 			brep_surf = BRepBuilderAPI_MakeFace(occ_face.GetHandle(), wire).Shape()
 			compound_builder.Add(compound, brep_surf)
 			n_faces += 1
@@ -193,8 +195,8 @@ def write(self, mesh_points, filename):
 		writer.AddShape(compound)
 
 		writer.Write(self.outfile)
-		
-		
+
+
 	def plot(self, plot_file=None, save_fig=False):
 		"""
 		Method to plot an iges file. If `plot_file` is not given it plots `self.infile`.
@@ -212,55 +214,54 @@ def plot(self, plot_file=None, save_fig=False):
 		reader = IGESControl_Reader()
 		reader.ReadFile(plot_file)
 		reader.TransferRoots()
-		shape = reader.Shape()		
-		
+		shape = reader.Shape()
+
 		stl_writer = StlAPI_Writer()
 		# Do not switch SetASCIIMode() from False to True.
 		stl_writer.SetASCIIMode(False)
-		stl_writer.Write(shape,'aux_figure.stl')
-		
+		stl_writer.Write(shape, 'aux_figure.stl')
+
 		# Create a new plot
 		figure = pyplot.figure()
 		axes = mplot3d.Axes3D(figure)
 
 		# Load the STL files and add the vectors to the plot
 		stl_mesh = mesh.Mesh.from_file('aux_figure.stl')
+		os.remove('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`.
 
-		:param string show_file: the iges filename you want to plot.
+		:param string show_file: the iges filename you want to show.
 		"""
 		if show_file is None:
 			show_file = self.infile
 		else:
 			self._check_filename_type(show_file)
 
-		## read in the IGES file
+		# read in the IGES file
 		reader = IGESControl_Reader()
 		reader.ReadFile(show_file)
 		reader.TransferRoots()
-		shape = reader.Shape()		
-		
-		display, start_display, add_menu, add_function_to_menu = init_display()
+		shape = reader.Shape()
+
+		display, start_display, __, __ = init_display()
 		display.FitAll()
 		display.DisplayShape(shape, update=True)
-		
+
 		# Show the plot to the screen
 		start_display()
 

pygem/stlhandler.py

@@ -73,7 +73,7 @@ def write(self, mesh_points, filename, write_bin=False):
 		for i in range(0, n_triplets):
 			for j in range(0, 3):
 				data['vectors'][i][j] = mesh_points[3*i + j]
-				
+
 		if not write_bin:
 			stl_mesh.save(self.outfile, mode=1, update_normals=True)
 		else:

pygem/vtkhandler.py

@@ -99,9 +99,9 @@ def write(self, mesh_points, filename):
 
 	def plot(self, plot_file=None, save_fig=False):
 		"""
-		Method to plot an stl file. If `plot_file` is not given it plots `self.infile`.
+		Method to plot a vtk file. If `plot_file` is not given it plots `self.infile`.
 
-		:param string plot_file: the stl filename you want to plot.
+		:param string plot_file: the vtk 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.
 		"""
@@ -147,3 +147,50 @@ def plot(self, plot_file=None, save_fig=False):
 			figure.savefig(plot_file.split('.')[0] + '.png')
 
 
+	def show(self, show_file=None):
+		"""
+		Method to show a vtk file. If `show_file` is not given it shows `self.infile`.
+
+		:param string show_file: the vtk filename you want to show.
+		"""
+		if show_file is None:
+			show_file = self.infile
+		else:
+			self._check_filename_type(show_file)
+
+		# Read the source file.
+		reader = vtk.vtkUnstructuredGridReader()
+		reader.SetFileName(show_file)
+		reader.Update() # Needed because of GetScalarRange
+		output = reader.GetOutput()
+		scalar_range = output.GetScalarRange()
+
+		# Create the mapper that corresponds the objects of the vtk file
+		# into graphics elements
+		mapper = vtk.vtkDataSetMapper()
+		if vtk.VTK_MAJOR_VERSION <= 5:
+			mapper.SetInput(output)
+		else:
+			mapper.SetInputData(output)
+		mapper.SetScalarRange(scalar_range)
+
+		# Create the Actor
+		actor = vtk.vtkActor()
+		actor.SetMapper(mapper)
+
+		# Create the Renderer
+		renderer = vtk.vtkRenderer()
+		renderer.AddActor(actor)
+		# Set background color (white is 1, 1, 1)
+		renderer.SetBackground(20, 20, 20)
+
+		# Create the RendererWindow
+		renderer_window = vtk.vtkRenderWindow()
+		renderer_window.AddRenderer(renderer)
+
+		# Create the RendererWindowInteractor and display the vtk_file
+		interactor = vtk.vtkRenderWindowInteractor()
+		interactor.SetRenderWindow(renderer_window)
+		interactor.Initialize()
+		interactor.Start()
+

tests/test_igeshandler.py

@@ -153,3 +153,8 @@ def test_iges_plot_failing_outfile_type(self):
 		with self.assertRaises(TypeError):
 			iges_handler.plot(plot_file=3)
 
+
+	def test_iges_show_failing_outfile_type(self):
+		iges_handler = ih.IgesHandler()
+		with self.assertRaises(TypeError):
+			iges_handler.show(show_file=1.1)

tests/test_stlhandler.py

@@ -215,6 +215,13 @@ def test_stl_plot_save_fig_bin(self):
 		os.remove('tests/test_datasets/test_sphere_bin.png')
 
 
+	def test_stl_plot_save_fig_plot_file(self):
+		stl_handler = sh.StlHandler()
+		stl_handler.plot(plot_file='tests/test_datasets/test_sphere.stl', save_fig=True)
+		self.assertTrue(os.path.isfile('tests/test_datasets/test_sphere.png'))
+		os.remove('tests/test_datasets/test_sphere.png')
+
+
 	def test_stl_plot_failing_outfile_type(self):
 		stl_handler = sh.StlHandler()
 		with self.assertRaises(TypeError):

tests/test_vtkhandler.py

@@ -141,10 +141,23 @@ def test_vtk_plot_failing_outfile_type(self):
 			vtk_handler.plot(plot_file=1.1)
 
 
-	def test_vtk_plot_save_fig(self):
+	def test_vtk_plot_save_fig_infile(self):
 		vtk_handler = vh.VtkHandler()
 		mesh_points = vtk_handler.parse('tests/test_datasets/test_red_blood_cell.vtk')
 		vtk_handler.plot(save_fig=True)
 		self.assertTrue(os.path.isfile('tests/test_datasets/test_red_blood_cell.png'))
 		os.remove('tests/test_datasets/test_red_blood_cell.png')
 
+
+	def test_vtk_plot_save_fig_plot_file(self):
+		vtk_handler = vh.VtkHandler()
+		vtk_handler.plot(plot_file='tests/test_datasets/test_red_blood_cell.vtk', save_fig=True)
+		self.assertTrue(os.path.isfile('tests/test_datasets/test_red_blood_cell.png'))
+		os.remove('tests/test_datasets/test_red_blood_cell.png')
+
+
+	def test_vtk_show_failing_outfile_type(self):
+		vtk_handler = vh.VtkHandler()
+		with self.assertRaises(TypeError):
+			vtk_handler.show(show_file=1.1)
+