Duplicate from nvidia/nemotron-ocr-v2

82f8ed3 1 day ago

6.92 kB

// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: Apache-2.0

#include "non_maximal_suppression.h"

#include <algorithm>
#include "../geometry.h"

using namespace std;


template<typename scalar_t>
void visit_node(
    const torch::TensorAccessor<scalar_t, 4> &quads,
    const torch::TensorAccessor<scalar_t, 2> &probs,
    const torch::TensorAccessor<int32_t, 3> &adjacency,
    MergeQuad_<scalar_t> &mQuad,
    unordered_set<int32_t> &visited,
    int64_t r, int64_t c, int32_t vIdx)
{
    if (visited.count(vIdx)) {
        return;
    }
    visited.insert(vIdx);

    int32_t *pAdj = adjacency[r][c].data();

    int32_t adjCt = pAdj[0];
    assert(adjCt > 0);

    mQuad.Append(Quad_<scalar_t>(quads[r][c].data()), probs[r][c]);

    int32_t *pOff = pAdj + 2;
    int32_t *pEnd = pAdj + adjCt + 1;

    const int32_t W = quads.size(1);

    for (; pOff != pEnd; ++pOff) {
        int32_t vIdx2 = *pOff;
        int32_t r2 = vIdx2 / W;
        int32_t c2 = vIdx2 % W;

        visit_node(quads, probs, adjacency, mQuad, visited, r2, c2, vIdx2);
    }
}

template<typename scalar_t>
std::vector<torch::Tensor> quad_nms_from_adjacency_impl(
    const torch::TensorAccessor<scalar_t, 5> &quads,
    const torch::TensorAccessor<scalar_t, 3> &probs,
    const torch::TensorAccessor<int32_t, 4> &adjacency,
    scalar_t probThreshold, scalar_t iouThreshold,
    int64_t maxRegions)
{
    const uint64_t B = quads.size((int)0);
    const int64_t H = quads.size((int)1);
    const int64_t W = quads.size((int)2);

    typedef MergeQuad_<scalar_t> MQuad;
    typedef EmbedQuad_<scalar_t> EFQuad;

    vector<vector<EFQuad>> batchQuads{ static_cast< const unsigned int >( B ) };
    vector<vector<EFQuad>> allQuads{ static_cast< const unsigned int >( B ) };
    vector<vector<vector<size_t>>> batchAdjIdxs{ static_cast< const unsigned int >( B ) };

    #pragma omp parallel num_threads (8)
    {
        #pragma omp for
        for (int64_t b = 0; b < B; ++b) {
            unordered_set<int32_t> visited;

            for (int64_t r = 0; r < H; ++r) {
                for (int64_t c = 0; c < W; ++c) {
                    auto currProb = probs[b][r][c];

                    if (currProb < probThreshold) {
                        continue;
                    }

                    int32_t vIdx = r * W + c;

                    // Ensure that this quad hasn't already been merged
                    if (visited.count(vIdx)) {
                        continue;
                    }

                    MQuad mQuad{ZeroInitTag{}};
                    visit_node(quads[b], probs[b], adjacency[b], mQuad, visited, r, c, vIdx);

                    batchQuads[b].push_back(mQuad.Commit());
                }
            }
        }

        #pragma omp single
        {
            for (size_t b = 0; b < B; ++b) {
                size_t numQuads = batchQuads[b].size();
                batchAdjIdxs[b].resize(numQuads);
                for (int64_t n = 0; n < numQuads; ++n) {
                    #pragma omp task default(none) shared(batchAdjIdxs, batchQuads, iouThreshold) firstprivate(b, numQuads, n)
                    {
                        for (int64_t m = n + 1; m < numQuads; ++m) {
                            vector<size_t> &adjIdxs = batchAdjIdxs[b][n];
                            vector<EFQuad> &quads = batchQuads[b];
                            auto iou = quads[n].IOU(quads[m]);

                            if (iou > iouThreshold) {
                                adjIdxs.push_back(m);
                            }
                        }
                    }
                }
            }

            #pragma omp taskwait
        }

        #pragma omp for
        for (int64_t batchIdx = 0; batchIdx < B; ++batchIdx) {
            vector<vector<size_t>> &adjIdxs = batchAdjIdxs[batchIdx];
            vector<EFQuad> &quads = batchQuads[batchIdx];
            vector<EFQuad> &finalQuads = allQuads[batchIdx];

            // Step 3: Using depth first search, merge the regions
            unordered_set<size_t> visited;
            for (int64_t n = 0; n < quads.size(); ++n) {
                EFQuad currQuad;
                visit_node(quads, n, adjIdxs, currQuad, visited);

                if (currQuad.NumQuads > 0) {
                    currQuad.Prepare();

                    finalQuads.push_back(currQuad);
                }
            }

            // Only sort the part that we want to keep
            partial_sort(begin(finalQuads),
                        begin(finalQuads) + std::min<int64_t>(finalQuads.size(), maxRegions),
                        end(finalQuads),
                [] (auto a, auto b) {
                    return a.Confidence > b.Confidence;
                }
            );

            // Truncate the low confidence regions
            if (finalQuads.size() > maxRegions) {
                finalQuads.resize(maxRegions);
            }

            //cout << "Ex " << batchIdx << " quads:" << endl << finalQuads << endl << endl;
        }

    } // End parallel

    int64_t numOutQuads = 0;
    for (int64_t batchIdx = 0; batchIdx < B; ++batchIdx) {
        numOutQuads += allQuads[batchIdx].size();
    }

    // Step 4: Convert the quads into tensor representation
    auto outQuadTensor = torch::empty({ numOutQuads, 4, 2 }, torch::kFloat32);
    auto outConfTensor = torch::empty({ numOutQuads }, torch::kFloat32);
    torch::Tensor outCountTensor = torch::empty({ static_cast<int64_t>( allQuads.size() ) }, torch::kInt64);

    auto outQuadAccess = outQuadTensor.accessor<float, 3>();
    auto outConfAccess = outConfTensor.accessor<float, 1>();
    auto outCountAccess = outCountTensor.accessor<int64_t, 1>();

    int64_t offset = 0;
    for (int64_t batchIdx = 0; batchIdx < allQuads.size(); ++batchIdx) {
        vector<EFQuad> &exQuads = allQuads[batchIdx];

        outCountAccess[batchIdx] = exQuads.size();

        for (int64_t qIdx = 0; qIdx < exQuads.size(); ++qIdx, ++offset) {
            copy_quad(exQuads[qIdx], outQuadAccess[offset].data());
            outConfAccess[offset] = exQuads[qIdx].Confidence;
        }
    }

    return { outQuadTensor, outConfTensor, outCountTensor };
}

std::vector<torch::Tensor> quad_nms_from_adjacency(
    torch::Tensor quads, torch::Tensor probs, torch::Tensor adjacency,
    float probThreshold, float iouThreshold,
    int64_t maxRegions)
{
    std::vector<torch::Tensor> ret;

    AT_DISPATCH_FLOATING_TYPES(
        quads.scalar_type(),
        "quad_nms_from_adjacency",
        ([&] {
            ret = quad_nms_from_adjacency_impl<scalar_t>(
                quads.accessor<scalar_t, 5>(),
                probs.accessor<scalar_t, 3>(),
                adjacency.accessor<int32_t, 4>(),
                probThreshold, iouThreshold,
                maxRegions
            );
        })
    );

    return ret;
}

	// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
	// SPDX-License-Identifier: Apache-2.0

	#include "non_maximal_suppression.h"

	#include <algorithm>
	#include "../geometry.h"

	using namespace std;


	template<typename scalar_t>
	void visit_node(
	const torch::TensorAccessor<scalar_t, 4> &quads,
	const torch::TensorAccessor<scalar_t, 2> &probs,
	const torch::TensorAccessor<int32_t, 3> &adjacency,
	MergeQuad_<scalar_t> &mQuad,
	unordered_set<int32_t> &visited,
	int64_t r, int64_t c, int32_t vIdx)
	{
	if (visited.count(vIdx)) {
	return;
	}
	visited.insert(vIdx);

	int32_t *pAdj = adjacency[r][c].data();

	int32_t adjCt = pAdj[0];
	assert(adjCt > 0);

	mQuad.Append(Quad_<scalar_t>(quads[r][c].data()), probs[r][c]);

	int32_t *pOff = pAdj + 2;
	int32_t *pEnd = pAdj + adjCt + 1;

	const int32_t W = quads.size(1);

	for (; pOff != pEnd; ++pOff) {
	int32_t vIdx2 = *pOff;
	int32_t r2 = vIdx2 / W;
	int32_t c2 = vIdx2 % W;

	visit_node(quads, probs, adjacency, mQuad, visited, r2, c2, vIdx2);
	}
	}

	template<typename scalar_t>
	std::vector<torch::Tensor> quad_nms_from_adjacency_impl(
	const torch::TensorAccessor<scalar_t, 5> &quads,
	const torch::TensorAccessor<scalar_t, 3> &probs,
	const torch::TensorAccessor<int32_t, 4> &adjacency,
	scalar_t probThreshold, scalar_t iouThreshold,
	int64_t maxRegions)
	{
	const uint64_t B = quads.size((int)0);
	const int64_t H = quads.size((int)1);
	const int64_t W = quads.size((int)2);

	typedef MergeQuad_<scalar_t> MQuad;
	typedef EmbedQuad_<scalar_t> EFQuad;

	vector<vector<EFQuad>> batchQuads{ static_cast< const unsigned int >( B ) };
	vector<vector<EFQuad>> allQuads{ static_cast< const unsigned int >( B ) };
	vector<vector<vector<size_t>>> batchAdjIdxs{ static_cast< const unsigned int >( B ) };

	#pragma omp parallel num_threads (8)
	{
	#pragma omp for
	for (int64_t b = 0; b < B; ++b) {
	unordered_set<int32_t> visited;

	for (int64_t r = 0; r < H; ++r) {
	for (int64_t c = 0; c < W; ++c) {
	auto currProb = probs[b][r][c];

	if (currProb < probThreshold) {
	continue;
	}

	int32_t vIdx = r * W + c;

	// Ensure that this quad hasn't already been merged
	if (visited.count(vIdx)) {
	continue;
	}

	MQuad mQuad{ZeroInitTag{}};
	visit_node(quads[b], probs[b], adjacency[b], mQuad, visited, r, c, vIdx);

	batchQuads[b].push_back(mQuad.Commit());
	}
	}
	}

	#pragma omp single
	{
	for (size_t b = 0; b < B; ++b) {
	size_t numQuads = batchQuads[b].size();
	batchAdjIdxs[b].resize(numQuads);
	for (int64_t n = 0; n < numQuads; ++n) {
	#pragma omp task default(none) shared(batchAdjIdxs, batchQuads, iouThreshold) firstprivate(b, numQuads, n)
	{
	for (int64_t m = n + 1; m < numQuads; ++m) {
	vector<size_t> &adjIdxs = batchAdjIdxs[b][n];
	vector<EFQuad> &quads = batchQuads[b];
	auto iou = quads[n].IOU(quads[m]);

	if (iou > iouThreshold) {
	adjIdxs.push_back(m);
	}
	}
	}
	}
	}

	#pragma omp taskwait
	}

	#pragma omp for
	for (int64_t batchIdx = 0; batchIdx < B; ++batchIdx) {
	vector<vector<size_t>> &adjIdxs = batchAdjIdxs[batchIdx];
	vector<EFQuad> &quads = batchQuads[batchIdx];
	vector<EFQuad> &finalQuads = allQuads[batchIdx];

	// Step 3: Using depth first search, merge the regions
	unordered_set<size_t> visited;
	for (int64_t n = 0; n < quads.size(); ++n) {
	EFQuad currQuad;
	visit_node(quads, n, adjIdxs, currQuad, visited);

	if (currQuad.NumQuads > 0) {
	currQuad.Prepare();

	finalQuads.push_back(currQuad);
	}
	}

	// Only sort the part that we want to keep
	partial_sort(begin(finalQuads),
	begin(finalQuads) + std::min<int64_t>(finalQuads.size(), maxRegions),
	end(finalQuads),
	[] (auto a, auto b) {
	return a.Confidence > b.Confidence;
	}
	);

	// Truncate the low confidence regions
	if (finalQuads.size() > maxRegions) {
	finalQuads.resize(maxRegions);
	}

	//cout << "Ex " << batchIdx << " quads:" << endl << finalQuads << endl << endl;
	}

	} // End parallel

	int64_t numOutQuads = 0;
	for (int64_t batchIdx = 0; batchIdx < B; ++batchIdx) {
	numOutQuads += allQuads[batchIdx].size();
	}

	// Step 4: Convert the quads into tensor representation
	auto outQuadTensor = torch::empty({ numOutQuads, 4, 2 }, torch::kFloat32);
	auto outConfTensor = torch::empty({ numOutQuads }, torch::kFloat32);
	torch::Tensor outCountTensor = torch::empty({ static_cast<int64_t>( allQuads.size() ) }, torch::kInt64);

	auto outQuadAccess = outQuadTensor.accessor<float, 3>();
	auto outConfAccess = outConfTensor.accessor<float, 1>();
	auto outCountAccess = outCountTensor.accessor<int64_t, 1>();

	int64_t offset = 0;
	for (int64_t batchIdx = 0; batchIdx < allQuads.size(); ++batchIdx) {
	vector<EFQuad> &exQuads = allQuads[batchIdx];

	outCountAccess[batchIdx] = exQuads.size();

	for (int64_t qIdx = 0; qIdx < exQuads.size(); ++qIdx, ++offset) {
	copy_quad(exQuads[qIdx], outQuadAccess[offset].data());
	outConfAccess[offset] = exQuads[qIdx].Confidence;
	}
	}

	return { outQuadTensor, outConfTensor, outCountTensor };
	}

	std::vector<torch::Tensor> quad_nms_from_adjacency(
	torch::Tensor quads, torch::Tensor probs, torch::Tensor adjacency,
	float probThreshold, float iouThreshold,
	int64_t maxRegions)
	{
	std::vector<torch::Tensor> ret;

	AT_DISPATCH_FLOATING_TYPES(
	quads.scalar_type(),
	"quad_nms_from_adjacency",
	([&] {
	ret = quad_nms_from_adjacency_impl<scalar_t>(
	quads.accessor<scalar_t, 5>(),
	probs.accessor<scalar_t, 3>(),
	adjacency.accessor<int32_t, 4>(),
	probThreshold, iouThreshold,
	maxRegions
	);
	})
	);

	return ret;
	}