milxLabelVisualisation.cpp

/*=========================================================================
  The Software is copyright (c) Commonwealth Scientific and Industrial Research Organisation (CSIRO)
  ABN 41 687 119 230.
  All rights reserved.

  Licensed under the CSIRO BSD 3-Clause License
  You may not use this file except in compliance with the License.
  You may obtain a copy of the License in the file LICENSE.md or at

  https://stash.csiro.au/projects/SMILI/repos/smili/browse/license.txt

  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.
=========================================================================*/
//Qt
#include <QApplication>

#include "milxQtImage.h"
#include "milxQtModel.h"
#include "milxQtPlot.h"
#include "milxQtFile.h"

#include <vtkWindowToImageFilter.h>

#include <tclap/CmdLine.h> //Command line parser library

using namespace TCLAP;

int main(int argc, char* argv[])
{
    QApplication app(argc,argv);
    QMainWindow mainWindow;//app takes ownership so need to be on stack, not heap

    CmdLine cmd("A visualisation tool for labelled images", ' ', milx::NumberToString(milx::Version));

    ValueArg<std::string> labelledImageArg("l", "labels", "Labelled image to visualise", true, "seg.nii.gz", "Labels");
    ValueArg<std::string> outputArg("o", "output", "Output Screenshot name", false, "screenshot.png", "Output");
    ValueArg<std::string> saveArg("", "save", "Save the mesh generated with name. Only for non-volume output.", false, "surface.vtk", "Save");
    ValueArg<std::string> surfaceArg("s", "surface", "Surface to overlay with labelled image", false, "surface.vtk", "Surface");
    ValueArg<std::string> transformArg("t", "transform", "Transform (ITK Format) to apply to objects being rendered", false, "rigid.txt", "Transform");
    ValueArg<std::string> loadViewFileArg("", "loadviewfile", "Load saved view from file (use onscreen mode to save view files)", false, "camera.cam", "Load View File");
    ValueArg<float> aboveArg("", "above", "Add above value to thresholding of the labels. Controls max of scalar range of colourmap also.", false, 255, "Above");
    ValueArg<float> belowArg("", "below", "Add below value to thresholding of the labels. Controls min of scalar range of colourmap also.", false, 0, "Below");
    ValueArg<float> opacityArg("", "opacity", "Set the opacity of the surface", false, 0.5, "Opacity");
    ValueArg<float> zoomArg("", "zoom", "Zoom in on the current view by given factor.", false, 2.0, "Zoom");
    ValueArg<int> heightArg("y", "height", "Set the height of the window.", false, 600, "Height");
    ValueArg<int> widthArg("x", "width", "Set the width of the window.", false, 800, "Width");
    ValueArg<int> linearDivisionArg("", "lineardivision", "Evenly space the label values across N to get full use of colourmap.", false, 16, "Linear Division");
    ValueArg<int> linearOffsetArg("", "linearoffset", "Offset the even spacing of the label values across N to get full use of colourmap.", false, 1, "Linear Offset");
    SwitchArg onscreenArg("", "onscreen", "Enable on screen rendering, i.e. display the rendering as an interactive window.", false);
    SwitchArg inverseArg("", "inverse", "Use the inverse transform when transform file is provided.", false);
    SwitchArg whiteArg("", "white", "Make background white rather than default gradient colour.", false);
    SwitchArg loadViewArg("", "loadview", "Load saved view (use smilx or onscreen render mode to view and save with Right Click->View->Save View", false);
    SwitchArg humanArg("", "nohuman", "Disable human orientation glyph.", false);
    SwitchArg processArg("", "postprocess", "Process the resulting surfaces from labels to reduce vertices and smooth etc.", false);
    SwitchArg oversmoothArg("", "oversmooth", "Process the resulting surfaces from labels with oversmoothing. Option is for Postprocess argument.", false);
    SwitchArg volumeArg("", "volume", "Display using volume rendering instead of marching cude iso-surfaces.", false);
    SwitchArg lineariseArg("", "linearise", "Evenly space the label values to get full use of colourmap.", false);
    SwitchArg wireframeArg("", "wireframe", "Display surface as a wireframe model.", false);
    SwitchArg axialArg("", "axial", "Show axial view based on position of first surface.", false);
    SwitchArg coronalArg("", "coronal", "Show coronal view based on position of first surface.", false);
    SwitchArg saggitalArg("", "saggital", "Show saggital view based on position of first surface.", false);
    SwitchArg jetArg("", "jet", "Change colourmap of the scalars to the Jet map (default)", false);
    SwitchArg vtkArg("", "vtk", "Change colourmap of the scalars to blue-red (rainbow VTK) map", false);
    SwitchArg hsvArg("", "hsv", "Change colourmap of the scalars to blue-red (rainbow HSV) map", false);
    SwitchArg rainbowArg("", "rainbow", "Change colourmap of the scalars to the rainbow map", false);
    SwitchArg spectralArg("", "spectral", "Change colourmap of the scalars to the spectral map", false);
    SwitchArg nihArg("", "NIH", "Change colourmap of the scalars to NIH", false);
    SwitchArg fireArg("", "NIH_FIRE", "Change colourmap of the scalars to NIH Fire", false);
    SwitchArg aalArg("", "AAL", "Change colourmap of the scalars to AAL", false);
    SwitchArg hotArg("", "HOT", "Change colourmap of the scalars to HOT", false);

    cmd.add( labelledImageArg );
    cmd.add( outputArg );
    cmd.add( saveArg );
    cmd.add( surfaceArg );
    cmd.add( transformArg );
    cmd.add( loadViewFileArg );
    cmd.add( aboveArg );
    cmd.add( belowArg );
    cmd.add( opacityArg );
    cmd.add( zoomArg );
    cmd.add( heightArg );
    cmd.add( widthArg );
    cmd.add( linearDivisionArg );
    cmd.add( linearOffsetArg );

    cmd.add( onscreenArg );
    cmd.add( inverseArg );
    cmd.add( whiteArg );
    cmd.add( loadViewArg );
    cmd.add( humanArg );
    cmd.add( processArg );
    cmd.add( oversmoothArg );
    cmd.add( volumeArg );
    cmd.add( lineariseArg );
    cmd.add( wireframeArg );
    cmd.add( axialArg );
    cmd.add( coronalArg );
    cmd.add( saggitalArg );

    cmd.add( jetArg );
    cmd.add( vtkArg );
    cmd.add( hsvArg );
    cmd.add( rainbowArg );
    cmd.add( spectralArg );
    cmd.add( nihArg );
    cmd.add( fireArg );
    cmd.add( aalArg );
    cmd.add( hotArg );

    cmd.parse( argc, argv );

    const std::string labelsName = labelledImageArg.getValue();
    const std::string screenName = outputArg.getValue();
    const std::string fileName = saveArg.getValue();
    const std::string surfaceName = surfaceArg.getValue();
    const std::string transformName = transformArg.getValue();
    const std::string loadViewName = loadViewFileArg.getValue();
    const float aboveValue = aboveArg.getValue();
    const float belowValue = belowArg.getValue();
    const float opacity = opacityArg.getValue();
    const float zoomFactor = zoomArg.getValue();
    const int windowHeight = heightArg.getValue();
    const int windowWidth = widthArg.getValue();
    int N = linearDivisionArg.getValue();
    int offset = linearOffsetArg.getValue();

    //Check arguments
    //Most of the checking is done by TCLAP
    if(inverseArg.isSet())
    {
        if(!transformArg.isSet())
        {
            cerr << "Error in arguments! Inverse argument needs to be used with the transform argument." << endl;
            exit(EXIT_FAILURE);
        }
    }

    //Read transform
    vtkSmartPointer<vtkTransform> transform = vtkSmartPointer<vtkTransform>::New(); //transform matrix
        transform->Identity();
        transform->PostMultiply();
    if(transformArg.isSet())
    {
        vtkSmartPointer<vtkMatrix4x4> matrix = milx::File::OpenITKTransform(transformName);
        transform->Concatenate(matrix);
        if(inverseArg.isSet())
            transform->Inverse();
    }

    std::cerr << "Reading Header for type info etc." << std::endl;
    std::string pixelType, componentType;
    size_t dimensions = 3;
    if( !milx::File::ReadImageInformation(labelsName, pixelType, componentType, dimensions) )
    {
        milx::PrintError("Failed Reading Image. Check the image type/file. Exiting.");
        exit(EXIT_FAILURE);
    }
    milx::PrintInfo("Pixel Type: " + pixelType);
    milx::PrintInfo("Component Type: " + componentType);
    milx::PrintInfo("Dimensions: " + milx::NumberToString(dimensions));

    //Reader object
    QScopedPointer<milxQtFile> reader(new milxQtFile);
    QScopedPointer<milxQtImage> labelledImage(new milxQtImage);  //smart deletion
    bool success = reader->openImage(labelsName.c_str(), labelledImage.data());
        labelledImage->setName(labelsName.c_str());
    if(aboveArg.isSet() || belowArg.isSet())
        labelledImage->threshold(0, belowValue, aboveValue);
        labelledImage->generateImage();

    QPointer<milxQtImage> eightBitImage;  //smart deletion
    if(componentType == "unsigned_char" || componentType == "unsigned char")
    {
        eightBitImage = labelledImage.data();
        milx::PrintInfo("Detected labelled images.");
    }
    else
    {
        milx::PrintWarning("8-bit image not found and will be cast to 8-bit image for visualisation.");
        eightBitImage = new milxQtImage(labelledImage.data());
        eightBitImage->setData( milx::Image<floatImageType>::CastImage<charImageType>(labelledImage->GetFloatImage()) );
    }
    eightBitImage->setName(labelsName.c_str());
    eightBitImage->generateImage();

    std::vector<unsigned char> values = milx::Image<charImageType>::LabelValues(eightBitImage->GetCharImage());
    std::cout << ">> " << values.size() << " Labels present: " << std::endl;
    for(size_t j = 0; j < values.size(); j ++)
        std::cout << static_cast<unsigned>(values[j]) << ", ";
    std::cout << std::endl;

    if(!linearDivisionArg.isSet())
        N = values.size();

    vtkSmartPointer<vtkLookupTable> lookupTable = vtkSmartPointer<vtkLookupTable>::New();
        lookupTable->SetTableRange(0.0, values.size()+1);
        lookupTable->Build();

    std::vector< QSharedPointer<milxQtModel> > isoSurfaces;
    QSharedPointer<milxQtRenderWindow> model(new milxQtRenderWindow);
        model->setName("Label Surface Rendering");
        model->SetSize(windowHeight, windowWidth);
        model->generateRender();

if(volumeArg.isSet())
{
    cout << ">> Overlay: Volume Rendering" << endl;
//    float minVolumeValue = numeric_limits<float>::max();
//    float maxVolumeValue = numeric_limits<float>::min();
    std::vector<float> colourValuesAsFloat, opacityValuesAsFloat;
    opacityValuesAsFloat.push_back(0.0); //background transparent
    colourValuesAsFloat.push_back(0.0); //background
    for(size_t j = 0; j < values.size(); j ++)
    {
        opacityValuesAsFloat.push_back(1.0);
        colourValuesAsFloat.push_back( static_cast<float>(values[j]) );
    }

    vtkSmartPointer<vtkImageFlip> imageReorient = vtkSmartPointer<vtkImageFlip>::New();
    #if VTK_MAJOR_VERSION <= 5
        imageReorient->SetInput(eightBitImage->GetOutput());
    #else
        imageReorient->SetInputData(eightBitImage->GetOutput());
    #endif
        imageReorient->SetFilteredAxis(1);
        imageReorient->FlipAboutOriginOn();
        imageReorient->Update(); //ITK image would have been flipped

    QScopedPointer<milxQtPlot> label(new milxQtPlot);  //smart deletion
        label->SetSource(imageReorient->GetOutput(), true);
        label->setOpacityTransferValues(opacityValuesAsFloat);
        label->setColourTransferValues(colourValuesAsFloat);
    //Colour maps
        label->colourMapToJet(belowValue, aboveValue); //default
    if(vtkArg.isSet())
        label->colourMapToVTK(belowValue, aboveValue);
    if(hsvArg.isSet())
        label->colourMapToHSV(belowValue, aboveValue);
    if(rainbowArg.isSet())
        label->colourMapToRainbow(belowValue, aboveValue);
    if(spectralArg.isSet())
        label->colourMapToSpectral(belowValue, aboveValue);
    if(nihArg.isSet())
        label->colourMapToNIH(belowValue, aboveValue);
    if(fireArg.isSet())
        label->colourMapToNIH_Fire(belowValue, aboveValue);
    if(aalArg.isSet())
        label->colourMapToAAL(belowValue, aboveValue);
    if(hotArg.isSet())
        label->colourMapToHOT(belowValue, aboveValue);
        label->volumePlot(imageReorient->GetOutput(), true, true);

    vtkVolumeCollection *volumeCollection = label->GetVolumes();

    volumeCollection->InitTraversal();
    vtkVolume * volume = volumeCollection->GetNextVolume();
    if(transformArg.isSet())
        volume->SetUserTransform(transform);
    model->AddVolume(volume);
}
else
{
    cout << ">> Overlay: Isosurfacing" << endl;
    const float linearScale = (aboveValue-belowValue)/N;
    for(size_t j = 0; j < values.size(); j ++)
    {
        if(lineariseArg.isSet())
            std::cout << ">> Iso surface label " << static_cast<unsigned>(values[j]) << " -> " << (j+offset)*linearScale << std::endl;
        else
            std::cout << ">> Iso surface label " << static_cast<unsigned>(values[j]) << std::endl;

        QScopedPointer<milxQtImage> label(new milxQtImage);  //smart deletion
            label->SetInput(eightBitImage->GetCharImage());
            label->binaryThreshold(1.0, values[j], values[j]);
            label->generateImage();

        QSharedPointer<milxQtModel> mdl(new milxQtModel);
            mdl->generateIsoSurface(label->GetOutput(), 0, 0.5);
        if(processArg.isSet())
        {
            mdl->quadricDecimate(0.25);
            if(oversmoothArg.isSet())
                mdl->smooth(500);
            else
                mdl->smoothSinc(15);
        }
        if(transformArg.isSet())
            mdl->SetTransform(transform);
            mdl->generateModel();

        //Colour maps
        model->colourMapToJet(belowValue, aboveValue); //default
        if(vtkArg.isSet())
            model->colourMapToVTK(belowValue, aboveValue);
        if(hsvArg.isSet())
            model->colourMapToHSV(belowValue, aboveValue);
        if(rainbowArg.isSet())
            model->colourMapToRainbow(belowValue, aboveValue);
        if(spectralArg.isSet())
            model->colourMapToSpectral(belowValue, aboveValue);
        if(nihArg.isSet())
            model->colourMapToNIH(belowValue, aboveValue);
        if(fireArg.isSet())
            model->colourMapToNIH_Fire(belowValue, aboveValue);
        if(aalArg.isSet())
            model->colourMapToAAL(belowValue, aboveValue);
        if(hotArg.isSet())
            model->colourMapToHOT(belowValue, aboveValue);

        double colour[3];
        if(lineariseArg.isSet())
            model->GetLookupTable()->GetColor((j+offset)*linearScale, colour); 
        else
            model->GetLookupTable()->GetColor(values[j], colour); 
        mdl->changeColour(colour[0], colour[1], colour[2]);

        model->AddActor(mdl->GetActor());
        isoSurfaces.push_back(mdl); //prevent deletion
    }
    
    if(saveArg.isSet())
    {
        QSharedPointer<milxQtModel> labelMesh(new milxQtModel);
        for(size_t j = 0; j < isoSurfaces.size(); j ++)
        {
            vtkSmartPointer<vtkPolyData> labelPolyData = vtkSmartPointer<vtkPolyData>::New();
            labelPolyData->DeepCopy(isoSurfaces[j].data()->GetOutput());
            vtkSmartPointer<vtkFloatArray> labelValueArray = vtkSmartPointer<vtkFloatArray>::New();
            labelValueArray->SetNumberOfComponents(1);
            labelValueArray->SetNumberOfTuples(labelPolyData->GetNumberOfPoints());
            labelValueArray->FillComponent(0, values[j]);
            labelPolyData->GetPointData()->SetScalars(labelValueArray);
            labelMesh->AddInput(labelPolyData);
        }
        labelMesh->generateModel();
        milx::File::SaveModel(fileName, labelMesh.data()->GetOutput());
    }
}
    cout << ">> Overlay: Generating Scene" << endl;
    model->generateRender();

    QScopedPointer<milxQtModel> surface(new milxQtModel); //smart deletion
    if(surfaceArg.isSet())
    {
        success = reader->openModel(surfaceName.c_str(), surface.data());
        if(!success)
        {
            milx::PrintError("Failed Reading surface. Check the surface type/file. Exiting.");
            exit(EXIT_FAILURE);
        }

        cout << ">> Overlaying surface" << endl;
        surface->setName(surfaceName.c_str());
        surface->generateModel();
        if(wireframeArg.isSet())
            surface->generateWireframe();
        if(opacityArg.isSet()) //needs to be after generate Model
            surface->SetOpacity(opacity);

        model->AddActor(surface->GetActor());
    }
    model->generateRender();

    mainWindow.setCentralWidget(model.data());
    mainWindow.resize(windowWidth, windowHeight);
    if(humanArg.isSet())
        model->disableOrient();
    if(whiteArg.isSet())
        model->background(true);
    if(axialArg.isSet())
        model->viewToAxial();
    if(coronalArg.isSet())
        model->viewToCoronal();
    if(saggitalArg.isSet())
        model->viewToSagittal();
    if(loadViewArg.isSet())
        model->loadView();
    if(loadViewFileArg.isSet())
        model->loadView(loadViewName.c_str());
    //Zoom
    if(zoomArg.isSet())
    {
        vtkCamera *camera = model->GetRenderer()->GetActiveCamera();
        camera->Zoom(zoomFactor);
    }
    cout << ">> Overlay: Rendering" << endl;
    if(!onscreenArg.isSet())
        model->OffScreenRenderingOn();
    else
        mainWindow.show();

    //Update view
    model->GetRenderWindow()->Render();
    qApp->processEvents();
    cout << ">> Complete" << endl;
        
    if(!onscreenArg.isSet())
    {
        vtkSmartPointer<vtkWindowToImageFilter> windowToImage = vtkSmartPointer<vtkWindowToImageFilter>::New();
        windowToImage->SetInput(model->GetRenderWindow());
        windowToImage->Update();

        QScopedPointer<milxQtFile> writer(new milxQtFile); //Smart deletion
        model->GetRenderWindow()->Render();
        writer->saveImage(screenName.c_str(), windowToImage->GetOutput());

        model->OffScreenRenderingOff(); //Required to prevent double-free
        return EXIT_SUCCESS;
    }
    else
        return app.exec();
}