ImageFusionProto/vtk_slice_backend.cpp at main · Nowhere4Nothing/ImageFusionProto · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <vtkSmartPointer.h>
#include <vtkImageData.h>
#include <vtkImageReslice.h>
#include <vtkImageBlend.h>
#include <vtkMatrix4x4.h>
#include <vtkTransform.h>
#include <vtkPointData.h>
#include <vtkDICOMImageReader.h>

namespace py = pybind11;

class VTKSliceExtractor {
public:
    VTKSliceExtractor() {
        reslice = vtkSmartPointer<vtkImageReslice>::New();
        reslice->SetInterpolationModeToLinear();
        reslice->SetAutoCropOutput(1);

        blend = vtkSmartPointer<vtkImageBlend>::New();
        blend->SetOpacity(0, 1.0);
        blend->SetOpacity(1, 0.5);

        transform = vtkSmartPointer<vtkTransform>::New();
        transform->PostMultiply();
    }

    void set_fixed(vtkImageData* img) {
        fixed = img;
        _wire_blend();
        _sync_reslice_output();
    }

    void set_moving(vtkImageData* img) {
        moving = img;
        reslice->SetInputData(moving);
        _apply_transform();
        _wire_blend();
        _sync_reslice_output();
    }

    void set_transform(double tx, double ty, double tz,
                       double rx, double ry, double rz) {
        if (!fixed || !moving) return;
        double center[3];
        fixed->GetCenter(center);

        vtkSmartPointer<vtkTransform> t = vtkSmartPointer<vtkTransform>::New();
        t->PostMultiply();
        t->Translate(-center[0], -center[1], -center[2]);
        t->RotateX(rx);
        t->RotateY(ry);
        t->RotateZ(rz);
        t->Translate(center[0], center[1], center[2]);
        t->Translate(tx, ty, tz);

        transform->DeepCopy(t);
        _apply_transform();
    }

    void set_opacity(double alpha) {
        blend->SetOpacity(1, alpha);
    }

    py::array_t<uint8_t> get_slice(const std::string& orientation, int slice_idx) {
        blend->Update();
        vtkImageData* img = blend->GetOutput();
        int extent[6];
        img->GetExtent(extent);
        int nx = extent[1]-extent[0]+1;
        int ny = extent[3]-extent[2]+1;
        int nz = extent[5]-extent[4]+1;

        uint8_t* out_ptr = new uint8_t[nx*ny]; // max size, will slice later
        int w,h;
        if (orientation == "axial") {
            int z = std::max(0, std::min(slice_idx - extent[4], nz-1));
            for(int j=0;j<ny;j++){
                for(int i=0;i<nx;i++){
                    double* pixel = static_cast<double*>(img->GetScalarPointer(i,j,z));
                    out_ptr[j*nx+i] = static_cast<uint8_t>(*pixel);
                }
            }
            w = nx; h = ny;
        } else if (orientation == "coronal") {
            int y = std::max(0,std::min(slice_idx - extent[2], ny-1));
            for(int k=0;k<nz;k++){
                for(int i=0;i<nx;i++){
                    double* pixel = static_cast<double*>(img->GetScalarPointer(i,y,k));
                    out_ptr[k*nx+i] = static_cast<uint8_t>(*pixel);
                }
            }
            w = nx; h = nz;
        } else if (orientation == "sagittal") {
            int x = std::max(0,std::min(slice_idx - extent[0], nx-1));
            for(int k=0;k<nz;k++){
                for(int j=0;j<ny;j++){
                    double* pixel = static_cast<double*>(img->GetScalarPointer(x,j,k));
                    out_ptr[k*ny+j] = static_cast<uint8_t>(*pixel);
                }
            }
            w = ny; h = nz;
        } else {
            delete[] out_ptr;
            return py::array_t<uint8_t>();
        }

        auto result = py::array_t<uint8_t>({h,w}, out_ptr);
        return result;
    }

private:
    vtkSmartPointer<vtkImageReslice> reslice;
    vtkSmartPointer<vtkImageBlend> blend;
    vtkSmartPointer<vtkTransform> transform;
    vtkImageData* fixed = nullptr;
    vtkImageData* moving = nullptr;

    void _apply_transform() {
        if (!fixed || !moving) return;
        reslice->SetResliceAxes(transform->GetMatrix());
        reslice->Modified();
    }

    void _wire_blend() {
        blend->RemoveAllInputs();
        if (fixed) blend->AddInputData(fixed);
        if (moving) blend->AddInputConnection(reslice->GetOutputPort());
    }

    void _sync_reslice_output() {
        if (!fixed) return;
        double sp[3]; fixed->GetSpacing(sp);
        double org[3]; fixed->GetOrigin(org);
        int ext[6]; fixed->GetExtent(ext);
        reslice->SetOutputSpacing(sp);
        reslice->SetOutputOrigin(org);
        reslice->SetOutputExtent(ext);
        reslice->Modified();
    }
};

PYBIND11_MODULE(vtk_slice_backend, m) {
    py::class_<VTKSliceExtractor>(m, "VTKSliceExtractor")
        .def(py::init<>())
        .def("set_fixed", &VTKSliceExtractor::set_fixed)
        .def("set_moving", &VTKSliceExtractor::set_moving)
        .def("set_transform", &VTKSliceExtractor::set_transform)
        .def("set_opacity", &VTKSliceExtractor::set_opacity)
        .def("get_slice", &VTKSliceExtractor::get_slice);
}