# A script to generate spherical rendering of STL Files

stl_render.py

#!/usr/bin/env python

# A script to generate spherical rendering of STL Files
# Use with python STL_viewer.py /path/to/folder
# or python STL_viewer.py /path/to/file

# Written by Billy Zhou (billyzs.728 # gmail.com)

import vtk
import argparse, os
from math import sin, cos, radians
parser = argparse.ArgumentParser()
parser.add_argument("path", help="path to folder or STL file")


class STLViewer:
    def __init__(self, path, origin, position=None):
        self.path = path
        self.file_name = self.path.split("/")[-1].split(".")[0]
        # self.rotation = angle
        self.origin = origin
        self.camera = None
        self.reader = None
        self.mapper = None
        self.actor = None
        self.ren = None
        self.renWin = None
        self.iren = None
        self.w2if = None
        self.writer = None
        self.set_reader()
        self.set_mapper(vtk.vtkPolyDataMapper())

        self.set_actor(vtk.vtkActor())
        self.set_renderer(vtk.vtkRenderer())
        self.set_camera(vtk.vtkCamera(), position)

        self.set_render_window(vtk.vtkRenderWindow())
        self.set_iren(vtk.vtkRenderWindowInteractor())
        self.set_w2if(vtk.vtkWindowToImageFilter())
        self.set_writer(vtk.vtkPNGWriter())

    def set_reader(self):
        self.reader = vtk.vtkSTLReader()
        self.reader.SetFileName(self.path)

    def set_mapper(self, i):
        self.mapper = i
        if vtk.VTK_MAJOR_VERSION <= 5:
            self.mapper.SetInput(self.reader.GetOutput())
        else:
            self.mapper.SetInputConnection(self.reader.GetOutputPort())

    def set_actor(self, a):
        self.actor = a
        self.actor.SetMapper(self.mapper)
        self.actor.SetOrigin(self.origin)

    def set_camera(self, c, position=None):
        self.camera = c
        if position is not None:
            self.camera.SetPosition(position[0], position[1], position[2])
        self.camera.SetFocalPoint(self.origin)
        self.ren.SetActiveCamera(self.camera)
         
    def set_renderer(self, r):
        self.ren = r
        self.ren.AddActor(self.actor)
    
    def set_render_window(self, rw):
        self.renWin = rw
        self.renWin.AddRenderer(self.ren)
                
    def set_iren(self, iren):
        self.iren = iren
        self.iren.SetRenderWindow(self.renWin)

    def set_w2if(self, w2if):
        self.w2if = w2if
        self.w2if.SetInput(self.renWin)

    def set_writer(self, w):
        self.writer = w
        if vtk.VTK_MAJOR_VERSION <= 6:
            self.writer.SetInputData(self.w2if.GetOutput())
        else:
            self.writer.SetInput(self.w2if.GetOutput())
        
    def start_ui(self):
        self.iren.Initialize()
        self.renWin.Render()
        self.iren.Start()
    
    def screenshot(self, name=None):
        file_name = self.file_name
        if name is not None:
            file_name += name

        # self.renWin.AddRenderer(self.ren)

        self.w2if.SetInput(self.renWin)
        self.renWin.Render()
        self.w2if.Update()

        self.writer.SetInputData(self.w2if.GetOutput())
        self.writer.SetFileName(file_name + ".png")
        self.writer.Write()

def spherical_to_cart(r, yr, zr):
    return (r * cos(yr) * sin(zr), r * sin(yr), r * cos(yr) * cos(zr))

def spherical_screen_shot(file, pole, r):
    for y in range(-90, 90, 15):
        for z in range(0, 360, 15):
            (yr, zr) = (radians(a) for a in [y, z])
            pos = spherical_to_cart(r, yr, zr)
            # print pos
            # (xr, yr, zr) = pos
            v = STLViewer(file, pole, pos)
            v.screenshot("_x=" + str(0) + "_y=" + str(y) + "_z=" + str(z))

if __name__ == '__main__':
    args = parser.parse_args()
    # print args.path
    pole = [0., 60., 0.]
    r = 180
    x = 0

    if os.path.isdir(args.path):
        (_, _, file_names) = os.walk(os.getcwd()).next()
        filter(lambda x: x[-4:] == ".STL" ,file_names)
        for file in file_names:
            spherical_screen_shot(file, pole, r)
    else:
        spherical_screen_shot(args.path, pole, r)
        # pos = spherical_to_cart(r, radians(-75), radians(345))
        # v = STLViewer(args.path, pole, pos)
        # v.screenshot("_x=" + str(0) + "_y=" + str(-75) + "_z=" + str(345))
        # v.start_ui()