from vtkmodules.web import camera, javascriptMapping, arrayTypesMapping from paraview.web import data_writer, data_converter from vtkmodules.web.query_data_model import DataHandler from paraview.web.camera import update_camera from vtkmodules.web import iteritems from paraview import simple from paraview import servermanager from vtkmodules.vtkCommonCore import ( vtkFloatArray, vtkIdList, vtkUnsignedCharArray, vtkTypeUInt32Array, ) from vtkmodules.vtkCommonDataModel import vtkDataSetAttributes import json, os, math, gzip, shutil, hashlib # Global helper variables encode_codes = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" # ----------------------------------------------------------------------------- # Basic Dataset Builder # ----------------------------------------------------------------------------- class DataSetBuilder(object): def __init__(self, location, camera_data, metadata={}, sections={}): self.dataHandler = DataHandler(location) self.cameraDescription = camera_data self.camera = None for key, value in iteritems(metadata): self.dataHandler.addMetaData(key, value) for key, value in iteritems(sections): self.dataHandler.addSection(key, value) # Update the can_write flag for MPI self.dataHandler.can_write = ( servermanager.vtkProcessModule.GetProcessModule().GetPartitionId() == 0 ) def getDataHandler(self): return self.dataHandler def getCamera(self): return self.camera def updateCamera(self, camera): update_camera(self.view, camera) def start(self, view=None): if view: # Keep track of the view self.view = view # Handle camera if any if self.cameraDescription: if self.cameraDescription["type"] == "spherical": self.camera = camera.SphericalCamera( self.dataHandler, view.CenterOfRotation, view.CameraPosition, view.CameraViewUp, self.cameraDescription["phi"], self.cameraDescription["theta"], ) elif self.cameraDescription["type"] == "cylindrical": self.camera = camera.CylindricalCamera( self.dataHandler, view.CenterOfRotation, view.CameraPosition, view.CameraViewUp, self.cameraDescription["phi"], self.cameraDescription["translation"], ) elif self.cameraDescription["type"] == "cube": self.camera = camera.CubeCamera( self.dataHandler, self.cameraDescription["viewForward"], self.cameraDescription["viewUp"], self.cameraDescription["positions"], ) elif self.cameraDescription["type"] == "cube-stereo": self.camera = camera.StereoCubeCamera( self.dataHandler, self.cameraDescription["viewForward"], self.cameraDescription["viewUp"], self.cameraDescription["positions"], self.cameraDescription["eyeSeparation"], ) # Update background color bgColor = view.Background bgColorString = "rgb(%d, %d, %d)" % tuple( int(bgColor[i] * 255) for i in range(3) ) if view.UseGradientBackground: bgColor2 = view.Background2 bgColor2String = "rgb(%d, %d, %d)" % tuple( int(bgColor2[i] * 255) for i in range(3) ) self.dataHandler.addMetaData( "backgroundColor", "linear-gradient(%s,%s)" % (bgColor2String, bgColorString), ) else: self.dataHandler.addMetaData("backgroundColor", bgColorString) # Update file patterns self.dataHandler.updateBasePattern() def stop(self): self.dataHandler.writeDataDescriptor() # ----------------------------------------------------------------------------- # Image Dataset Builder # ----------------------------------------------------------------------------- class ImageDataSetBuilder(DataSetBuilder): def __init__(self, location, imageMimeType, cameraInfo, metadata={}): DataSetBuilder.__init__(self, location, cameraInfo, metadata) imageExtenstion = "." + imageMimeType.split("/")[1] self.dataHandler.registerData( name="image", type="blob", mimeType=imageMimeType, fileName=imageExtenstion ) def writeImages(self): for cam in self.camera: update_camera(self.view, cam) simple.WriteImage(self.dataHandler.getDataAbsoluteFilePath("image")) # ----------------------------------------------------------------------------- # Data Prober Dataset Builder # ----------------------------------------------------------------------------- class DataProberDataSetBuilder(DataSetBuilder): def __init__( self, input, location, sampling_dimesions, fields_to_keep, custom_probing_bounds=None, metadata={}, ): DataSetBuilder.__init__(self, location, None, metadata) self.fieldsToWrite = fields_to_keep self.resamplerFilter = simple.ResampleToImage(Input=input) self.resamplerFilter.SamplingDimensions = sampling_dimesions if custom_probing_bounds: self.resamplerFilter.UseInputBounds = 0 self.resamplerFilter.SamplingBounds = custom_probing_bounds else: self.resamplerFilter.UseInputBounds = 1 # Register all fields self.dataHandler.addTypes("data-prober", "binary") self.DataProber = { "types": {}, "dimensions": sampling_dimesions, "ranges": {}, "spacing": [1, 1, 1], } for field in self.fieldsToWrite: self.dataHandler.registerData( name=field, type="array", rootFile=True, fileName="%s.array" % field ) def writeData(self, time=0): if not self.dataHandler.can_write: return self.resamplerFilter.UpdatePipeline(time) imageData = self.resamplerFilter.GetClientSideObject().GetOutput() self.DataProber["spacing"] = imageData.GetSpacing() arrays = imageData.GetPointData() maskArray = arrays.GetArray(vtkDataSetAttributes.GhostArrayName()) for field in self.fieldsToWrite: array = arrays.GetArray(field) if array: if array.GetNumberOfComponents() == 1: # Push NaN when no value are present instead of 0 for idx in range(maskArray.GetNumberOfTuples()): if maskArray.GetValue(idx) == 2: # Hidden point array.SetValue(idx, float("NaN")) with open( self.dataHandler.getDataAbsoluteFilePath(field), "wb" ) as f: f.write(memoryview(array)) self.expandRange(array) else: magarray = array.NewInstance() magarray.SetNumberOfTuples(array.GetNumberOfTuples()) magarray.SetName(field) for idx in range(magarray.GetNumberOfTuples()): if maskArray.GetValue(idx) == 2: # Hidden point # Push NaN when no value are present magarray.SetValue(idx, float("NaN")) else: entry = array.GetTuple(idx) mag = self.magnitude(entry) magarray.SetValue(idx, mag) with open( self.dataHandler.getDataAbsoluteFilePath(field), "wb" ) as f: f.write(memoryview(magarray)) self.expandRange(magarray) else: print("No array for", field) print(self.resamplerFilter.GetOutput()) def magnitude(self, tuple): value = 0 for item in tuple: value += item * item value = value ** 0.5 return value def expandRange(self, array): field = array.GetName() self.DataProber["types"][field] = javascriptMapping[ arrayTypesMapping[array.GetDataType()] ] if field in self.DataProber["ranges"]: dataRange = array.GetRange() if dataRange[0] < self.DataProber["ranges"][field][0]: self.DataProber["ranges"][field][0] = dataRange[0] if dataRange[1] > self.DataProber["ranges"][field][1]: self.DataProber["ranges"][field][1] = dataRange[1] else: self.DataProber["ranges"][field] = [ array.GetRange()[0], array.GetRange()[1], ] def stop(self, compress=True): # Rescale spacing to have the smaller value to be 1.0 smallerValue = min(self.DataProber["spacing"]) if smallerValue < 1.0: self.DataProber["spacing"] = tuple( i / smallerValue for i in self.DataProber["spacing"] ) # Push metadata self.dataHandler.addSection("DataProber", self.DataProber) # Write metadata DataSetBuilder.stop(self) if compress: for root, dirs, files in os.walk(self.dataHandler.getBasePath()): print("Compress", root) for name in files: if ".array" in name and ".gz" not in name: with open(os.path.join(root, name), "rb") as f_in: with gzip.open( os.path.join(root, name + ".gz"), "wb" ) as f_out: shutil.copyfileobj(f_in, f_out) os.remove(os.path.join(root, name)) # ----------------------------------------------------------------------------- # Float Image with Layer Dataset Builder # ----------------------------------------------------------------------------- class LayerDataSetBuilder(DataSetBuilder): def __init__(self, input, location, cameraInfo, imageSize=[500, 500], metadata={}): DataSetBuilder.__init__(self, location, cameraInfo, metadata) self.dataRenderer = data_writer.ScalarRenderer( isWriter=self.dataHandler.can_write ) self.view = self.dataRenderer.getView() self.view.ViewSize = imageSize self.floatImage = {"dimensions": imageSize, "layers": [], "ranges": {}} self.layerMap = {} self.input = input self.activeLayer = None self.activeField = None self.layerChanged = False self.lastTime = -1 # Update data type self.dataHandler.addTypes("float-image") def getView(self): return self.view def setActiveLayer(self, layer, field, hasMesh=False, activeSource=None): if activeSource: self.activeSource = activeSource else: self.activeSource = self.input needDataRegistration = False if layer not in self.layerMap: layerObj = { "name": layer, "array": field, "arrays": [field], "active": True, "type": "Float32Array", "hasMesh": hasMesh, } self.layerMap[layer] = layerObj self.floatImage["layers"].append(layerObj) needDataRegistration = True # Register layer lighting self.dataHandler.registerData( name="%s__light" % layer, type="array", rootFile=True, fileName="%s__light.array" % layer, categories=["%s__light" % layer], ) # Register layer mesh if hasMesh: self.dataHandler.registerData( name="%s__mesh" % layer, type="array", rootFile=True, fileName="%s__mesh.array" % layer, categories=["%s__mesh" % layer], ) elif field not in self.layerMap[layer]["arrays"]: self.layerMap[layer]["arrays"].append(field) needDataRegistration = True # Keep track of the active data if self.activeLayer != layer: self.layerChanged = True self.activeLayer = layer self.activeField = field if needDataRegistration: self.dataHandler.registerData( name="%s_%s" % (layer, field), type="array", rootFile=True, fileName="%s_%s.array" % (layer, field), categories=["%s_%s" % (layer, field)], ) def writeLayerData(self, time=0): dataRange = [0, 1] self.activeSource.UpdatePipeline(time) if self.activeField and self.activeLayer: if self.layerChanged or self.lastTime != time: self.layerChanged = False self.lastTime = time # Capture lighting information for camPos in self.getCamera(): self.view.CameraFocalPoint = camPos["focalPoint"] self.view.CameraPosition = camPos["position"] self.view.CameraViewUp = camPos["viewUp"] self.dataRenderer.writeLightArray( self.dataHandler.getDataAbsoluteFilePath( "%s__light" % self.activeLayer ), self.activeSource, ) # Capture mesh information if self.layerMap[self.activeLayer]["hasMesh"]: for camPos in self.getCamera(): self.view.CameraFocalPoint = camPos["focalPoint"] self.view.CameraPosition = camPos["position"] self.view.CameraViewUp = camPos["viewUp"] self.dataRenderer.writeMeshArray( self.dataHandler.getDataAbsoluteFilePath( "%s__mesh" % self.activeLayer ), self.activeSource, ) for camPos in self.getCamera(): self.view.CameraFocalPoint = camPos["focalPoint"] self.view.CameraPosition = camPos["position"] self.view.CameraViewUp = camPos["viewUp"] dataName = "%s_%s" % (self.activeLayer, self.activeField) dataRange = self.dataRenderer.writeArray( self.dataHandler.getDataAbsoluteFilePath(dataName), self.activeSource, self.activeField, ) if self.activeField not in self.floatImage["ranges"]: self.floatImage["ranges"][self.activeField] = [ dataRange[0], dataRange[1], ] else: # Expand the ranges if dataRange[0] < self.floatImage["ranges"][self.activeField][0]: self.floatImage["ranges"][self.activeField][0] = dataRange[0] if dataRange[1] > self.floatImage["ranges"][self.activeField][1]: self.floatImage["ranges"][self.activeField][1] = dataRange[1] def start(self): DataSetBuilder.start(self, self.view) def stop(self, compress=True): if not self.dataHandler.can_write: return # Push metadata self.dataHandler.addSection("FloatImage", self.floatImage) # Write metadata DataSetBuilder.stop(self) if compress: for root, dirs, files in os.walk(self.dataHandler.getBasePath()): print("Compress", root) for name in files: if ".array" in name and ".gz" not in name: with open(os.path.join(root, name), "rb") as f_in: with gzip.open( os.path.join(root, name + ".gz"), "wb" ) as f_out: shutil.copyfileobj(f_in, f_out) os.remove(os.path.join(root, name)) # ----------------------------------------------------------------------------- # Composite Dataset Builder # ----------------------------------------------------------------------------- class CompositeDataSetBuilder(DataSetBuilder): def __init__( self, location, sceneConfig, cameraInfo, metadata={}, sections={}, view=None ): DataSetBuilder.__init__(self, location, cameraInfo, metadata, sections) if view: self.view = view else: self.view = simple.CreateView("RenderView") self.view.ViewSize = sceneConfig["size"] self.view.CenterAxesVisibility = 0 self.view.OrientationAxesVisibility = 0 self.view.UpdatePropertyInformation() if "camera" in sceneConfig and "CameraFocalPoint" in sceneConfig["camera"]: self.view.CenterOfRotation = sceneConfig["camera"]["CameraFocalPoint"] # Initialize camera for key, value in iteritems(sceneConfig["camera"]): self.view.GetProperty(key).SetData(value) # Create a representation for all scene sources self.config = sceneConfig self.representations = [] for data in self.config["scene"]: rep = simple.Show(data["source"], self.view) self.representations.append(rep) if "representation" in data: for key in data["representation"]: rep.GetProperty(key).SetData(data["representation"][key]) # Add directory path self.dataHandler.registerData( name="directory", rootFile=True, fileName="file.txt", categories=["trash"] ) def start(self): DataSetBuilder.start(self, self.view) def stop(self, compress=True, clean=True): DataSetBuilder.stop(self) if not self.dataHandler.can_write: return # Make the config serializable for item in self.config["scene"]: del item["source"] # Write the scene to disk with open( os.path.join(self.dataHandler.getBasePath(), "config.json"), "w" ) as f: f.write(json.dumps(self.config)) dataConverter = data_converter.ConvertCompositeDataToSortedStack( self.dataHandler.getBasePath() ) dataConverter.convert() # Remove tmp files os.remove(os.path.join(self.dataHandler.getBasePath(), "index.json")) os.remove(os.path.join(self.dataHandler.getBasePath(), "config.json")) # Composite pipeline meta description compositePipeline = { "default_pipeline": "", "layers": [], "fields": {}, "layer_fields": {}, "pipeline": [], } rootItems = {} fieldNameMapping = {} # Clean scene in config and gather ranges dataRanges = {} layerIdx = 0 for layer in self.config["scene"]: # Create group node if any if "parent" in layer and layer["parent"] not in rootItems: rootItems[layer["parent"]] = { "name": layer["parent"], "ids": [], "children": [], } compositePipeline["pipeline"].append(rootItems[layer["parent"]]) # Create layer entry layerCode = encode_codes[layerIdx] layerItem = {"name": layer["name"], "ids": [layerCode]} compositePipeline["layers"].append(layerCode) compositePipeline["layer_fields"][layerCode] = [] compositePipeline["default_pipeline"] += layerCode # Register layer entry in pipeline if "parent" in layer: rootItems[layer["parent"]]["children"].append(layerItem) rootItems[layer["parent"]]["ids"].append(layerCode) else: compositePipeline["pipeline"].append(layerItem) # Handle color / field colorByList = [] for color in layer["colors"]: # Find color code if color not in fieldNameMapping: colorCode = encode_codes[len(fieldNameMapping)] fieldNameMapping[color] = colorCode compositePipeline["fields"][colorCode] = color else: colorCode = fieldNameMapping[color] # Register color code compositePipeline["layer_fields"][layerCode].append(colorCode) if len(colorByList) == 0: compositePipeline["default_pipeline"] += colorCode values = None if "constant" in layer["colors"][color]: value = layer["colors"][color]["constant"] values = [value, value] colorByList.append({"name": color, "type": "const", "value": value}) elif "range" in layer["colors"][color]: values = layer["colors"][color]["range"] colorByList.append({"name": color, "type": "field"}) if values: if color not in dataRanges: dataRanges[color] = values else: dataRanges[color][0] = min( dataRanges[color][0], values[0], values[1] ) dataRanges[color][1] = max( dataRanges[color][1], values[0], values[1] ) layer["colorBy"] = colorByList del layer["colors"] layerIdx += 1 sortedCompositeSection = { "dimensions": self.config["size"], "pipeline": self.config["scene"], "ranges": dataRanges, "layers": len(self.config["scene"]), "light": self.config["light"], } self.dataHandler.addSection("SortedComposite", sortedCompositeSection) self.dataHandler.addSection("CompositePipeline", compositePipeline) self.dataHandler.addTypes("sorted-composite", "multi-color-by") self.dataHandler.removeData("directory") for dataToRegister in dataConverter.listData(): self.dataHandler.registerData(**dataToRegister) self.dataHandler.writeDataDescriptor() if clean: for root, dirs, files in os.walk(self.dataHandler.getBasePath()): print("Clean", root) for name in files: if name in ["camera.json"]: os.remove(os.path.join(root, name)) if compress: for root, dirs, files in os.walk(self.dataHandler.getBasePath()): print("Compress", root) for name in files: if (".float32" in name or ".uint8" in name) and ".gz" not in name: with open(os.path.join(root, name), "rb") as f_in: with gzip.open( os.path.join(root, name + ".gz"), "wb" ) as f_out: shutil.copyfileobj(f_in, f_out) os.remove(os.path.join(root, name)) def writeData(self): composite_size = len(self.representations) self.view.UpdatePropertyInformation() self.view.Background = [0, 0, 0] imageSize = self.view.ViewSize[0] * self.view.ViewSize[1] # Generate the heavy data for camPos in self.getCamera(): self.view.CameraFocalPoint = camPos["focalPoint"] self.view.CameraPosition = camPos["position"] self.view.CameraViewUp = camPos["viewUp"] # Show all representations for compositeIdx in range(composite_size): rep = self.representations[compositeIdx] rep.Visibility = 1 # Fix camera bounds self.view.LockBounds = 0 simple.Render(self.view) self.view.LockBounds = 1 # Update destination directory dest_path = os.path.dirname( self.dataHandler.getDataAbsoluteFilePath("directory") ) # Write camera information if self.dataHandler.can_write: with open(os.path.join(dest_path, "camera.json"), "w") as f: f.write(json.dumps(camPos)) # Hide all representations for compositeIdx in range(composite_size): rep = self.representations[compositeIdx] rep.Visibility = 0 # Show only active Representation # Extract images for each fields for compositeIdx in range(composite_size): rep = self.representations[compositeIdx] if compositeIdx > 0: self.representations[compositeIdx - 1].Visibility = 0 rep.Visibility = 1 # capture Z simple.Render() zBuffer = self.view.CaptureDepthBuffer() with open( os.path.join(dest_path, "depth_%d.float32" % compositeIdx), "wb" ) as f: f.write(memoryview(zBuffer)) # Prevent color interference rep.DiffuseColor = [1, 1, 1] # Handle light for lightType in self.config["light"]: if lightType == "intensity": rep.AmbientColor = [1, 1, 1] rep.SpecularColor = [1, 1, 1] self.view.StartCaptureLuminance() image = self.view.CaptureWindow(1) imagescalars = image.GetPointData().GetScalars() self.view.StopCaptureLuminance() # Extract specular information specularOffset = 1 # [diffuse, specular, ?] imageSize = imagescalars.GetNumberOfTuples() specularComponent = vtkUnsignedCharArray() specularComponent.SetNumberOfComponents(1) specularComponent.SetNumberOfTuples(imageSize) for idx in range(imageSize): specularComponent.SetValue( idx, imagescalars.GetValue(idx * 3 + specularOffset) ) with open( os.path.join( dest_path, "intensity_%d.uint8" % compositeIdx ), "wb", ) as f: f.write(memoryview(specularComponent)) # Free memory image.UnRegister(None) if lightType == "normal": if rep.Representation in [ "Point Gaussian", "Points", "Outline", "Wireframe", ]: uniqNormal = [ (camPos["position"][i] - camPos["focalPoint"][i]) for i in range(3) ] tmpNormalArray = vtkFloatArray() tmpNormalArray.SetNumberOfComponents(1) tmpNormalArray.SetNumberOfTuples(imageSize) for comp in range(3): tmpNormalArray.FillComponent(0, uniqNormal[comp]) with open( os.path.join( dest_path, "normal_%d_%d.float32" % (compositeIdx, comp), ), "wb", ) as f: f.write(memoryview(tmpNormalArray)) else: for comp in range(3): # Configure view to handle POINT_DATA / CELL_DATA self.view.DrawCells = 0 self.view.ArrayNameToDraw = "Normals" self.view.ArrayComponentToDraw = comp self.view.ScalarRange = [-1.0, 1.0] self.view.StartCaptureValues() image = self.view.CaptureWindow(1) imagescalars = image.GetPointData().GetScalars() self.view.StopCaptureValues() # Convert RGB => Float => Write floatArray = data_converter.convertRGBArrayToFloatArray( imagescalars, [-1.0, 1.0] ) with open( os.path.join( dest_path, "normal_%d_%d.float32" % (compositeIdx, comp), ), "wb", ) as f: f.write(memoryview(floatArray)) # Free memory image.UnRegister(None) # Handle color by for fieldName, fieldConfig in iteritems( self.config["scene"][compositeIdx]["colors"] ): if "constant" in fieldConfig: # Skip nothing to render continue # Configure view to handle POINT_DATA / CELL_DATA if fieldConfig["location"] == "POINT_DATA": self.view.DrawCells = 0 else: self.view.DrawCells = 1 self.view.ArrayNameToDraw = fieldName self.view.ArrayComponentToDraw = 0 self.view.ScalarRange = fieldConfig["range"] self.view.StartCaptureValues() image = self.view.CaptureWindow(1) imagescalars = image.GetPointData().GetScalars() self.view.StopCaptureValues() floatArray = data_converter.convertRGBArrayToFloatArray( imagescalars, fieldConfig["range"] ) with open( os.path.join( dest_path, "%d_%s.float32" % (compositeIdx, fieldName) ), "wb", ) as f: f.write(memoryview(floatArray)) self.dataHandler.registerData( name="%d_%s" % (compositeIdx, fieldName), fileName="/%d_%s.float32" % (compositeIdx, fieldName), type="array", categories=["%d_%s" % (compositeIdx, fieldName)], ) # Free memory image.UnRegister(None) # ----------------------------------------------------------------------------- # VTKGeometryDataSetBuilder Dataset Builder # ----------------------------------------------------------------------------- def writeCellArray(dataHandler, currentData, cellName, inputCellArray): nbValues = inputCellArray.GetNumberOfTuples() if nbValues == 0: return outputCells = vtkTypeUInt32Array() outputCells.SetNumberOfTuples(nbValues) for valueIdx in range(nbValues): outputCells.SetValue(valueIdx, inputCellArray.GetValue(valueIdx)) iBuffer = memoryview(outputCells) iMd5 = hashlib.md5(iBuffer).hexdigest() iPath = os.path.join(dataHandler.getBasePath(), "data", "%s.Uint32Array" % iMd5) currentData["polys"] = "data/%s.Uint32Array" % iMd5 with open(iPath, "wb") as f: f.write(iBuffer) class VTKGeometryDataSetBuilder(DataSetBuilder): def __init__(self, location, sceneConfig, metadata={}, sections={}): DataSetBuilder.__init__(self, location, None, metadata, sections) # Update data type self.dataHandler.addTypes("vtk-geometry") # Create a representation for all scene sources self.config = sceneConfig # Processing pipeline self.surfaceExtract = None # Add directory path self.dataHandler.registerData( priority=0, name="scene", rootFile=True, fileName="scene.json", type="json" ) # Create directory containers dataPath = os.path.join(location, "data") for p in [dataPath]: if not os.path.exists(p): os.makedirs(p) # Create metadata structure colorToCodeMap = {} parentNodes = {} pipelineMeta = {"layers": [], "pipeline": [], "layer_fields": {}, "fields": {}} geometryMeta = { "ranges": {}, "layer_map": {}, } self.ranges = geometryMeta["ranges"] for item in sceneConfig["scene"]: # Handle layer layerCode = encode_codes[len(pipelineMeta["layers"])] pipelineMeta["layers"].append(layerCode) geometryMeta["layer_map"][layerCode] = item["name"] # Handle colors pipelineMeta["layer_fields"][layerCode] = [] for fieldName in item["colors"]: colorCode = None if fieldName in colorToCodeMap: colorCode = colorToCodeMap[fieldName] else: colorCode = encode_codes[len(colorToCodeMap)] colorToCodeMap[fieldName] = colorCode geometryMeta["ranges"][fieldName] = [ 1, -1, ] # FIXME we don't know the range pipelineMeta["layer_fields"][layerCode].append(colorCode) pipelineMeta["fields"][colorCode] = fieldName # Handle pipeline if "parent" in item: # Need to handle hierarchy if item["parent"] in parentNodes: # Fill children rootNode = parentNodes[item["parent"]] rootNode["ids"].append(layerCode) rootNode["children"].append( {"name": item["name"], "ids": [layerCode]} ) else: # Create root + register rootNode = { "name": item["parent"], "ids": [layerCode], "children": [{"name": item["name"], "ids": [layerCode]}], } parentNodes[item["parent"]] = rootNode pipelineMeta["pipeline"].append(rootNode) else: # Add item info as a new pipeline node pipelineMeta["pipeline"].append( {"name": item["name"], "ids": [layerCode]} ) # Register metadata to be written in index.json self.dataHandler.addSection("Geometry", geometryMeta) self.dataHandler.addSection("CompositePipeline", pipelineMeta) def writeData(self, time=0): if not self.dataHandler.can_write: return currentScene = [] for data in self.config["scene"]: currentData = {"name": data["name"], "fields": {}, "cells": {}} currentScene.append(currentData) if self.surfaceExtract: self.merge.Input = data["source"] else: self.merge = simple.MergeBlocks(Input=data["source"], MergePoints=0) self.surfaceExtract = simple.ExtractSurface(Input=self.merge) # Extract surface self.surfaceExtract.UpdatePipeline(time) ds = self.surfaceExtract.SMProxy.GetClientSideObject().GetOutputDataObject( 0 ) originalDS = ( data["source"].SMProxy.GetClientSideObject().GetOutputDataObject(0) ) originalPoints = ds.GetPoints() # Points points = vtkFloatArray() nbPoints = originalPoints.GetNumberOfPoints() points.SetNumberOfComponents(3) points.SetNumberOfTuples(nbPoints) for idx in range(nbPoints): coord = originalPoints.GetPoint(idx) points.SetTuple3(idx, coord[0], coord[1], coord[2]) pBuffer = memoryview(points) pMd5 = hashlib.md5(pBuffer).hexdigest() pPath = os.path.join( self.dataHandler.getBasePath(), "data", "%s.Float32Array" % pMd5 ) currentData["points"] = "data/%s.Float32Array" % pMd5 with open(pPath, "wb") as f: f.write(pBuffer) # Handle cells writeCellArray( self.dataHandler, currentData["cells"], "verts", ds.GetVerts().GetData() ) writeCellArray( self.dataHandler, currentData["cells"], "lines", ds.GetLines().GetData() ) writeCellArray( self.dataHandler, currentData["cells"], "polys", ds.GetPolys().GetData() ) writeCellArray( self.dataHandler, currentData["cells"], "strips", ds.GetStrips().GetData(), ) # Fields for fieldName, fieldInfo in iteritems(data["colors"]): array = None if "constant" in fieldInfo: currentData["fields"][fieldName] = fieldInfo continue elif "POINT_DATA" in fieldInfo["location"]: array = ds.GetPointData().GetArray(fieldName) elif "CELL_DATA" in fieldInfo["location"]: array = ds.GetCellData().GetArray(fieldName) jsType = javascriptMapping[arrayTypesMapping[array.GetDataType()]] arrayRange = array.GetRange(-1) tupleSize = array.GetNumberOfComponents() arraySize = array.GetNumberOfTuples() if tupleSize == 1: outputField = array else: # compute magnitude outputField = array.NewInstance() outputField.SetNumberOfTuples(arraySize) tupleIdxs = range(tupleSize) for i in range(arraySize): magnitude = 0 for j in tupleIdxs: magnitude += math.pow(array.GetValue(i * tupleSize + j), 2) outputField.SetValue(i, math.sqrt(magnitude)) fBuffer = memoryview(outputField) fMd5 = hashlib.md5(fBuffer).hexdigest() fPath = os.path.join( self.dataHandler.getBasePath(), "data", "%s_%s.%s" % (fieldName, fMd5, jsType), ) with open(fPath, "wb") as f: f.write(fBuffer) currentRange = self.ranges[fieldName] if currentRange[1] < currentRange[0]: currentRange[0] = arrayRange[0] currentRange[1] = arrayRange[1] else: currentRange[0] = ( arrayRange[0] if arrayRange[0] < currentRange[0] else currentRange[0] ) currentRange[1] = ( arrayRange[1] if arrayRange[1] > currentRange[1] else currentRange[1] ) currentData["fields"][fieldName] = { "array": "data/%s_%s.%s" % (fieldName, fMd5, jsType), "location": fieldInfo["location"], "range": outputField.GetRange(), } # Write scene with open(self.dataHandler.getDataAbsoluteFilePath("scene"), "w") as f: f.write(json.dumps(currentScene, indent=4)) def stop(self, compress=True, clean=True): if not self.dataHandler.can_write: return DataSetBuilder.stop(self) if compress: for dirName in ["fields", "index", "points"]: for root, dirs, files in os.walk( os.path.join(self.dataHandler.getBasePath(), dirName) ): print("Compress", root) for name in files: if "Array" in name and ".gz" not in name: with open(os.path.join(root, name), "rb") as f_in: with gzip.open( os.path.join(root, name + ".gz"), "wb" ) as f_out: shutil.copyfileobj(f_in, f_out) os.remove(os.path.join(root, name)) # ----------------------------------------------------------------------------- # GeometryDataSetBuilder Dataset Builder # ----------------------------------------------------------------------------- class GeometryDataSetBuilder(DataSetBuilder): def __init__(self, location, sceneConfig, metadata={}, sections={}): DataSetBuilder.__init__(self, location, None, metadata, sections) # Update data type self.dataHandler.addTypes("geometry") # Create a representation for all scene sources self.config = sceneConfig # Processing pipeline self.surfaceExtract = None # Add directory path self.dataHandler.registerData( priority=0, name="scene", rootFile=True, fileName="scene.json", type="json" ) # Create directory containers pointsPath = os.path.join(location, "points") polyPath = os.path.join(location, "index") colorPath = os.path.join(location, "fields") for p in [pointsPath, polyPath, colorPath]: if not os.path.exists(p): os.makedirs(p) # Create metadata structure colorToCodeMap = {} parentNodes = {} pipelineMeta = {"layers": [], "pipeline": [], "layer_fields": {}, "fields": {}} geometryMeta = {"ranges": {}, "layer_map": {}, "object_size": {}} self.objSize = geometryMeta["object_size"] for item in sceneConfig["scene"]: # Handle layer layerCode = encode_codes[len(pipelineMeta["layers"])] pipelineMeta["layers"].append(layerCode) geometryMeta["layer_map"][layerCode] = item["name"] geometryMeta["object_size"][item["name"]] = {"points": 0, "index": 0} # Handle colors pipelineMeta["layer_fields"][layerCode] = [] for fieldName in item["colors"]: colorCode = None if fieldName in colorToCodeMap: colorCode = colorToCodeMap[fieldName] else: colorCode = encode_codes[len(colorToCodeMap)] colorToCodeMap[fieldName] = colorCode geometryMeta["ranges"][fieldName] = [ 0, 1, ] # FIXME we don't know the range pipelineMeta["layer_fields"][layerCode].append(colorCode) pipelineMeta["fields"][colorCode] = fieldName # Handle pipeline if "parent" in item: # Need to handle hierarchy if item["parent"] in parentNodes: # Fill children rootNode = parentNodes[item["parent"]] rootNode["ids"].append(layerCode) rootNode["children"].append( {"name": item["name"], "ids": [layerCode]} ) else: # Create root + register rootNode = { "name": item["parent"], "ids": [layerCode], "children": [{"name": item["name"], "ids": [layerCode]}], } parentNodes[item["parent"]] = rootNode pipelineMeta["pipeline"].append(rootNode) else: # Add item info as a new pipeline node pipelineMeta["pipeline"].append( {"name": item["name"], "ids": [layerCode]} ) # Register metadata to be written in index.json self.dataHandler.addSection("Geometry", geometryMeta) self.dataHandler.addSection("CompositePipeline", pipelineMeta) def writeData(self, time=0): if not self.dataHandler.can_write: return currentScene = [] for data in self.config["scene"]: currentData = {"name": data["name"], "fields": {}} currentScene.append(currentData) if self.surfaceExtract: self.merge.Input = data["source"] else: self.merge = simple.MergeBlocks(Input=data["source"], MergePoints=0) self.surfaceExtract = simple.ExtractSurface(Input=self.merge) # Extract surface self.surfaceExtract.UpdatePipeline(time) ds = self.surfaceExtract.SMProxy.GetClientSideObject().GetOutputDataObject( 0 ) originalDS = ( data["source"].SMProxy.GetClientSideObject().GetOutputDataObject(0) ) originalPoints = ds.GetPoints() # Points points = vtkFloatArray() nbPoints = originalPoints.GetNumberOfPoints() points.SetNumberOfComponents(3) points.SetNumberOfTuples(nbPoints) for idx in range(nbPoints): coord = originalPoints.GetPoint(idx) points.SetTuple3(idx, coord[0], coord[1], coord[2]) pBuffer = memoryview(points) pMd5 = hashlib.md5(pBuffer).hexdigest() pPath = os.path.join( self.dataHandler.getBasePath(), "points", "%s.Float32Array" % pMd5 ) currentData["points"] = "points/%s.Float32Array" % pMd5 with open(pPath, "wb") as f: f.write(pBuffer) # Polys poly = ds.GetPolys() nbCells = poly.GetNumberOfCells() cellLocation = 0 idList = vtkIdList() topo = vtkTypeUInt32Array() topo.Allocate(poly.GetData().GetNumberOfTuples()) for cellIdx in range(nbCells): poly.GetCell(cellLocation, idList) cellSize = idList.GetNumberOfIds() cellLocation += cellSize + 1 if cellSize == 3: topo.InsertNextValue(idList.GetId(0)) topo.InsertNextValue(idList.GetId(1)) topo.InsertNextValue(idList.GetId(2)) elif cellSize == 4: topo.InsertNextValue(idList.GetId(0)) topo.InsertNextValue(idList.GetId(1)) topo.InsertNextValue(idList.GetId(3)) topo.InsertNextValue(idList.GetId(1)) topo.InsertNextValue(idList.GetId(2)) topo.InsertNextValue(idList.GetId(3)) else: print("Cell size of", cellSize, "not supported") iBuffer = memoryview(topo) iMd5 = hashlib.md5(iBuffer).hexdigest() iPath = os.path.join( self.dataHandler.getBasePath(), "index", "%s.Uint32Array" % iMd5 ) currentData["index"] = "index/%s.Uint32Array" % iMd5 with open(iPath, "wb") as f: f.write(iBuffer) # Grow object side self.objSize[data["name"]]["points"] = max( self.objSize[data["name"]]["points"], nbPoints ) self.objSize[data["name"]]["index"] = max( self.objSize[data["name"]]["index"], topo.GetNumberOfTuples() ) # Colors / FIXME for fieldName, fieldInfo in iteritems(data["colors"]): array = ds.GetPointData().GetArray(fieldName) tupleSize = array.GetNumberOfComponents() arraySize = array.GetNumberOfTuples() outputField = vtkFloatArray() outputField.SetNumberOfTuples(arraySize) if tupleSize == 1: for i in range(arraySize): outputField.SetValue(i, array.GetValue(i)) else: # compute magnitude tupleIdxs = range(tupleSize) for i in range(arraySize): magnitude = 0 for j in tupleIdxs: magnitude += math.pow(array.GetValue(i * tupleSize + j), 2) outputField.SetValue(i, math.sqrt(magnitude)) fBuffer = memoryview(outputField) fMd5 = hashlib.md5(fBuffer).hexdigest() fPath = os.path.join( self.dataHandler.getBasePath(), "fields", "%s_%s.Float32Array" % (fieldName, fMd5), ) with open(fPath, "wb") as f: f.write(fBuffer) currentData["fields"][fieldName] = "fields/%s_%s.Float32Array" % ( fieldName, fMd5, ) # Write scene with open(self.dataHandler.getDataAbsoluteFilePath("scene"), "w") as f: f.write(json.dumps(currentScene, indent=2)) def stop(self, compress=True, clean=True): if not self.dataHandler.can_write: return DataSetBuilder.stop(self) if compress: for dirName in ["fields", "index", "points"]: for root, dirs, files in os.walk( os.path.join(self.dataHandler.getBasePath(), dirName) ): print("Compress", root) for name in files: if "Array" in name and ".gz" not in name: with open(os.path.join(root, name), "rb") as f_in: with gzip.open( os.path.join(root, name + ".gz"), "wb" ) as f_out: shutil.copyfileobj(f_in, f_out) os.remove(os.path.join(root, name))