| import numpy as np
|
| import torch
|
| import torch.distributed as dist
|
|
|
|
|
| def reduce(tensor, world_size):
|
| if isinstance(tensor, torch.Tensor):
|
| tensor = tensor.clone()
|
| dist.all_reduce(tensor, dist.ReduceOp.SUM)
|
| tensor.div_(world_size)
|
| return tensor
|
|
|
|
|
| def expand(mask, num=1):
|
|
|
|
|
|
|
| for _ in range(num):
|
| mask[..., 1:, :] = mask[..., 1:, :] | mask[..., :-1, :]
|
| mask[..., :-1, :] = mask[..., :-1, :] | mask[..., 1:, :]
|
| mask[..., :, 1:] = mask[..., :, 1:] | mask[..., :, :-1]
|
| mask[..., :, :-1] = mask[..., :, :-1] | mask[..., :, 1:]
|
| return mask
|
|
|
|
|
| def differentiate(mask):
|
|
|
|
|
|
|
| diff = torch.zeros_like(mask).bool()
|
| diff_y = mask[..., 1:, :] != mask[..., :-1, :]
|
| diff_x = mask[..., :, 1:] != mask[..., :, :-1]
|
| diff[..., 1:, :] = diff[..., 1:, :] | diff_y
|
| diff[..., :-1, :] = diff[..., :-1, :] | diff_y
|
| diff[..., :, 1:] = diff[..., :, 1:] | diff_x
|
| diff[..., :, :-1] = diff[..., :, :-1] | diff_x
|
| return diff
|
|
|
|
|
| def sample_points(step, boundaries, num_samples):
|
| if boundaries.ndim == 3:
|
| points = []
|
| for boundaries_k in boundaries:
|
| points_k = sample_points(step, boundaries_k, num_samples)
|
| points.append(points_k)
|
| points = torch.stack(points)
|
| else:
|
| H, W = boundaries.shape
|
| boundary_points, _ = sample_mask_points(step, boundaries, num_samples // 2)
|
| num_boundary_points = boundary_points.shape[0]
|
| num_random_points = num_samples - num_boundary_points
|
| random_points = sample_random_points(step, H, W, num_random_points)
|
| random_points = random_points.to(boundary_points.device)
|
| points = torch.cat((boundary_points, random_points), dim=0)
|
| return points
|
|
|
|
|
| def sample_mask_points(step, mask, num_points):
|
| num_nonzero = int(mask.sum())
|
| i, j = torch.nonzero(mask, as_tuple=True)
|
| if num_points < num_nonzero:
|
| sample = np.random.choice(num_nonzero, size=num_points, replace=False)
|
| i, j = i[sample], j[sample]
|
| t = torch.ones_like(i) * step
|
| x, y = j, i
|
| points = torch.stack((t, x, y), dim=-1)
|
| return points.float(), (i, j)
|
|
|
|
|
| def sample_random_points(step, height, width, num_points):
|
| x = torch.randint(width, size=[num_points])
|
| y = torch.randint(height, size=[num_points])
|
| t = torch.ones(num_points) * step
|
| points = torch.stack((t, x, y), dim=-1)
|
| return points.float()
|
|
|
|
|
| def get_grid(height, width, shape=None, dtype="torch", device="cpu", align_corners=True, normalize=True):
|
| H, W = height, width
|
| S = shape if shape else []
|
| if align_corners:
|
| x = torch.linspace(0, 1, W, device=device)
|
| y = torch.linspace(0, 1, H, device=device)
|
| if not normalize:
|
| x = x * (W - 1)
|
| y = y * (H - 1)
|
| else:
|
| x = torch.linspace(0.5 / W, 1.0 - 0.5 / W, W, device=device)
|
| y = torch.linspace(0.5 / H, 1.0 - 0.5 / H, H, device=device)
|
| if not normalize:
|
| x = x * W
|
| y = y * H
|
| x_view, y_view, exp = [1 for _ in S] + [1, -1], [1 for _ in S] + [-1, 1], S + [H, W]
|
| x = x.view(*x_view).expand(*exp)
|
| y = y.view(*y_view).expand(*exp)
|
| grid = torch.stack([x, y], dim=-1)
|
| if dtype == "numpy":
|
| grid = grid.numpy()
|
| return grid
|
|
|
|
|
| def get_sobel_kernel(kernel_size):
|
| K = kernel_size
|
| sobel = torch.tensor(list(range(K))) - K // 2
|
| sobel_x, sobel_y = sobel.view(-1, 1), sobel.view(1, -1)
|
| sum_xy = sobel_x ** 2 + sobel_y ** 2
|
| sum_xy[sum_xy == 0] = 1
|
| sobel_x, sobel_y = sobel_x / sum_xy, sobel_y / sum_xy
|
| sobel_kernel = torch.stack([sobel_x.unsqueeze(0), sobel_y.unsqueeze(0)], dim=0)
|
| return sobel_kernel
|
|
|
|
|
| def to_device(data, device):
|
| data = {k: v.to(device) for k, v in data.items()}
|
| return data
|
|
|
|
|
| def get_alpha_consistency(bflow, fflow, thresh_1=0.01, thresh_2=0.5, thresh_mul=1):
|
| norm = lambda x: x.pow(2).sum(dim=-1).sqrt()
|
| B, H, W, C = bflow.shape
|
|
|
| mag = norm(fflow) + norm(bflow)
|
| grid = get_grid(H, W, shape=[B], device=fflow.device)
|
| grid[..., 0] = grid[..., 0] + bflow[..., 0] / (W - 1)
|
| grid[..., 1] = grid[..., 1] + bflow[..., 1] / (H - 1)
|
| grid = grid * 2 - 1
|
| fflow_warped = torch.nn.functional.grid_sample(fflow.permute(0, 3, 1, 2), grid, mode="bilinear", align_corners=True)
|
| flow_diff = bflow + fflow_warped.permute(0, 2, 3, 1)
|
| occ_thresh = thresh_1 * mag + thresh_2
|
| occ_thresh = occ_thresh * thresh_mul
|
| alpha = norm(flow_diff) < occ_thresh
|
| alpha = alpha.float()
|
| return alpha |
