| # SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| import shutil | |
| import torch | |
| import torch.nn.functional as F | |
| from torch.autograd import gradcheck | |
| from torchvision.transforms import ToTensor, ToPILImage | |
| from PIL import Image | |
| from nemotron_ocr.inference.post_processing.quad_rectify import QuadRectify | |
| example = 3 | |
| isotropic = True | |
| # root_dir = '/mnt/fsx-datasets-a-1-new/mranzinger/ocr/scene-text/icdar/incidental_text/train' | |
| root_dir = "/home/dcg-adlr-mranzinger-data.cosmos1100/ocr/scene-text/icdar/incidental_text/train" | |
| # root_dir = '/raid/local_datasets/scene-text/icdar/focused_text/train' | |
| image_path = "{}/images/img_{}.jpg".format(root_dir, example) | |
| label_path = "{}/gt/gt_img_{}.txt".format(root_dir, example) | |
| shutil.copyfile(image_path, "original.jpg") | |
| image = Image.open(image_path) | |
| image = ToTensor()(image) | |
| print("image", image.shape) | |
| quads = [] | |
| with open(label_path, "r") as fd: | |
| for i, line in enumerate(fd.readlines()): | |
| line = line.strip() | |
| print(line) | |
| # Skip the unicode character | |
| if i == 0: | |
| line = line[1:] | |
| line = line.split(",") | |
| coords = [float(t) for t in line[:8]] | |
| word = line[-1] | |
| if word != "###": | |
| quads.append(coords) | |
| quads = torch.tensor(quads).reshape(-1, 4, 2) | |
| # quads[:,:,0] /= image.shape[-1] | |
| # quads[:,:,1] /= image.shape[-2] | |
| # print(quads) | |
| image = image.unsqueeze(0).repeat(quads.shape[0], 1, 1, 1) | |
| image = image.cuda() | |
| quads = quads.cuda() | |
| aspect = image.shape[-1] / image.shape[-2] | |
| qr = QuadRectify(60, 400, 0, isotropic=isotropic) | |
| # qr = QuadRectify(60, -1000, aspect) | |
| grid = qr(quads, *image.shape[2:]) | |
| print("grid", grid.shape) | |
| resampled = F.grid_sample(image, grid, align_corners=False) | |
| print(resampled.shape) | |
| resampled = resampled.permute(1, 0, 2, 3).contiguous().reshape(3, -1, resampled.shape[-1]) | |
| resampled = resampled.cpu() | |
| pilOutput = ToPILImage()(resampled) | |
| pilOutput.save("rectified.jpg") | |
| # sys.exit(0) | |
| print("checking gradients") | |
| quads = quads.double() | |
| cuda_quads = quads.cuda() | |
| quads.requires_grad_() | |
| cuda_quads.requires_grad_() | |
| qr = QuadRectify(9, 12, isotropic=isotropic) | |
| print("check GPU gradients:") | |
| # gradcheck(qr.forward, cuda_quads, eps=1e-3, atol=1e-3) | |
| gradcheck(qr.forward, (cuda_quads, *image.shape[2:])) | |
| print("check CPU gradients:") | |
| gradcheck(qr.forward, (quads, *image.shape[2:])) | |
| # #gradcheck(qr.forward, quads, eps=1e-3, atol=1e-3) | |
| # gradcheck(qr.forward, quads) | |
| print("checking performance...") | |
| quads = quads.detach() | |
| quads = torch.rand(256, 4, 2, dtype=torch.float64) | |
| cuda_quads = quads.cuda() | |
| qr = QuadRectify(8, 60, isotropic=isotropic) | |
| # num_passes = 500 | |
| # for target_device in [torch.device('cuda:0'), torch.device('cpu')]: | |
| # for target_type in [torch.float16, torch.float32, torch.float64]: | |
| # test_quads = quads.to(target_type).to(target_device).clone() | |
| # test_quads.requires_grad_() | |
| # print('type:', target_type, 'device:', target_device) | |
| # try: | |
| # with torch.autograd.profiler.profile(use_cuda=test_quads.is_cuda) as prof: | |
| # for _ in range(num_passes): | |
| # g = qr(test_quads) | |
| # l = (1 - g.mean()) | |
| # l.backward() | |
| # print(prof.key_averages().table(sort_by="cuda_time_total", row_limit=10)) | |
| # # print('\tforward...') | |
| # # torch.cuda.synchronize() | |
| # # start_time = time.time() | |
| # # for i in range(num_passes): | |
| # # qr(test_quads).sum() | |
| # # torch.cuda.synchronize() | |
| # # end_time = time.time() | |
| # # fp_sec_per_call = (end_time - start_time) / num_passes | |
| # # print('\t\tTook', fp_sec_per_call, "sec/call") | |
| # # print('\tbackward...') | |
| # # torch.cuda.synchronize() | |
| # # start_time = time.time() | |
| # # for i in range(num_passes): | |
| # # g = qr(test_quads) | |
| # # g.sum().backward() | |
| # # torch.cuda.synchronize() | |
| # # end_time = time.time() | |
| # # bp_sec_per_call = (end_time - start_time) / num_passes - fp_sec_per_call | |
| # # print('\t\tTook', bp_sec_per_call, "sec/call") | |
| # except Exception as e: | |
| # print('\t\t', e) | |
| # # for target_type in [torch.float16, torch.float32]: | |
| # # print('DType:', target_type) | |
| # # h_cuda_quads = cuda_quads.clone().to(target_type) | |
| # # h_cuda_quads.requires_grad_() | |
| # # v = qr(h_cuda_quads).sum() | |
| # # v.backward() | |
| # # print(h_cuda_quads.grad) | |