From abe934dc0604bbcab483ffbe9573246fe18ec52c Mon Sep 17 00:00:00 2001 From: Erik Smistad Date: Mon, 6 Nov 2023 15:01:20 +0100 Subject: [PATCH] Implemented initial support for NN compression and decompression in image pyramids --- .../InferenceEngines/ONNXRuntimeEngine.cpp | 7 +- .../InferenceEngines/OpenVINOEngine.cpp | 7 +- .../InferenceEngines/TensorRTEngine.cpp | 16 ++- .../FAST/Data/Access/ImagePyramidAccess.cpp | 112 ++++++++++++++---- .../FAST/Data/Access/ImagePyramidAccess.hpp | 9 +- source/FAST/Data/ImagePyramid.cpp | 62 ++++++++++ source/FAST/Data/ImagePyramid.hpp | 11 ++ .../Exporters/TIFFImagePyramidExporter.cpp | 3 + 8 files changed, 202 insertions(+), 25 deletions(-) diff --git a/source/FAST/Algorithms/NeuralNetwork/InferenceEngines/ONNXRuntimeEngine.cpp b/source/FAST/Algorithms/NeuralNetwork/InferenceEngines/ONNXRuntimeEngine.cpp index 0a66b71e2..6bd13ae33 100644 --- a/source/FAST/Algorithms/NeuralNetwork/InferenceEngines/ONNXRuntimeEngine.cpp +++ b/source/FAST/Algorithms/NeuralNetwork/InferenceEngines/ONNXRuntimeEngine.cpp @@ -129,6 +129,7 @@ void ONNXRuntimeEngine::load() { const bool outputsDefined = !mOutputNodes.empty(); int inputCount = 0; int outputCount = 0; + bool imageOrderingFoundOnInput = false; for (size_t i = 0; i < m_session->GetInputCount(); i++) { std::string name = m_session->GetInputNameAllocated(i, allocator).get(); reportInfo() << "Found input node: " << name << " : " << print_shape(m_session->GetInputTypeInfo(i).GetTensorTypeAndShapeInfo().GetShape()) << reportEnd(); @@ -138,8 +139,10 @@ void ONNXRuntimeEngine::load() { } NodeType type = detectNodeType(shape); - if(type == NodeType::IMAGE) + if(type == NodeType::IMAGE) { m_imageOrdering = detectImageOrdering(shape); + imageOrderingFoundOnInput = true; + } if(inputsDefined) { if(mInputNodes.count(name) > 0) { reportInfo() << "Node was defined by user at id " << mInputNodes[name].id << reportEnd(); @@ -162,6 +165,8 @@ void ONNXRuntimeEngine::load() { shape.addDimension(x); } NodeType type = detectNodeType(shape); + if(!imageOrderingFoundOnInput && type == NodeType::IMAGE) + m_imageOrdering = detectImageOrdering(shape); if(outputsDefined) { if(mOutputNodes.count(name) > 0) { reportInfo() << "Node was defined by user at id " << mOutputNodes[name].id << reportEnd(); diff --git a/source/FAST/Algorithms/NeuralNetwork/InferenceEngines/OpenVINOEngine.cpp b/source/FAST/Algorithms/NeuralNetwork/InferenceEngines/OpenVINOEngine.cpp index 1cd9d3458..bd946b079 100644 --- a/source/FAST/Algorithms/NeuralNetwork/InferenceEngines/OpenVINOEngine.cpp +++ b/source/FAST/Algorithms/NeuralNetwork/InferenceEngines/OpenVINOEngine.cpp @@ -71,6 +71,7 @@ void OpenVINOEngine::load() { int outputCount = 0; bool dynamicInputShapes = false; int index = 0; + bool imageOrderingFoundOnInput = false; for(auto inputNode : model->inputs()) { auto name = inputNode.get_any_name(); reportInfo() << "Found input node " << inputNode.get_any_name() << " with shape " << inputNode.get_partial_shape().to_string() << " and index " << inputNode.get_index() << reportEnd(); @@ -88,8 +89,10 @@ void OpenVINOEngine::load() { m_inputIndices[name] = index;// inputNode.get_index(); // get_index always return 0 for some reason? ++index; NodeType type = detectNodeType(shape); - if(type == NodeType::IMAGE) + if(type == NodeType::IMAGE) { m_imageOrdering = detectImageOrdering(shape); + imageOrderingFoundOnInput = true; + } if(inputsDefined) { if(mInputNodes.count(name) > 0) { reportInfo() << "Node was defined by user at id " << mInputNodes[name].id << reportEnd(); @@ -150,6 +153,8 @@ void OpenVINOEngine::load() { m_outputIndices[name] = index;// outputNode.get_index(); // get_index always returns 0 for some reason.. ++index; NodeType type = detectNodeType(shape); + if(!imageOrderingFoundOnInput && type == NodeType::IMAGE) + m_imageOrdering = detectImageOrdering(shape); if(outputsDefined) { if(mOutputNodes.count(name) > 0) { reportInfo() << "Node was defined by user at id " << mOutputNodes[name].id << reportEnd(); diff --git a/source/FAST/Algorithms/NeuralNetwork/InferenceEngines/TensorRTEngine.cpp b/source/FAST/Algorithms/NeuralNetwork/InferenceEngines/TensorRTEngine.cpp index ea7379a74..a919357ea 100644 --- a/source/FAST/Algorithms/NeuralNetwork/InferenceEngines/TensorRTEngine.cpp +++ b/source/FAST/Algorithms/NeuralNetwork/InferenceEngines/TensorRTEngine.cpp @@ -395,13 +395,16 @@ void TensorRTEngine::load() { if(filename.substr(filename.size()-4) != ".uff") { int inputCount = 0; int outputCount = 0; + bool imageOrderingFoundOnInput = false; for(auto binding : bindings) { auto name = binding.first; int i = binding.second; auto shape = getTensorShape(m_engine->getBindingDimensions(i)); NodeType type = detectNodeType(shape); - if(type == NodeType::IMAGE && m_engine->bindingIsInput(i)) + if(type == NodeType::IMAGE && m_engine->bindingIsInput(i)) { m_imageOrdering = detectImageOrdering(shape); + imageOrderingFoundOnInput = true; + } if(m_engine->bindingIsInput(i)) { reportInfo() << "Found input node " << name << " with shape " << shape.toString() << reportEnd(); auto dims = m_engine->getBindingDimensions(i); @@ -446,6 +449,17 @@ void TensorRTEngine::load() { } } } + if(!imageOrderingFoundOnInput) { + for(auto binding : bindings) { + auto name = binding.first; + int i = binding.second; + auto shape = getTensorShape(m_engine->getBindingDimensions(i)); + NodeType type = detectNodeType(shape); + if(type == NodeType::IMAGE) { + m_imageOrdering = detectImageOrdering(shape); + } + } + } } m_context = m_engine->createExecutionContext(); diff --git a/source/FAST/Data/Access/ImagePyramidAccess.cpp b/source/FAST/Data/Access/ImagePyramidAccess.cpp index 3937d91e5..6a1cc46a3 100644 --- a/source/FAST/Data/Access/ImagePyramidAccess.cpp +++ b/source/FAST/Data/Access/ImagePyramidAccess.cpp @@ -10,6 +10,9 @@ #include #include #include +#include +#include +#include namespace fast { @@ -147,13 +150,13 @@ std::unique_ptr ImagePyramidAccess::getPatchData(int level, int x, int } if(width == tileWidth && height == tileHeight && x % tileWidth == 0 && y % tileHeight == 0) { // From TIFFReadTile documentation: Return the data for the tile containing the specified coordinates. - int bytesRead = TIFFReadTile(m_tiffHandle, (void *) data.get(), x, y, 0, 0); + int bytesRead = readTileFromTIFF((void *) data.get(), x, y); } else if((width < tileWidth || height < tileHeight) && x % tileWidth == 0 && y % tileHeight == 0) { auto tileData = std::make_unique(tileWidth*tileHeight*channels); { // From TIFFReadTile documentation: Return the data for the tile containing the specified coordinates. // In TIFF all tiles have the same size, thus they are padded.. - int bytesRead = TIFFReadTile(m_tiffHandle, (void *) tileData.get(), x, y, 0, 0); + int bytesRead = readTileFromTIFF((void *) tileData.get(), x, y); } // Remove extra for(int dy = 0; dy < height; ++dy) { @@ -188,7 +191,7 @@ std::unique_ptr ImagePyramidAccess::getPatchData(int level, int x, int auto tileData = std::make_unique(tileWidth*tileHeight*channels); int tileX = i*tileWidth; int tileY = j*tileHeight; - int bytesRead = TIFFReadTile(m_tiffHandle, (void *) tileData.get(), firstTileX*tileWidth+tileX, firstTileY*tileHeight+tileY, 0, 0); + int bytesRead = readTileFromTIFF((void *) tileData.get(), firstTileX*tileWidth+tileX, firstTileY*tileHeight+tileY); // Stitch tile into full buffer for(int cy = 0; cy < tileHeight; ++cy) { for(int cx = 0; cx < tileWidth; ++cx) { @@ -427,7 +430,55 @@ std::shared_ptr ImagePyramidAccess::getPatchAsImage(int level, int tileX, } } -void ImagePyramidAccess::setPatch(int level, int x, int y, Image::pointer patch) { +uint32_t ImagePyramidAccess::writeTileToTIFF(int level, int x, int y, uchar *data, int width, int height, int channels) { + if(m_image->getCompression() == ImageCompression::NEURAL_NETWORK) { + auto image = Image::create(width, height, TYPE_UINT8, channels, data); // TODO this seems unnecessary + return writeTileToTIFF(level, x, y, image); + } else { + return writeTileToTIFF(level, x, y, data); + } +} + +uint32_t ImagePyramidAccess::writeTileToTIFF(int level, int x, int y, Image::pointer image) { + if(m_image->getCompression() == ImageCompression::NEURAL_NETWORK) { + return writeTileToTIFFNeuralNetwork(level, x, y, image); + } else { + auto access = image->getImageAccess(ACCESS_READ); + return writeTileToTIFF(level, x, y, (uchar*)access->get()); + } +} + +uint32_t ImagePyramidAccess::writeTileToTIFF(int level, int x, int y, uchar *data) { + std::lock_guard lock(m_readMutex); + TIFFSetDirectory(m_tiffHandle, level); + TIFFWriteTile(m_tiffHandle, (void *) data, x, y, 0, 0); + TIFFCheckpointDirectory(m_tiffHandle); + uint32_t tile_id = TIFFComputeTile(m_tiffHandle, x, y, 0, 0); + return tile_id; +} + +uint32_t ImagePyramidAccess::writeTileToTIFFNeuralNetwork(int level, int x, int y, Image::pointer image) { + std::lock_guard lock(m_readMutex); + TIFFSetDirectory(m_tiffHandle, level); + uint32_t tile_id = TIFFComputeTile(m_tiffHandle, x, y, 0, 0); + if(m_image->getCompression() != ImageCompression::NEURAL_NETWORK) + throw Exception("Compression is not neural network type"); + + auto compressionModel = m_image->getCompressionModel(); + compressionModel->connect(image); + auto tensor = compressionModel->runAndGetOutputData(); + auto access = tensor->getAccess(ACCESS_READ); + float* data = access->getRawData(); + uint32_t size = tensor->getShape().getTotalSize()*4; + TIFFSetWriteOffset(m_tiffHandle, 0); // Set write offset to 0, so that we dont appen data + TIFFWriteRawTile(m_tiffHandle, tile_id, (void *) data, size); // This appends data.. + //TIFFWriteTile(m_tiffHandle, (void *) data, x, y, 0,0); // This does not append, but tries to compress data + + TIFFCheckpointDirectory(m_tiffHandle); + return tile_id; +} + +void ImagePyramidAccess::setPatch(int level, int x, int y, Image::pointer patch, bool propagate) { if(m_tiffHandle == nullptr) throw Exception("setPatch only available for TIFF backend ImagePyramids"); @@ -452,16 +503,7 @@ void ImagePyramidAccess::setPatch(int level, int x, int y, Image::pointer patch) } // Write tile to this level - auto patchAccess = patch->getImageAccess(ACCESS_READ); - auto data = (uchar*)patchAccess->get(); - uint32_t tile_id; - { - std::lock_guard lock(m_readMutex); - TIFFSetDirectory(m_tiffHandle, level); - TIFFWriteTile(m_tiffHandle, (void *) data, x, y, 0, 0); - TIFFCheckpointDirectory(m_tiffHandle); - tile_id = TIFFComputeTile(m_tiffHandle, x, y, 0, 0); - } + uint32_t tile_id = writeTileToTIFF(level, x, y, patch); m_initializedPatchList.insert(std::to_string(level) + "-" + std::to_string(tile_id)); // Add patch to list of dirty patches, so the renderer can update it if needed @@ -474,7 +516,11 @@ void ImagePyramidAccess::setPatch(int level, int x, int y, Image::pointer patch) m_image->setDirtyPatch(level, patchIdX, patchIdY); // Propagate upwards + if(!propagate) + return; auto previousData = std::make_unique(patch->getNrOfVoxels()*patch->getNrOfChannels()); + auto patchAccess = patch->getImageAccess(ACCESS_READ); + auto data = (uchar*)patchAccess->get(); std::memcpy(previousData.get(), data, patch->getNrOfVoxels()*patch->getNrOfChannels()); const auto channels = m_image->getNrOfChannels(); while(level < m_image->getNrOfLevels()-1) { @@ -540,13 +586,7 @@ void ImagePyramidAccess::setPatch(int level, int x, int y, Image::pointer patch) } } } - { - std::lock_guard lock(m_readMutex); - TIFFSetDirectory(m_tiffHandle, level); - TIFFWriteTile(m_tiffHandle, (void *) newData.get(), x, y, 0, 0); - TIFFCheckpointDirectory(m_tiffHandle); - tile_id = TIFFComputeTile(m_tiffHandle, x, y, 0, 0); - } + tile_id = writeTileToTIFF(level, x, y, newData.get(), tileWidth, tileHeight, channels); previousData = std::move(newData); int levelWidth = m_image->getLevelWidth(level); @@ -569,4 +609,34 @@ bool ImagePyramidAccess::isPatchInitialized(uint level, uint x, uint y) { return m_initializedPatchList.count(std::to_string(level) + "-" + std::to_string(tile)) > 0; } +int ImagePyramidAccess::readTileFromTIFF(void *data, int x, int y) { + // Assumes level (directory is already set) + if(m_compressionFormat == ImageCompression::NEURAL_NETWORK) { + auto decompressionModel = m_image->getDecompressionModel(); + // TODO The logic here must be improved + // TODO this assumes fixed size code + auto shape = decompressionModel->getInputNodes().begin()->second.shape; + shape[0] = 1; + int64_t size = shape.getTotalSize()*4; + float* buffer = new float[shape.getTotalSize()]; + uint32_t tile_id = TIFFComputeTile(m_tiffHandle, x, y, 0, 0); + int bytesRead = TIFFReadRawTile(m_tiffHandle, tile_id, buffer, size); + auto tensor = Tensor::create(buffer, shape); + decompressionModel->connect(tensor); + // TODO TensorToImage not really needed.. + auto outputTensor = decompressionModel->runAndGetOutputData(); + auto tensorToImage = TensorToImage::create()->connect(outputTensor); + auto image = tensorToImage->runAndGetOutputData(); + // Have to go from float image to uint8 image + image = ImageCaster::create(TYPE_UINT8, m_image->getDecompressionOutputScaleFactor())->connect(image)->runAndGetOutputData(); + auto access = image->getImageAccess(ACCESS_READ); + std::memcpy(data, access->get(), image->getNrOfVoxels()*image->getNrOfChannels()); + return bytesRead; + } else { + int bytesRead = TIFFReadTile(m_tiffHandle, data, x, y, 0, 0); + return bytesRead; + } +} + + } diff --git a/source/FAST/Data/Access/ImagePyramidAccess.hpp b/source/FAST/Data/Access/ImagePyramidAccess.hpp index 4f4a89a9c..6b9e9f89c 100644 --- a/source/FAST/Data/Access/ImagePyramidAccess.hpp +++ b/source/FAST/Data/Access/ImagePyramidAccess.hpp @@ -13,6 +13,7 @@ namespace fast { class Image; class ImagePyramid; +class NeuralNetwork; /** @@ -26,6 +27,7 @@ enum class ImageCompression { JPEG, JPEG2000, LZW, // Lossless compression + NEURAL_NETWORK, // Use neural network to do the compression and decompression. See ImagePyramid::setCompressionModels }; struct vsi_tile_header { @@ -68,7 +70,7 @@ class FAST_EXPORT ImagePyramidAccess : Object { public: typedef std::unique_ptr pointer; ImagePyramidAccess(std::vector levels, openslide_t* fileHandle, TIFF* tiffHandle, std::ifstream* stream, std::vector& vsiTiles, std::shared_ptr imagePyramid, bool writeAccess, std::unordered_set& initializedPatchList, std::mutex& readMutex, ImageCompression compressionFormat); - void setPatch(int level, int x, int y, std::shared_ptr patch); + void setPatch(int level, int x, int y, std::shared_ptr patch, bool propagate = true); bool isPatchInitialized(uint level, uint x, uint y); std::unique_ptr getPatchData(int level, int x, int y, int width, int height); ImagePyramidPatch getPatch(std::string tile); @@ -90,6 +92,11 @@ class FAST_EXPORT ImagePyramidAccess : Object { ImageCompression m_compressionFormat; std::vector m_vsiTiles; void readVSITileToBuffer(vsi_tile_header tile, uchar* data); + uint32_t writeTileToTIFF(int level, int x, int y, std::shared_ptr image); + uint32_t writeTileToTIFF(int level, int x, int y, uchar* data, int width, int height, int channels); + uint32_t writeTileToTIFF(int level, int x, int y, uchar* data); + uint32_t writeTileToTIFFNeuralNetwork(int level, int x, int y, std::shared_ptr image); + int readTileFromTIFF(void* data, int x, int y); }; } \ No newline at end of file diff --git a/source/FAST/Data/ImagePyramid.cpp b/source/FAST/Data/ImagePyramid.cpp index dce60759d..d350bb5a9 100644 --- a/source/FAST/Data/ImagePyramid.cpp +++ b/source/FAST/Data/ImagePyramid.cpp @@ -76,6 +76,7 @@ ImagePyramid::ImagePyramid(int width, int height, int channels, int patchWidth, photometric = PHOTOMETRIC_RGB; samplesPerPixel = 3; // RGBA image pyramid is converted to RGB with getPatchAsImage } + m_compressionFormat = compression; while(true) { currentWidth = width / std::pow(2, currentLevel); @@ -127,6 +128,9 @@ ImagePyramid::ImagePyramid(int width, int height, int channels, int patchWidth, throw NotImplementedException(); TIFFSetField(tiff, TIFFTAG_COMPRESSION, COMPRESSION_JP2000); break; + case ImageCompression::NEURAL_NETWORK: + TIFFSetField(tiff, TIFFTAG_COMPRESSION, 34666); // TODO What should the value be? + break; } TIFFSetField(tiff, TIFFTAG_TILEWIDTH, levelData.tileWidth); @@ -365,6 +369,30 @@ ImagePyramid::ImagePyramid(TIFF *fileHandle, std::vector leve m_levels[i].tilesX = std::ceil((float)m_levels[i].width / m_levels[i].tileWidth); m_levels[i].tilesY = std::ceil((float)m_levels[i].height / m_levels[i].tileHeight); } + // Get compression + uint compressionTag; + TIFFGetField(fileHandle, TIFFTAG_COMPRESSION, &compressionTag); + ImageCompression compression; + switch(compressionTag) { + case COMPRESSION_NONE: + compression = ImageCompression::RAW; + break; + case COMPRESSION_JPEG: + compression = ImageCompression::JPEG; + break; + case COMPRESSION_LZW: + compression = ImageCompression::LZW; + break; + case COMPRESSION_JP2000: + throw Exception("JPEG 2000 TIFF not supported yet"); + break; + case 34666: + compression = ImageCompression::NEURAL_NETWORK; + break; + default: + reportWarning() << "Unrecognized compression in TIFF: " << compressionTag << reportEnd(); + } + m_compressionFormat = compression; // Get spacing from TIFF float spacingX; float spacingY; @@ -470,4 +498,38 @@ bool ImagePyramid::isOMETIFF() const { return m_isOMETIFF; } +ImageCompression ImagePyramid::getCompression() const { + return m_compressionFormat; +} + +void ImagePyramid::setCompressionModels(std::shared_ptr compressionModel, std::shared_ptr decompressionModel, float decompressionOutputScaleFactor) { + setCompressionModel(compressionModel); + setDecompressionModel(decompressionModel, decompressionOutputScaleFactor); +} + +std::shared_ptr ImagePyramid::getCompressionModel() const { + if(!m_compressionModel) + throw Exception("Image pyramid has no compression model"); + return m_compressionModel; +} + +std::shared_ptr ImagePyramid::getDecompressionModel() const { + if(!m_decompressionModel) + throw Exception("Image pyramid has no decompression model"); + return m_decompressionModel; +} + +void ImagePyramid::setCompressionModel(std::shared_ptr compressionModel) { + m_compressionModel = compressionModel; +} + +void ImagePyramid::setDecompressionModel(std::shared_ptr decompressionModel, float outputScaleFactor) { + m_decompressionModel = decompressionModel; + m_decompressionOutputScaleFactor = outputScaleFactor; +} + +float ImagePyramid::getDecompressionOutputScaleFactor() const { + return m_decompressionOutputScaleFactor; +} + } diff --git a/source/FAST/Data/ImagePyramid.hpp b/source/FAST/Data/ImagePyramid.hpp index 8a5d35a65..a6d2db795 100644 --- a/source/FAST/Data/ImagePyramid.hpp +++ b/source/FAST/Data/ImagePyramid.hpp @@ -67,6 +67,13 @@ class FAST_EXPORT ImagePyramid : public SpatialDataObject { // Override DataBoundingBox getTransformedBoundingBox() const override; DataBoundingBox getBoundingBox() const override; + ImageCompression getCompression() const; + void setCompressionModels(std::shared_ptr compressionModel, std::shared_ptr decompressionModel, float outputScaleFactor = 1.0f); + void setCompressionModel(std::shared_ptr compressionModel); + void setDecompressionModel(std::shared_ptr decompressionModel, float outputScaleFactor = 1.0f); + std::shared_ptr getCompressionModel() const; + std::shared_ptr getDecompressionModel() const; + float getDecompressionOutputScaleFactor() const; private: ImagePyramid(); std::vector m_levels; @@ -98,6 +105,10 @@ class FAST_EXPORT ImagePyramid : public SpatialDataObject { // A mutex needed to control multi-threaded reading of VSI and TIFF files std::mutex m_readMutex; + + std::shared_ptr m_compressionModel; + std::shared_ptr m_decompressionModel; + float m_decompressionOutputScaleFactor = 1.0f; }; } diff --git a/source/FAST/Exporters/TIFFImagePyramidExporter.cpp b/source/FAST/Exporters/TIFFImagePyramidExporter.cpp index b7a8e69f8..c9b0ea75e 100644 --- a/source/FAST/Exporters/TIFFImagePyramidExporter.cpp +++ b/source/FAST/Exporters/TIFFImagePyramidExporter.cpp @@ -104,6 +104,9 @@ void TIFFImagePyramidExporter::execute() { throw NotImplementedException(); TIFFSetField(tiff, TIFFTAG_COMPRESSION, COMPRESSION_JP2000); break; + case ImageCompression::NEURAL_NETWORK: + TIFFSetField(tiff, TIFFTAG_COMPRESSION, 34666); // TODO What should the value be? + break; } TIFFSetField(tiff, TIFFTAG_TILEWIDTH, imagePyramid->getLevelTileWidth(level));