File size: 936 Bytes
e05eed1 98a67a0 ea7747b 98a67a0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 | # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
import torch
from nemotron_ocr_cpp import quad_non_maximal_suppression
quads = [
[[0, 2], [2, 2], [2, 0], [0, 0]],
[[1, 3], [3, 3], [3, 1], [1, 1]],
[[6, 6], [10, 6], [10, 5], [6, 5]],
[[6, 5.5], [8, 6], [10, 5.5], [8, 5]],
[[7, 6], [9, 6], [9, 5], [7, 5]],
]
quads = torch.tensor([[quads]], dtype=torch.float32).cuda()
probs = [0.7, 0.8, 0.55, 0.9, 0.85]
probs = torch.tensor([[probs]], dtype=torch.float32).cuda()
print("in_quads")
print(quads)
print("in_probs")
print(probs)
out_quads, out_probs, out_region_counts = quad_non_maximal_suppression(
quads, probs, prob_threshold=0.5, iou_threshold=0.1, kernel_height=2, kernel_width=2
)[:3]
print("out_quads")
print(out_quads)
print("out_probs")
print(out_probs)
print("out_region_counts")
print(out_region_counts)
|