Duplicate from nvidia/nemotron-ocr-v2

82f8ed3 1 day ago

6.94 kB

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


#include "quad_rectify_cpu.h"

#include <iostream>

#include "../geometry.h"
#include "quad_rectify_shared.h"

using namespace std;

template<typename quads_accessor_t, typename output_accessor_t, typename scalar_t>
void quad_rectify_calc_quad_width_impl(const quads_accessor_t &quads,
                                       output_accessor_t output,
                                       const scalar_t outputHeight,
                                       const scalar_t roundFactor,
                                       const scalar_t maxWidth)
{
    const int64_t numQuads = quads.size(0);

    for (int64_t quadIdx = 0; quadIdx < numQuads; ++quadIdx) {
        auto quadWidth = calc_quad_width(quads[quadIdx], outputHeight, roundFactor, maxWidth);

        output[quadIdx] = Convert<scalar_t, int64_t>::LeftToRight(quadWidth);
    }
}

template<typename quads_accessor_t, typename output_accessor_t, typename scalar_t>
void quad_rectify_cpu_forward_impl(const quads_accessor_t &quads,
                                   output_accessor_t output,
                                   const scalar_t imageHeight,
                                   const scalar_t imageWidth,
                                   bool isotropic)
{
    typedef Point_<scalar_t> Point_t;

    const int64_t numQuads = quads.size(0);
    const int64_t outputHeight = output.size(1);
    const int64_t outputWidth = output.size(2);

    for (int64_t quadIdx = 0; quadIdx < numQuads; ++quadIdx) {
        auto currQuad = quads[quadIdx];

        scalar_t quadWidth = isotropic ? calc_quad_width<scalar_t>(currQuad, outputHeight, 1, outputWidth) : scalar_t(outputWidth);

        for (int64_t row = 0; row < outputHeight; ++row) {
            for (int64_t col = 0; col < outputWidth; ++col) {
                Point_t outputPoint = calc_rect_value<scalar_t>(currQuad,
                                                                quadWidth,
                                                                outputHeight,
                                                                col,
                                                                row,
                                                                imageWidth,
                                                                imageHeight);

                auto currOutput = output[quadIdx][row][col];
                currOutput[0] = outputPoint.X;
                currOutput[1] = outputPoint.Y;
            }
        }
    }
}

/*template<typename scalar_t>
void quad_rectify_cpu_backward_impl(torch::Tensor quads,
                                    torch::Tensor gradOutput,
                                    torch::Tensor gradInput)
{
    const int64_t batchSize = gradOutput.size(0);
    const int64_t outputHeight = gradOutput.size(1);
    const int64_t outputWidth = gradOutput.size(2);
    auto gradInputAccess = gradInput.accessor<scalar_t, 3>();
    auto gradOutputAccess = gradOutput.accessor<scalar_t, 4>();
    for (int64_t batchIdx = 0; batchIdx < batchSize; ++batchIdx) {
        auto batchInputAccess = gradInputAccess[batchIdx];
        auto batchOutputAccess = gradOutputAccess[batchIdx];
        for (int64_t rowIdx = 0; rowIdx < outputHeight; ++rowIdx) {
            for (int64_t colIdx = 0; colIdx < outputWidth; ++colIdx) {
                const scalar_t fRow = scalar_t(rowIdx) / outputHeight;
                const scalar_t fCol = scalar_t(colIdx) / outputWidth;
                const scalar_t fRowCol = fRow * fCol;
                for (int64_t dim = 0; dim < 2; ++dim) {
                    const scalar_t dOut = batchOutputAccess[rowIdx][colIdx][dim];
                    const scalar_t gradIns[] = {
                        dOut * (fRowCol - fCol - fRow + 1),
                        dOut * (fCol - fRowCol),
                        dOut * fRowCol,
                        dOut * (fRow - fRowCol)
                    };
                    for (int64_t quadIdx = 0; quadIdx < 4; ++quadIdx) {
                        batchInputAccess[quadIdx][dim] += 2.0f * gradIns[quadIdx];
                    }
                }
            }
        }
    }
}*/

torch::Tensor quad_rectify_cpu_calc_quad_width(torch::Tensor quads,
                                               int64_t outputHeight,
                                               int64_t roundFactor,
                                               float maxWidth)
{
    auto output = torch::empty({ quads.size(0) },
                               quads.options().dtype(torch::kInt64));

    AT_DISPATCH_FLOATING_TYPES(
        quads.scalar_type(),
        "quad_rectify_cpu_calc_quad_width",
        ([&] {
            quad_rectify_calc_quad_width_impl(
                quads.accessor<scalar_t, 3>(),
                output.accessor<int64_t, 1>(),
                Convert<scalar_t, int64_t>::RightToLeft(outputHeight),
                Convert<scalar_t, int64_t>::RightToLeft(roundFactor),
                Convert<scalar_t, float>::RightToLeft(maxWidth)
            );
        })
    );

    return output;
}

torch::Tensor quad_rectify_cpu_forward(torch::Tensor quads,
                                       int64_t imageHeight,
                                       int64_t imageWidth,
                                       int64_t outputHeight,
                                       int64_t outputWidth,
                                       bool isotropic)
{
    auto output = torch::empty({ quads.size(0), outputHeight, outputWidth, 2 },
                               quads.options());

    AT_DISPATCH_FLOATING_TYPES(
        quads.scalar_type(),
        "quad_rectify_cpu_forward",
        ([&] {
            quad_rectify_cpu_forward_impl(
                quads.accessor<scalar_t, 3>(),
                output.accessor<scalar_t, 4>(),
                Convert<scalar_t, int64_t>::RightToLeft(imageHeight),
                Convert<scalar_t, int64_t>::RightToLeft(imageWidth),
                isotropic
            );
        })
    );

    return output;
}

torch::Tensor quad_rectify_cpu_backward(torch::Tensor quads,
                                        torch::Tensor gradOutput,
                                        int64_t imageHeight,
                                        int64_t imageWidth,
                                        bool isotropic)
{
    auto gradInput = torch::zeros_like(quads);

    throw std::runtime_error("Calling backward, and it's not implemented!");

    /*AT_DISPATCH_FLOATING_TYPES_AND_HALF(
        quads.scalar_type(),
        "quad_rectify_cpu_backward",
        ([&] {
            quad_rectify_cpu_backward_impl<scalar_t>(quads,
                                                     gradOutput,
                                                     gradInput);
        })
    );*/

    return gradInput;
}

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


	#include "quad_rectify_cpu.h"

	#include <iostream>

	#include "../geometry.h"
	#include "quad_rectify_shared.h"

	using namespace std;

	template<typename quads_accessor_t, typename output_accessor_t, typename scalar_t>
	void quad_rectify_calc_quad_width_impl(const quads_accessor_t &quads,
	output_accessor_t output,
	const scalar_t outputHeight,
	const scalar_t roundFactor,
	const scalar_t maxWidth)
	{
	const int64_t numQuads = quads.size(0);

	for (int64_t quadIdx = 0; quadIdx < numQuads; ++quadIdx) {
	auto quadWidth = calc_quad_width(quads[quadIdx], outputHeight, roundFactor, maxWidth);

	output[quadIdx] = Convert<scalar_t, int64_t>::LeftToRight(quadWidth);
	}
	}

	template<typename quads_accessor_t, typename output_accessor_t, typename scalar_t>
	void quad_rectify_cpu_forward_impl(const quads_accessor_t &quads,
	output_accessor_t output,
	const scalar_t imageHeight,
	const scalar_t imageWidth,
	bool isotropic)
	{
	typedef Point_<scalar_t> Point_t;

	const int64_t numQuads = quads.size(0);
	const int64_t outputHeight = output.size(1);
	const int64_t outputWidth = output.size(2);

	for (int64_t quadIdx = 0; quadIdx < numQuads; ++quadIdx) {
	auto currQuad = quads[quadIdx];

	scalar_t quadWidth = isotropic ? calc_quad_width<scalar_t>(currQuad, outputHeight, 1, outputWidth) : scalar_t(outputWidth);

	for (int64_t row = 0; row < outputHeight; ++row) {
	for (int64_t col = 0; col < outputWidth; ++col) {
	Point_t outputPoint = calc_rect_value<scalar_t>(currQuad,
	quadWidth,
	outputHeight,
	col,
	row,
	imageWidth,
	imageHeight);

	auto currOutput = output[quadIdx][row][col];
	currOutput[0] = outputPoint.X;
	currOutput[1] = outputPoint.Y;
	}
	}
	}
	}

	/*template<typename scalar_t>
	void quad_rectify_cpu_backward_impl(torch::Tensor quads,
	torch::Tensor gradOutput,
	torch::Tensor gradInput)
	{
	const int64_t batchSize = gradOutput.size(0);
	const int64_t outputHeight = gradOutput.size(1);
	const int64_t outputWidth = gradOutput.size(2);

	auto gradInputAccess = gradInput.accessor<scalar_t, 3>();
	auto gradOutputAccess = gradOutput.accessor<scalar_t, 4>();

	for (int64_t batchIdx = 0; batchIdx < batchSize; ++batchIdx) {
	auto batchInputAccess = gradInputAccess[batchIdx];
	auto batchOutputAccess = gradOutputAccess[batchIdx];

	for (int64_t rowIdx = 0; rowIdx < outputHeight; ++rowIdx) {
	for (int64_t colIdx = 0; colIdx < outputWidth; ++colIdx) {

	const scalar_t fRow = scalar_t(rowIdx) / outputHeight;
	const scalar_t fCol = scalar_t(colIdx) / outputWidth;
	const scalar_t fRowCol = fRow * fCol;

	for (int64_t dim = 0; dim < 2; ++dim) {
	const scalar_t dOut = batchOutputAccess[rowIdx][colIdx][dim];

	const scalar_t gradIns[] = {
	dOut * (fRowCol - fCol - fRow + 1),
	dOut * (fCol - fRowCol),
	dOut * fRowCol,
	dOut * (fRow - fRowCol)
	};

	for (int64_t quadIdx = 0; quadIdx < 4; ++quadIdx) {
	batchInputAccess[quadIdx][dim] += 2.0f * gradIns[quadIdx];
	}
	}
	}
	}
	}
	}*/

	torch::Tensor quad_rectify_cpu_calc_quad_width(torch::Tensor quads,
	int64_t outputHeight,
	int64_t roundFactor,
	float maxWidth)
	{
	auto output = torch::empty({ quads.size(0) },
	quads.options().dtype(torch::kInt64));

	AT_DISPATCH_FLOATING_TYPES(
	quads.scalar_type(),
	"quad_rectify_cpu_calc_quad_width",
	([&] {
	quad_rectify_calc_quad_width_impl(
	quads.accessor<scalar_t, 3>(),
	output.accessor<int64_t, 1>(),
	Convert<scalar_t, int64_t>::RightToLeft(outputHeight),
	Convert<scalar_t, int64_t>::RightToLeft(roundFactor),
	Convert<scalar_t, float>::RightToLeft(maxWidth)
	);
	})
	);

	return output;
	}

	torch::Tensor quad_rectify_cpu_forward(torch::Tensor quads,
	int64_t imageHeight,
	int64_t imageWidth,
	int64_t outputHeight,
	int64_t outputWidth,
	bool isotropic)
	{
	auto output = torch::empty({ quads.size(0), outputHeight, outputWidth, 2 },
	quads.options());

	AT_DISPATCH_FLOATING_TYPES(
	quads.scalar_type(),
	"quad_rectify_cpu_forward",
	([&] {
	quad_rectify_cpu_forward_impl(
	quads.accessor<scalar_t, 3>(),
	output.accessor<scalar_t, 4>(),
	Convert<scalar_t, int64_t>::RightToLeft(imageHeight),
	Convert<scalar_t, int64_t>::RightToLeft(imageWidth),
	isotropic
	);
	})
	);

	return output;
	}

	torch::Tensor quad_rectify_cpu_backward(torch::Tensor quads,
	torch::Tensor gradOutput,
	int64_t imageHeight,
	int64_t imageWidth,
	bool isotropic)
	{
	auto gradInput = torch::zeros_like(quads);

	throw std::runtime_error("Calling backward, and it's not implemented!");

	/*AT_DISPATCH_FLOATING_TYPES_AND_HALF(
	quads.scalar_type(),
	"quad_rectify_cpu_backward",
	([&] {
	quad_rectify_cpu_backward_impl<scalar_t>(quads,
	gradOutput,
	gradInput);
	})
	);*/

	return gradInput;
	}