| 78 | }; |
| 79 | |
| 80 | class ArgmaxImageWriter : public ImageBase |
| 81 | { |
| 82 | public: |
| 83 | ArgmaxImageWriter(const std::string& filename, const nvinfer1::Dims& dims, const std::vector<int>& palette, const int num_classes); |
| 84 | void process(const int* buffer); |
| 85 | private: |
| 86 | int mNumClasses; |
| 87 | std::vector<int> mPalette; |
| 88 | }; |
| 89 | |
| 90 | }; // namespace util |
| 91 |