I am using OpenCV (4.6.0) DNN module to generate semantic segmentation of images and the output of the network is a cv::Mat with the size of (numberOfClasses x image_height x image_width) that contains the class probabilities for every pixel.
I want to calculate the class ID that has the highest probability for every pixel.
In Python, the function numpy.argmax(src_matrix, axis=0) gives the desired output.
In C++ OpenCV, the function cv::reduceArgMax(src_, dst_, axis_) can calculate the same thing only on the 2D matrices. Therefore, I tried to get 2D slices ( (image_width x numberOfClasses) or ( numberOfClasses x image_height) ) from the 3D matrix and calculate the argmax on those slices in a loop. However I could not get the correct slices.
Example Code
int sizes[] = {numberOfClasses, imageHeight, imageWidth};
cv::Mat probabilityMatrix(3, sizes, CV_32F);
cv::Mat argMaxOfSlice(image_width, 1);
for(int r = 0; r < probabilityMatrix.size[1]; r++){
// cv::Mat Slice = get a 2D slice of the size (image_width x numberOfClasses) from the row r
// cv::reduceArgMax(Slice, argMaxOfSlice, 1);
...
}
Preferably, I just want to use OpenCV libraries but I can also use Eigen (3.2.10).
EDIT:
Python Example Code along with example input:
import numpy as np
# Shape of the example_input (3x3x4) where (ch, row, col)
example_input = np.array([[[ -1, 0, -1, -1],
[ 0, -1, -1, 0],
[ 0, -1, -1, -1]],
[[ -1, -1, 1, 1],
[ -1, -1, -1, -1],
[ 1, -1, 1, -1]],
[[ 2, -1, -1, -1],
[ -1, 2, 2, -1],
[ -1, 2, -1, 2]]])
expected_output = np.array([[ 2, 0, 1, 1],
[ 0, 2, 2, 0],
[ 1, 2, 1, 2]])
function_output = np.argmax(example_input, axis=0)
if np.count_nonzero(expected_output - function_output) > 0 :
print("Something wrong")
else:
print("Correct")
C++ OpenCV Example Input and Expected Output
int example_size[3] = {3, 3, 4};
float example_input_data[36] = { -1, 0, -1, 0, 0, -1, -1, 0, 0, -1, -1, -1, -1, -1, 1, 1, -1, -1, -1, -1,
1, -1, 1, -1, 2, -1, -1, -1, -1, 2, 2, -1, -1, 2, -1, 2};
cv::Mat example_input (3, example_size, CV_32F, example_input_data);
int expected_output_data[12] = { 2, 0, 1, 1, 0, 2, 2, 0, 1, 2, 1, 2};
cv::Mat expected_output (3, 4, CV_16U, expected_output_data);
Thank you