VTK point cloud animation viewer

viewer.py

"""
Helper script to view a .ply point cloud animation (or a single point cloud).

Info:
    Takes as input a directory, which must contain point clouds in .ply format,
        as well as a sketchfab.timeframe file specifying the order of the files, 
        and the duration of each frame (although the duration is not used here).
    More information on the sketchfab.timeframe format can be obtained here:
        https://help.sketchfab.com/hc/en-us/articles/203058018-Animations#timeframe
    Please note that the input files must be generated else where (manually, C++, Python...)

Usage:
    python viewer.py /path/to/the/pointclouds/directory
    OR
    python viewer.py /path/to/the/pointcloud.ply
    python viewer.py --help

Installation:
    This script relies on multiple libraries, which you'll need to install:
        * vtk     (install through pip or conda): https://pypi.org/project/vtk/
        * numpy   (install through pip or conda): https://numpy.org/
        * plyfile (pip or source file, no conda): https://github.com/dranjan/python-plyfile

License:
    WTFPL: do What the Fuck you Want Public License: http://www.wtfpl.net/ :)

Examples of models this script was used to help create:
    * "A Windy Day":               https://skfb.ly/6OGwB
    * "Need some Space?":          https://skfb.ly/6JyVO
    * "Galactic Incident":         https://skfb.ly/6JBWp
    * "24h of crime in LA":        https://skfb.ly/6LVWM
    * "Earthquakes - 2010 & 2011": https://skfb.ly/6Mn8Y
"""

import os
import sys
import argparse

import vtk
import numpy as np
import plyfile


class VTKPointCloud:
    """A class to store points from a .ply point cloud"""

    def __init__(self):

        # Create all the usual VTK pipeline
        self.points   = vtk.vtkPoints()
        self.vertices = vtk.vtkCellArray()
        self.colors   = vtk.vtkUnsignedCharArray()
        self.polydata = vtk.vtkPolyData()
        self.mapper   = vtk.vtkPolyDataMapper()
        self.actor    = vtk.vtkActor()

        # Prepare stuff for the color
        self.colors.SetNumberOfComponents(3)
        self.colors.SetName("Colors")
        self.actor.GetProperty().SetPointSize(2)

        # Create the links for the polydata
        self.polydata.SetPoints(self.points)
        self.polydata.SetVerts(self.vertices)
        self.polydata.GetPointData().SetScalars(self.colors)
        self.polydata.Modified()

        # Set the mapper / actor links
        self.mapper.SetInputData(self.polydata)
        self.actor.SetMapper(self.mapper)

    def addpoints(self, path):

        try:
            # Read the vertices
            data = plyfile.PlyData.read(path)['vertex']
            xyz = np.c_[data['x'], data['y'], data['z']]
            rgb = np.c_[data['red'], data['green'], data['blue']]

            # Add the data to VTK structures
            for i in range(0, len(xyz)):
                p = xyz[i]
                id = self.points.InsertNextPoint(p)
                self.vertices.InsertNextCell(1)
                self.vertices.InsertCellPoint(id)
                self.colors.InsertNextTuple3(rgb[i][0], rgb[i][1], rgb[i][2])
            
            # Update
            self.polydata.Modified()

        except:
            print("ERROR: cannot read %s" % path)
            sys.exit(1)

class TimerCallback():
    """Callback class to update the visualization"""

    def __init__(self, renderer, listsOfClouds):

        self.listsOfClouds = listsOfClouds
        self.renderer     = renderer
        self.counter      = 0

    def execute(self, iren, event):
 
        # Remove the previous actors before adding new ones, in a looping way
        for actor in self.renderer.GetActors():
            self.renderer.RemoveActor(actor)
        for cloud in self.listsOfClouds[self.counter%len(self.listsOfClouds)]:
            self.renderer.AddActor(cloud.actor)

        # Update the interactive renderer and frame counter
        iren.GetRenderWindow().Render()
        self.counter += 1


def parse_arguments():

    # Create the parser
    desc = 'Displays an animated point cloud for Sketchfab export\n'
    desc += "https://help.sketchfab.com/hc/en-us/articles/203058018-Animations#timeframe"
    parser = argparse.ArgumentParser(description=desc)
    parser.add_argument("input", help="Directory containing .ply files and sketchfab.timeframe")
    args   = parser.parse_args()

    # Check if the directory is valid
    if os.path.exists(args.input):
        if os.path.isdir(args.input):
            time_file = os.path.join(args.input, "sketchfab.timeframe")
            if not os.path.exists(time_file):
                print("ERROR: %s must contain a file named sketchfab.timeframe" % args.input)
                sys.exit(1)
            if len([f for f in os.listdir(args.input) if ".ply" in f]) == 0:
                print("ERROR: %s must contain .ply files" % args.input)
                sys.exit(1)
        elif os.path.isfile(args.input):
            if not args.input.endswith(".ply"):
                print("ERROR: %s is not a .ply file" % args.input)
                sys.exit(1)
        else:
            print("ERROR: %s is neither or a file nor a directory" % args.input)
            sys.exit(1)
    else:
        print("ERROR: %s does not exist" % args.input)
        sys.exit(1)

    return args

def parse_data(_input):
    """Create VTK point clouds for every .ply file in the input directory"""

    if os.path.isdir(_input): # _input is a directory

        # Get the list of files to parse
        FILENAMES = []
        with open(os.path.join(_input, "sketchfab.timeframe")) as f:
            lines = [l.strip() for l in f.readlines() if len(l)>20]
            for l in lines:
                duration = float(l.split(" ")[0])
                files    = (l.split(" ")[1]).split("+")
                FILENAMES.append(files)

        # Make a set from the list (small optimization)
        FILES = list(set([item for sublist in FILENAMES for item in sublist]))
        FILES.sort()

        # Check that the files exist (sketchfab.timeframe was correct)
        for f in FILES:
            if not os.path.exists(os.path.join(_input, f)):
                print("ERROR: %s does not exist in %s" % (f, _input))
                sys.exit(1)

        # Create a VTK point cloud for every file
        VTKCLOUDS = {}
        for f in FILES:
            print("---- Reading %s" % f)
            PC = VTKPointCloud()
            PC.addpoints(os.path.join(_input, f))
            VTKCLOUDS[f] = PC

        # Return the VTK clouds, ordered by frame
        POINTCLOUDS = [ [ VTKCLOUDS[f] for f in files ] for files in FILENAMES ]
        return POINTCLOUDS

    else: # _input is a .ply file

        print("---- Reading %s" % _input)
        PC = VTKPointCloud()
        PC.addpoints(_input)
        return [[PC]]


if __name__ == "__main__":

    # Get the argument
    args   = parse_arguments()

    # Parse the data
    print("-- Reading and parsing file(s)")
    actors = parse_data(args.input)

    # Create and link the usual VTK rendering stuff
    print("-- Initializing VTK")
    renderer               = vtk.vtkRenderer()
    renderWindow           = vtk.vtkRenderWindow()
    renderWindowInteractor = vtk.vtkRenderWindowInteractor()
    renderWindow.AddRenderer(renderer)
    renderWindowInteractor.SetRenderWindow(renderWindow)
    renderWindowInteractor.Initialize()

    # Set background color
    renderer.SetBackground(0.1, 0.1, 0.1)

    # Initialize the render with the first point clouds
    for cloud in actors[0]:
        renderer.AddActor(cloud.actor)
    renderer.ResetCamera()

    # Link the timer with 30 fps
    if len(actors) > 1:
        timerCB = TimerCallback(renderer, actors)
        renderWindowInteractor.AddObserver('TimerEvent', timerCB.execute)
        timerId = renderWindowInteractor.CreateRepeatingTimer(int(1000/30.))
        timerCB.timerId = timerId

    # Run the application
    print("-- Running the viewer")
    renderWindow.Render()
    renderWindowInteractor.Start()
    

Hi there Norgeotloic. I've successfully implemented your code. Thank you very much. It's great. I was wondering if there was a way of improving the look? Is it possible to increase or colour the point size and / or make them into voxcels whilst maintaining the animation? I'm also finding it difficult to control the visualisation. Would it be easy to add simple key controls? I realise my question reveals my lack of knowledge. Any pointers on how to do this would great. Best wishes, Thingsontopofotherthings.

A newbee here. How can I create .ply files from these datas (https://psl.noaa.gov/data/gridded/data.ncep.reanalysis.html)? Any hint will be very helpful.