Spaces:
Running on Zero
Running on Zero
File size: 2,825 Bytes
4b35c4e | 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 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 | import os
import numpy as np
import torch
import torch.nn.functional as F
import torch.distributed as dist
def get_padding(orig_H, orig_W):
""" returns how the input of shape (orig_H, orig_W) should be padded
this ensures that both H and W are divisible by 32
"""
if orig_W % 32 == 0:
l = 0
r = 0
else:
new_W = 32 * ((orig_W // 32) + 1)
l = (new_W - orig_W) // 2
r = (new_W - orig_W) - l
if orig_H % 32 == 0:
t = 0
b = 0
else:
new_H = 32 * ((orig_H // 32) + 1)
t = (new_H - orig_H) // 2
b = (new_H - orig_H) - t
return l, r, t, b
def pad_input(img, intrins, lrtb=(0,0,0,0)):
""" pad input image
img should be a torch tensor of shape (B, 3, H, W)
intrins should be a torch tensor of shape (B, 3, 3)
"""
l, r, t, b = lrtb
if l+r+t+b != 0:
pad_value_R = (0 - 0.485) / 0.229
pad_value_G = (0 - 0.456) / 0.224
pad_value_B = (0 - 0.406) / 0.225
img_R = F.pad(img[:,0:1,:,:], (l, r, t, b), mode="constant", value=pad_value_R)
img_G = F.pad(img[:,1:2,:,:], (l, r, t, b), mode="constant", value=pad_value_G)
img_B = F.pad(img[:,2:3,:,:], (l, r, t, b), mode="constant", value=pad_value_B)
img = torch.cat([img_R, img_G, img_B], dim=1)
if intrins is not None:
intrins[:, 0, 2] += l
intrins[:, 1, 2] += t
return img, intrins
def compute_normal_error(pred_norm, gt_norm):
""" compute per-pixel surface normal error in degrees
NOTE: pred_norm and gt_norm should be torch tensors of shape (B, 3, ...)
"""
pred_error = torch.cosine_similarity(pred_norm, gt_norm, dim=1)
pred_error = torch.clamp(pred_error, min=-1.0, max=1.0)
pred_error = torch.acos(pred_error) * 180.0 / np.pi
pred_error = pred_error.unsqueeze(1) # (B, 1, ...)
return pred_error
def compute_normal_metrics(total_normal_errors):
""" compute surface normal metrics (used for benchmarking)
NOTE: total_normal_errors should be a 1D torch tensor of errors in degrees
"""
total_normal_errors = total_normal_errors.detach().cpu().numpy()
num_pixels = total_normal_errors.shape[0]
metrics = {
'mean': np.average(total_normal_errors),
'median': np.median(total_normal_errors),
'rmse': np.sqrt(np.sum(total_normal_errors * total_normal_errors) / num_pixels),
'a1': 100.0 * (np.sum(total_normal_errors < 5) / num_pixels),
'a2': 100.0 * (np.sum(total_normal_errors < 7.5) / num_pixels),
'a3': 100.0 * (np.sum(total_normal_errors < 11.25) / num_pixels),
'a4': 100.0 * (np.sum(total_normal_errors < 22.5) / num_pixels),
'a5': 100.0 * (np.sum(total_normal_errors < 30) / num_pixels)
}
return metrics |