Implemented initial support for NN compression and decompression in i…

…mage pyramids

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();

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();

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();

source/FAST/Data/Access/ImagePyramidAccess.cpp

@@ -10,6 +10,9 @@
 #include <tiffio.h>
 #include <FAST/Data/Image.hpp>
 #include <jpeglib.h>
+#include <FAST/Algorithms/NeuralNetwork/NeuralNetwork.hpp>
+#include <FAST/Algorithms/NeuralNetwork/TensorToImage.hpp>
+#include <FAST/Algorithms/ImageCaster/ImageCaster.hpp>
 
 namespace fast {
 
@@ -147,13 +150,13 @@ std::unique_ptr<uchar[]> 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<uchar[]>(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<uchar[]> ImagePyramidAccess::getPatchData(int level, int x, int
                     auto tileData = std::make_unique<uchar[]>(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<Image> 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<std::mutex> 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<std::mutex> 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<Tensor>();
+    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<std::mutex> 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<uchar[]>(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<std::mutex> 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<Tensor>();
+        auto tensorToImage = TensorToImage::create()->connect(outputTensor);
+        auto image = tensorToImage->runAndGetOutputData<Image>();
+        // Have to go from float image to uint8 image
+        image = ImageCaster::create(TYPE_UINT8, m_image->getDecompressionOutputScaleFactor())->connect(image)->runAndGetOutputData<Image>();
+        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;
+    }
+}
+
+
 }

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<ImagePyramidAccess> pointer;
 	ImagePyramidAccess(std::vector<ImagePyramidLevel> levels, openslide_t* fileHandle, TIFF* tiffHandle, std::ifstream* stream, std::vector<vsi_tile_header>& vsiTiles, std::shared_ptr<ImagePyramid> imagePyramid, bool writeAccess, std::unordered_set<std::string>& initializedPatchList, std::mutex& readMutex, ImageCompression compressionFormat);
-	void setPatch(int level, int x, int y, std::shared_ptr<Image> patch);
+	void setPatch(int level, int x, int y, std::shared_ptr<Image> patch, bool propagate = true);
 	bool isPatchInitialized(uint level, uint x, uint y);
 	std::unique_ptr<uchar[]> 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<vsi_tile_header> m_vsiTiles;
     void readVSITileToBuffer(vsi_tile_header tile, uchar* data);
+    uint32_t writeTileToTIFF(int level, int x, int y, std::shared_ptr<Image> 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> image);
+    int readTileFromTIFF(void* data, int x, int y);
 };
 
 }

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<ImagePyramidLevel> 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<NeuralNetwork> compressionModel, std::shared_ptr<NeuralNetwork> decompressionModel, float decompressionOutputScaleFactor) {
+    setCompressionModel(compressionModel);
+    setDecompressionModel(decompressionModel, decompressionOutputScaleFactor);
+}
+
+std::shared_ptr<NeuralNetwork> ImagePyramid::getCompressionModel() const {
+    if(!m_compressionModel)
+        throw Exception("Image pyramid has no compression model");
+    return m_compressionModel;
+}
+
+std::shared_ptr<NeuralNetwork> ImagePyramid::getDecompressionModel() const {
+    if(!m_decompressionModel)
+        throw Exception("Image pyramid has no decompression model");
+    return m_decompressionModel;
+}
+
+void ImagePyramid::setCompressionModel(std::shared_ptr<NeuralNetwork> compressionModel) {
+    m_compressionModel = compressionModel;
+}
+
+void ImagePyramid::setDecompressionModel(std::shared_ptr<NeuralNetwork> decompressionModel, float outputScaleFactor) {
+    m_decompressionModel = decompressionModel;
+    m_decompressionOutputScaleFactor = outputScaleFactor;
+}
+
+float ImagePyramid::getDecompressionOutputScaleFactor() const {
+    return m_decompressionOutputScaleFactor;
+}
+
 }