Spaces:
Running
Running
File size: 1,512 Bytes
a80d6bb |
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 |
import torch
import torch.distributed as dist
import numpy as np
import cv2
def parse_pair_seq(pair_num_list):
#generate pair_seq_list: [#pair_num]:seq
# accu_pair_num: dict{seq_name:accumulated_pair}
pair_num=int(pair_num_list[0,1])
pair_num_list=pair_num_list[1:]
pair_seq_list=[]
cursor=0
accu_pair_num={}
for line in pair_num_list:
seq,seq_pair_num=line[0],int(line[1])
for _ in range(seq_pair_num):
pair_seq_list.append(seq)
accu_pair_num[seq]=cursor
cursor+=seq_pair_num
assert pair_num==cursor
return pair_seq_list,accu_pair_num
def tocuda(data):
# convert tensor data in dictionary to cuda when it is a tensor
for key in data.keys():
if type(data[key]) == torch.Tensor:
data[key] = data[key].cuda()
return data
def reduce_tensor(tensor,op='mean'):
rt = tensor.detach()
dist.all_reduce(rt, op=dist.ReduceOp.SUM)
if op=='mean':
rt /= dist.get_world_size()
return rt
def get_rnd_homography(batch_size, pert_ratio=0.25):
corners = np.array([[-1, 1], [1, 1], [-1, -1], [1, -1]], dtype=np.float32)
homo_tower = []
for _ in range(batch_size):
rnd_pert = np.random.uniform(-2 * pert_ratio, 2 * pert_ratio, (4, 2)).astype(np.float32)
pert_corners = corners + rnd_pert
M = cv2.getPerspectiveTransform(corners, pert_corners)
homo_tower.append(M)
homo_tower = np.stack(homo_tower, axis=0)
return homo_tower
|