template<class T>
T get_value(const ncnn::Mat &_to_process_mat,
            const int row_index,
            const int col_index,
            const int channel_index) {
    return *((T *) _to_process_mat.data + (channel_index * _to_process_mat.cstep) + row_index * _to_process_mat.w +
             col_index);
}

template<class T>
void set_value(const ncnn::Mat &_to_process_mat,
               const int row_index,
               const int col_index,
               const int channel_index,
               T value) {
    *((T *) _to_process_mat.data + (channel_index * _to_process_mat.cstep) + row_index * _to_process_mat.w +
      col_index) = value;
}

template<class T>
void argmax_first_axis(const ncnn::Mat &_to_process_mat, ncnn::Mat &_result_mat) {
    assert(sizeof(T) == _to_process_mat.elemsize);
    int classes = _to_process_mat.c;
    for (int h = 0; h < _to_process_mat.h; ++h) {
        for (int w = 0; w < _to_process_mat.w; ++w) {
            int max_index = 0;
            T max_value = get_value<T>(_to_process_mat, h, w, 0);
            for (int i = 1; i < classes; ++i) {
                T current_value = get_value<T>(_to_process_mat, h, w, i);
                if (current_value > max_value) {
                    max_index = i;
                    max_value = current_value;
                }
            }
            set_value<T>(_result_mat, h, w, 0, max_index);
        }
    }
}