Spaces:
Running
Running
| 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 | |