|
import torch |
|
import torch.distributed as dist |
|
import numpy as np |
|
import cv2 |
|
|
|
def parse_pair_seq(pair_num_list): |
|
|
|
|
|
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): |
|
|
|
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 |
|
|