import json, os, gzip, shutil from vtkmodules.vtkRenderingCore import vtkWindowToImageFilter from vtkmodules.vtkIOImage import vtkPNGReader, vtkPNGWriter, vtkJPEGWriter from vtkmodules.vtkCommonDataModel import vtkImageData from vtkmodules.vtkCommonCore import vtkUnsignedCharArray from vtkmodules.vtkFiltersParallel import vtkPResampleFilter from vtkmodules.web import iteritems, getJSArrayType from vtkmodules.web.camera import ( update_camera, create_spherical_camera, create_cylindrical_camera, ) from vtkmodules.web.query_data_model import DataHandler # Global helper variables encode_codes = "ABCDEFGHIJKLMNOPQRSTUVWXYZ" # ----------------------------------------------------------------------------- # Capture image from render window # ----------------------------------------------------------------------------- class CaptureRenderWindow(object): def __init__(self, magnification=1): self.windowToImage = vtkWindowToImageFilter() self.windowToImage.SetScale(magnification) self.windowToImage.SetInputBufferTypeToRGB() self.windowToImage.ReadFrontBufferOn() self.writer = None def SetRenderWindow(self, renderWindow): self.windowToImage.SetInput(renderWindow) def SetFormat(self, mimeType): if mimeType == "image/png": self.writer = vtkPNGWriter() self.writer.SetInputConnection(self.windowToImage.GetOutputPort()) elif mimeType == "image/jpg": self.writer = vtkJPEGWriter() self.writer.SetInputConnection(self.windowToImage.GetOutputPort()) def writeImage(self, path): if self.writer: self.windowToImage.Modified() self.windowToImage.Update() self.writer.SetFileName(path) self.writer.Write() # ----------------------------------------------------------------------------- # 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 self.imageCapture = CaptureRenderWindow() for key, value in iteritems(metadata): self.dataHandler.addMetaData(key, value) for key, value in iteritems(sections): self.dataHandler.addSection(key, value) def getDataHandler(self): return self.dataHandler def getCamera(self): return self.camera def updateCamera(self, camera): update_camera(self.renderer, camera) self.renderWindow.Render() def start(self, renderWindow=None, renderer=None): if renderWindow: # Keep track of renderWindow and renderer self.renderWindow = renderWindow self.renderer = renderer # Initialize image capture self.imageCapture.SetRenderWindow(renderWindow) # Handle camera if any if self.cameraDescription: if self.cameraDescription["type"] == "spherical": self.camera = create_spherical_camera( renderer, self.dataHandler, self.cameraDescription["phi"], self.cameraDescription["theta"], ) elif self.cameraDescription["type"] == "cylindrical": self.camera = create_cylindrical_camera( renderer, self.dataHandler, self.cameraDescription["phi"], self.cameraDescription["translation"], ) # Update background color bgColor = renderer.GetBackground() bgColorString = "rgb(%d, %d, %d)" % tuple( int(bgColor[i] * 255) for i in range(3) ) 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={}, sections={}): DataSetBuilder.__init__(self, location, cameraInfo, metadata, sections) imageExtenstion = "." + imageMimeType.split("/")[1] self.dataHandler.registerData( name="image", type="blob", mimeType=imageMimeType, fileName=imageExtenstion ) self.imageCapture.SetFormat(imageMimeType) def writeImage(self): self.imageCapture.writeImage(self.dataHandler.getDataAbsoluteFilePath("image")) def writeImages(self): for cam in self.camera: update_camera(self.renderer, cam) self.renderWindow.Render() self.imageCapture.writeImage( self.dataHandler.getDataAbsoluteFilePath("image") ) # ----------------------------------------------------------------------------- # Volume Composite Dataset Builder # ----------------------------------------------------------------------------- class VolumeCompositeDataSetBuilder(DataSetBuilder): def __init__(self, location, imageMimeType, cameraInfo, metadata={}, sections={}): DataSetBuilder.__init__(self, location, cameraInfo, metadata, sections) self.dataHandler.addTypes("volume-composite", "rgba+depth") self.imageMimeType = imageMimeType self.imageExtenstion = "." + imageMimeType.split("/")[1] if imageMimeType == "image/png": self.imageWriter = vtkPNGWriter() if imageMimeType == "image/jpg": self.imageWriter = vtkJPEGWriter() self.imageDataColor = vtkImageData() self.imageWriter.SetInputData(self.imageDataColor) self.imageDataDepth = vtkImageData() self.depthToWrite = None self.layerInfo = {} self.colorByMapping = {} self.compositePipeline = { "layers": [], "dimensions": [], "fields": {}, "layer_fields": {}, "pipeline": [], } self.activeDepthKey = "" self.activeRGBKey = "" self.nodeWithChildren = {} def _getColorCode(self, colorBy): if colorBy in self.colorByMapping: # The color code exist return self.colorByMapping[colorBy] else: # No color code assigned yet colorCode = encode_codes[len(self.colorByMapping)] # Assign color code self.colorByMapping[colorBy] = colorCode # Register color code with color by value self.compositePipeline["fields"][colorCode] = colorBy # Return the color code return colorCode def _getLayerCode(self, parent, layerName): if layerName in self.layerInfo: # Layer already exist return (self.layerInfo[layerName]["code"], False) else: layerCode = encode_codes[len(self.layerInfo)] self.layerInfo[layerName] = { "code": layerCode, "name": layerName, "parent": parent, } self.compositePipeline["layers"].append(layerCode) self.compositePipeline["layer_fields"][layerCode] = [] # Let's register it in the pipeline if parent: if parent not in self.nodeWithChildren: # Need to create parent rootNode = {"name": parent, "ids": [], "children": []} self.nodeWithChildren[parent] = rootNode self.compositePipeline["pipeline"].append(rootNode) # Add node to its parent self.nodeWithChildren[parent]["children"].append( {"name": layerName, "ids": [layerCode]} ) self.nodeWithChildren[parent]["ids"].append(layerCode) else: self.compositePipeline["pipeline"].append( {"name": layerName, "ids": [layerCode]} ) return (layerCode, True) def _needToRegisterColor(self, layerCode, colorCode): if colorCode in self.compositePipeline["layer_fields"][layerCode]: return False else: self.compositePipeline["layer_fields"][layerCode].append(colorCode) return True def activateLayer(self, parent, name, colorBy): layerCode, needToRegisterDepth = self._getLayerCode(parent, name) colorCode = self._getColorCode(colorBy) needToRegisterColor = self._needToRegisterColor(layerCode, colorCode) # Update active keys self.activeDepthKey = "%s_depth" % layerCode self.activeRGBKey = "%s%s_rgb" % (layerCode, colorCode) # Need to register data if needToRegisterDepth: self.dataHandler.registerData( name=self.activeDepthKey, type="array", fileName="/%s_depth.uint8" % layerCode, categories=[layerCode], ) if needToRegisterColor: self.dataHandler.registerData( name=self.activeRGBKey, type="blob", fileName="/%s%s_rgb%s" % (layerCode, colorCode, self.imageExtenstion), categories=["%s%s" % (layerCode, colorCode)], mimeType=self.imageMimeType, ) def writeData(self, mapper): width = self.renderWindow.GetSize()[0] height = self.renderWindow.GetSize()[1] if not self.depthToWrite: self.depthToWrite = bytearray(width * height) for cam in self.camera: self.updateCamera(cam) imagePath = self.dataHandler.getDataAbsoluteFilePath(self.activeRGBKey) depthPath = self.dataHandler.getDataAbsoluteFilePath(self.activeDepthKey) # ----------------------------------------------------------------- # Write Image # ----------------------------------------------------------------- mapper.GetColorImage(self.imageDataColor) self.imageWriter.SetFileName(imagePath) self.imageWriter.Write() # ----------------------------------------------------------------- # Write Depth # ----------------------------------------------------------------- mapper.GetDepthImage(self.imageDataDepth) inputArray = self.imageDataDepth.GetPointData().GetArray(0) size = inputArray.GetNumberOfTuples() for idx in range(size): self.depthToWrite[idx] = int(inputArray.GetValue(idx)) with open(depthPath, "wb") as f: f.write(self.depthToWrite) def start(self, renderWindow, renderer): DataSetBuilder.start(self, renderWindow, renderer) self.camera.updatePriority([2, 1]) def stop(self, compress=True): # Push metadata self.compositePipeline["dimensions"] = self.renderWindow.GetSize() self.compositePipeline["default_pipeline"] = ( "A".join(self.compositePipeline["layers"]) + "A" ) self.dataHandler.addSection("CompositePipeline", self.compositePipeline) # 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 ".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)) # ----------------------------------------------------------------------------- # Data Prober Dataset Builder # ----------------------------------------------------------------------------- class DataProberDataSetBuilder(DataSetBuilder): def __init__( self, location, sampling_dimesions, fields_to_keep, custom_probing_bounds=None, metadata={}, ): DataSetBuilder.__init__(self, location, None, metadata) self.fieldsToWrite = fields_to_keep self.resamplerFilter = vtkPResampleFilter() self.resamplerFilter.SetSamplingDimension(sampling_dimesions) if custom_probing_bounds: self.resamplerFilter.SetUseInputBounds(0) self.resamplerFilter.SetCustomSamplingBounds(custom_probing_bounds) else: self.resamplerFilter.SetUseInputBounds(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", fileName="/%s.array" % field ) def setDataToProbe(self, dataset): self.resamplerFilter.SetInputData(dataset) def setSourceToProbe(self, source): self.resamplerFilter.SetInputConnection(source.GetOutputPort()) def writeData(self): self.resamplerFilter.Update() arrays = self.resamplerFilter.GetOutput().GetPointData() for field in self.fieldsToWrite: array = arrays.GetArray(field) if array: b = memoryview(array) with open(self.dataHandler.getDataAbsoluteFilePath(field), "wb") as f: f.write(b) self.DataProber["types"][field] = getJSArrayType(array) 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], ] else: print("No array for", field) print(self.resamplerFilter.GetOutput()) def stop(self, compress=True): # 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)) # ----------------------------------------------------------------------------- # Sorted Composite Dataset Builder # ----------------------------------------------------------------------------- class ConvertVolumeStackToSortedStack(object): def __init__(self, width, height): self.width = width self.height = height self.layers = 0 def convert(self, directory): imagePaths = {} depthPaths = {} layerNames = [] for fileName in os.listdir(directory): if "_rgb" in fileName or "_depth" in fileName: fileId = fileName.split("_")[0][0] if "_rgb" in fileName: imagePaths[fileId] = os.path.join(directory, fileName) else: layerNames.append(fileId) depthPaths[fileId] = os.path.join(directory, fileName) layerNames.sort() if len(layerNames) == 0: return # Load data in Memory depthArrays = [] imageReader = vtkPNGReader() numberOfValues = self.width * self.height * len(layerNames) imageSize = self.width * self.height self.layers = len(layerNames) # Write all images as single memoryview opacity = vtkUnsignedCharArray() opacity.SetNumberOfComponents(1) opacity.SetNumberOfTuples(numberOfValues) intensity = vtkUnsignedCharArray() intensity.SetNumberOfComponents(1) intensity.SetNumberOfTuples(numberOfValues) for layer in range(self.layers): imageReader.SetFileName(imagePaths[layerNames[layer]]) imageReader.Update() rgbaArray = imageReader.GetOutput().GetPointData().GetArray(0) for idx in range(imageSize): intensity.SetValue( (layer * imageSize) + idx, rgbaArray.GetValue(idx * 4) ) opacity.SetValue( (layer * imageSize) + idx, rgbaArray.GetValue(idx * 4 + 3) ) with open(depthPaths[layerNames[layer]], "rb") as depthFile: depthArrays.append(depthFile.read()) # Apply pixel sorting destOrder = vtkUnsignedCharArray() destOrder.SetNumberOfComponents(1) destOrder.SetNumberOfTuples(numberOfValues) opacityOrder = vtkUnsignedCharArray() opacityOrder.SetNumberOfComponents(1) opacityOrder.SetNumberOfTuples(numberOfValues) intensityOrder = vtkUnsignedCharArray() intensityOrder.SetNumberOfComponents(1) intensityOrder.SetNumberOfTuples(numberOfValues) for pixelIdx in range(imageSize): depthStack = [] for depthArray in depthArrays: depthStack.append((depthArray[pixelIdx], len(depthStack))) depthStack.sort(key=lambda tup: tup[0]) for destLayerIdx in range(len(depthStack)): sourceLayerIdx = depthStack[destLayerIdx][1] # Copy Idx destOrder.SetValue( (imageSize * destLayerIdx) + pixelIdx, sourceLayerIdx ) opacityOrder.SetValue( (imageSize * destLayerIdx) + pixelIdx, opacity.GetValue((imageSize * sourceLayerIdx) + pixelIdx), ) intensityOrder.SetValue( (imageSize * destLayerIdx) + pixelIdx, intensity.GetValue((imageSize * sourceLayerIdx) + pixelIdx), ) with open(os.path.join(directory, "alpha.uint8"), "wb") as f: f.write(memoryview(opacityOrder)) with open(os.path.join(directory, "intensity.uint8"), "wb") as f: f.write(memoryview(intensityOrder)) with open(os.path.join(directory, "order.uint8"), "wb") as f: f.write(memoryview(destOrder)) class SortedCompositeDataSetBuilder(VolumeCompositeDataSetBuilder): def __init__(self, location, cameraInfo, metadata={}, sections={}): VolumeCompositeDataSetBuilder.__init__( self, location, "image/png", cameraInfo, metadata, sections ) self.dataHandler.addTypes("sorted-composite", "rgba") # Register order and color textures self.layerScalars = [] self.dataHandler.registerData( name="order", type="array", fileName="/order.uint8" ) self.dataHandler.registerData( name="alpha", type="array", fileName="/alpha.uint8" ) self.dataHandler.registerData( name="intensity", type="array", fileName="/intensity.uint8", categories=["intensity"], ) def start(self, renderWindow, renderer): VolumeCompositeDataSetBuilder.start(self, renderWindow, renderer) imageSize = self.renderWindow.GetSize() self.dataConverter = ConvertVolumeStackToSortedStack(imageSize[0], imageSize[1]) def activateLayer(self, colorBy, scalar): VolumeCompositeDataSetBuilder.activateLayer( self, "root", "%s" % scalar, colorBy ) self.layerScalars.append(scalar) def writeData(self, mapper): VolumeCompositeDataSetBuilder.writeData(self, mapper) # Fill data pattern self.dataHandler.getDataAbsoluteFilePath("order") self.dataHandler.getDataAbsoluteFilePath("alpha") self.dataHandler.getDataAbsoluteFilePath("intensity") def stop(self, clean=True, compress=True): VolumeCompositeDataSetBuilder.stop(self, compress=False) # Go through all directories and convert them for root, dirs, files in os.walk(self.dataHandler.getBasePath()): for name in dirs: print("Process", os.path.join(root, name)) self.dataConverter.convert(os.path.join(root, name)) # Rename index.json to info_origin.json os.rename( os.path.join(self.dataHandler.getBasePath(), "index.json"), os.path.join(self.dataHandler.getBasePath(), "index_origin.json"), ) # Update index.json with open( os.path.join(self.dataHandler.getBasePath(), "index_origin.json"), "r" ) as infoFile: metadata = json.load(infoFile) metadata["SortedComposite"] = { "dimensions": metadata["CompositePipeline"]["dimensions"], "layers": self.dataConverter.layers, "scalars": self.layerScalars[0 : self.dataConverter.layers], } # Clean metadata dataToKeep = [] del metadata["CompositePipeline"] for item in metadata["data"]: if item["name"] in ["order", "alpha", "intensity"]: dataToKeep.append(item) metadata["data"] = dataToKeep metadata["type"] = ["tonic-query-data-model", "sorted-composite", "alpha"] # Override index.json with open( os.path.join(self.dataHandler.getBasePath(), "index.json"), "w" ) as newMetaFile: newMetaFile.write(json.dumps(metadata)) # Clean temporary data if clean: for root, dirs, files in os.walk(self.dataHandler.getBasePath()): print("Clean", root) for name in files: if ( "_rgb.png" in name or "_depth.uint8" in name or name == "index_origin.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 ".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))