| |
| |
|
|
| 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) |
|
|
