Repository: opencv/opencv
File: modules/dnn/src/layers/topk2_layer.cpp (C++)

Code before:

            auto kind = outputs_arr.kind();
            if (kind == _InputArray::STD_VECTOR_MAT) {
                std::vector<Mat>& outs = outputs_arr.getMatVecRef();
                CV_Assert(outs.size() == 2);
                outs[0].fit(outShape, input.type());
                outs[1].fit(outShape, CV_64S);
                output_value = outs[0];
                output_index = outs[1];
            } else if (kind == _InputArray::STD_VECTOR_UMAT) {
                std::vector<UMat>& uouts = outputs_arr.getUMatVecRef();
                CV_Assert(uouts.size() == 2);
                uouts[0].fit(outShape, input.type());
                uouts[1].fit(outShape, CV_64S);
                output_value = uouts[0].getMat(ACCESS_WRITE);
                output_index = uouts[1].getMat(ACCESS_WRITE);
            } else {
                // Fallback – recreate directly on provided Mat refs.
                output_value.create(outShape, input.type());
                output_index.create(outShape, CV_64S);
            }
        }

        if (largest) {
            switch (input.depth()) {
                case CV_8U: FindTopK<ComparatorGreater<uint8_t>, uint8_t>(input, output_value, output_index, normalized_axis); break;
                case CV_8S: FindTopK<ComparatorGreater<int8_t>, int8_t>(input, output_value, output_index, normalized_axis); break;
                case CV_16U: FindTopK<ComparatorGreater<uint16_t>, uint16_t>(input, output_value, output_index, normalized_axis); break;
                case CV_16S: FindTopK<ComparatorGreater<int16_t>, int16_t>(input, output_value, output_index, normalized_axis); break;
                case CV_16F: FindTopK<ComparatorGreater<hfloat>, hfloat>(input, output_value, output_index, normalized_axis); break;
                case CV_32U: FindTopK<ComparatorGreater<unsigned>, unsigned>(input, output_value, output_index, normalized_axis); break;
                case CV_32S: FindTopK<ComparatorGreater<int>, int>(input, output_value, output_index, normalized_axis); break;
                case CV_32F: FindTopK<ComparatorGreater<float>, float>(input, output_value, output_index, normalized_axis); break;
                case CV_64U: FindTopK<ComparatorGreater<uint64_t>, uint64_t>(input, output_value, output_index, normalized_axis); break;
                case CV_64S: FindTopK<ComparatorGreater<int64_t>, int64_t>(input, output_value, output_index, normalized_axis); break;
                case CV_64F: FindTopK<ComparatorGreater<double>, double>(input, output_value, output_index, normalized_axis); break;
                default: CV_Error(Error::BadDepth, "Unsupported input data type");
            }
        } else {
            switch (input.depth()) {
                case CV_8U: FindTopK<ComparatorLess<uint8_t>, uint8_t>(input, output_value, output_index, normalized_axis); break;
                case CV_8S: FindTopK<ComparatorLess<int8_t>, int8_t>(input, output_value, output_index, normalized_axis); break;
                case CV_16U: FindTopK<ComparatorLess<uint16_t>, uint16_t>(input, output_value, output_index, normalized_axis); break;
                case CV_16S: FindTopK<ComparatorLess<int16_t>, int16_t>(input, output_value, output_index, normalized_axis); break;
                case CV_16F: FindTopK<ComparatorLess<hfloat>, hfloat>(input, output_value, output_index, normalized_axis); break;
                case CV_32U: FindTopK<ComparatorLess<unsigned>, unsigned>(input, output_value, output_index, normalized_axis); break;
                case CV_32S: FindTopK<ComparatorLess<int>, int>(input, output_value, output_index, normalized_axis); break;
                case CV_32F: FindTopK<ComparatorLess<float>, float>(input, output_value, output_index, normalized_axis); break;
                case CV_64U: FindTopK<ComparatorLess<uint64_t>, uint64_t>(input, output_value, output_index, normalized_axis); break;
                case CV_64S: FindTopK<ComparatorLess<int64_t>, int64_t>(input, output_value, output_index, normalized_axis); break;
                case CV_64F: FindTopK<ComparatorLess<double>, double>(input, output_value, output_index, normalized_axis); break;

Review comment (suggestion):
I suggest you to construct `ComparatorGreater<T>` or `ComparatorLess<T>` (actually, it should be WT instead of T, see the comment below) inside FindTopK(), depending on the `largest` flag. This way you reduce the amount of code with basically no loss in efficiency, you will have just one switch

Code after:
            uouts[1].fit(outShape, CV_64S);
            output_value = uouts[0].getMat(ACCESS_WRITE);
            output_index = uouts[1].getMat(ACCESS_WRITE);
        } else {
            CV_Error(cv::Error::StsBadArg, cv::format("Unsupported output array kind: %d", kind));
        }

        switch (input.depth()) {
            case CV_8U:   FindTopK<uint8_t          >(input, output_value, output_index, normalized_axis, kVal); break;
            case CV_8S:   FindTopK<int8_t           >(input, output_value, output_index, normalized_axis, kVal); break;
            case CV_16U:  FindTopK<uint16_t         >(input, output_value, output_index, normalized_axis, kVal); break;
            case CV_16S:  FindTopK<int16_t          >(input, output_value, output_index, normalized_axis, kVal); break;
            case CV_16F:  FindTopK<hfloat, float    >(input, output_value, output_index, normalized_axis, kVal); break;
            case CV_16BF: FindTopK<bfloat, float    >(input, output_value, output_index, normalized_axis, kVal); break;
            case CV_32U:  FindTopK<uint32_t         >(input, output_value, output_index, normalized_axis, kVal); break;
            case CV_32S:  FindTopK<int32_t          >(input, output_value, output_index, normalized_axis, kVal); break;
            case CV_32F:  FindTopK<float            >(input, output_value, output_index, normalized_axis, kVal); break;
            case CV_64U:  FindTopK<uint64_t         >(input, output_value, output_index, normalized_axis, kVal); break;
            case CV_64S:  FindTopK<int64_t          >(input, output_value, output_index, normalized_axis, kVal); break;
            case CV_64F:  FindTopK<double           >(input, output_value, output_index, normalized_axis, kVal); break;
            default: CV_Error(Error::BadDepth, "Unsupported input data type");
        }
    }
};

Ptr<TopK2Layer> TopK2Layer::create(const LayerParams& params)
{
    return makePtr<TopK2LayerImpl>(params);
}

}} // namespace cv::dnn