Spaces:
Runtime error
Runtime error
| import torch | |
| import numpy as np | |
| from .sobel2 import SobelLayer, SeperateSobelLayer | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| def image_warp(image, flow): | |
| ''' | |
| image: 上一帧的图片,torch.Size([1, 3, 256, 256]) | |
| flow: 光流, torch.Size([1, 2, 256, 256]) | |
| final_grid: torch.Size([1, 2, 256, 256]) | |
| ''' | |
| b, c, h, w = image.size() | |
| device = image.device | |
| flow = torch.cat([flow[:, 0:1, :, :] / ((w - 1.0) / 2.0), flow[:, 1:2, :, :] / ((h - 1.0) / 2.0)], | |
| dim=1) # normalize to [-1~1](from upper left to lower right | |
| flow = flow.permute(0, 2, 3, | |
| 1) # if you wanna use grid_sample function, the channel(band) shape of show must be in the last dimension | |
| x = np.linspace(-1, 1, w) | |
| y = np.linspace(-1, 1, h) | |
| X, Y = np.meshgrid(x, y) | |
| grid = torch.cat((torch.from_numpy(X.astype('float32')).unsqueeze(0).unsqueeze(3), | |
| torch.from_numpy(Y.astype('float32')).unsqueeze(0).unsqueeze(3)), 3).to(device) | |
| output = torch.nn.functional.grid_sample(image, grid + flow, mode='bilinear', padding_mode='zeros') | |
| return output | |
| def length_sq(x): | |
| return torch.sum(torch.square(x), dim=1, keepdim=True) | |
| def fbConsistencyCheck(flow_fw, flow_bw, alpha1=0.01, alpha2=0.5): | |
| flow_bw_warped = image_warp(flow_bw, flow_fw) # wb(wf(x)) | |
| flow_fw_warped = image_warp(flow_fw, flow_bw) # wf(wb(x)) | |
| flow_diff_fw = flow_fw + flow_bw_warped # wf + wb(wf(x)) | |
| flow_diff_bw = flow_bw + flow_fw_warped # wb + wf(wb(x)) | |
| mag_sq_fw = length_sq(flow_fw) + length_sq(flow_bw_warped) # |wf| + |wb(wf(x))| | |
| mag_sq_bw = length_sq(flow_bw) + length_sq(flow_fw_warped) # |wb| + |wf(wb(x))| | |
| occ_thresh_fw = alpha1 * mag_sq_fw + alpha2 | |
| occ_thresh_bw = alpha1 * mag_sq_bw + alpha2 | |
| fb_occ_fw = (length_sq(flow_diff_fw) > occ_thresh_fw).float() | |
| fb_occ_bw = (length_sq(flow_diff_bw) > occ_thresh_bw).float() | |
| return fb_occ_fw, fb_occ_bw # fb_occ_fw -> frame2 area occluded by frame1, fb_occ_bw -> frame1 area occluded by frame2 | |
| def rgb2gray(image): | |
| gray_image = image[:, 0] * 0.299 + image[:, 1] * 0.587 + 0.110 * image[:, 2] | |
| gray_image = gray_image.unsqueeze(1) | |
| return gray_image | |
| def ternary_transform(image, max_distance=1): | |
| device = image.device | |
| patch_size = 2 * max_distance + 1 | |
| intensities = rgb2gray(image) * 255 | |
| out_channels = patch_size * patch_size | |
| w = np.eye(out_channels).reshape(out_channels, 1, patch_size, patch_size) | |
| weights = torch.from_numpy(w).float().to(device) | |
| patches = F.conv2d(intensities, weights, stride=1, padding=1) | |
| transf = patches - intensities | |
| transf_norm = transf / torch.sqrt(0.81 + torch.square(transf)) | |
| return transf_norm | |
| def hamming_distance(t1, t2): | |
| dist = torch.square(t1 - t2) | |
| dist_norm = dist / (0.1 + dist) | |
| dist_sum = torch.sum(dist_norm, dim=1, keepdim=True) | |
| return dist_sum | |
| def create_mask(mask, paddings): | |
| """ | |
| padding: [[top, bottom], [left, right]] | |
| """ | |
| shape = mask.shape | |
| inner_height = shape[2] - (paddings[0][0] + paddings[0][1]) | |
| inner_width = shape[3] - (paddings[1][0] + paddings[1][1]) | |
| inner = torch.ones([inner_height, inner_width]) | |
| mask2d = F.pad(inner, pad=[paddings[1][0], paddings[1][1], paddings[0][0], paddings[0][1]]) # mask最外边一圈都pad成0了 | |
| mask3d = mask2d.unsqueeze(0) | |
| mask4d = mask3d.unsqueeze(0).repeat(shape[0], 1, 1, 1) | |
| return mask4d.detach() | |
| def ternary_loss2(frame1, warp_frame21, confMask, masks, max_distance=1): | |
| """ | |
| Args: | |
| frame1: torch tensor, with shape [b * t, c, h, w] | |
| warp_frame21: torch tensor, with shape [b * t, c, h, w] | |
| confMask: confidence mask, with shape [b * t, c, h, w] | |
| masks: torch tensor, with shape [b * t, c, h, w] | |
| max_distance: maximum distance. | |
| Returns: ternary loss | |
| """ | |
| t1 = ternary_transform(frame1) | |
| t21 = ternary_transform(warp_frame21) | |
| dist = hamming_distance(t1, t21) # 近似求解，其实利用了mask区域和外界边缘交叉的那一部分像素 | |
| loss = torch.mean(dist * confMask * masks) / torch.mean(masks) | |
| return loss | |
| def gradient_loss(frame1, frame2, confMask): | |
| device = frame1.device | |
| frame1_edge = SobelLayer(device)(frame1) | |
| frame2_edge = SobelLayer(device)(frame2) | |
| loss = torch.sum(torch.abs(frame1_edge * confMask - frame2_edge * confMask)) / (torch.sum(confMask) + 1) # escape divide 0 | |
| return loss | |
| def seperate_gradient_loss(frame1, warp_frame21, confMask): | |
| device = frame1.device | |
| mask_x = create_mask(frame1, [[0, 0], [1, 1]]).to(device) | |
| mask_y = create_mask(frame1, [[1, 1], [0, 0]]).to(device) | |
| gradient_mask = torch.cat([mask_x, mask_y], dim=1).repeat(1, 3, 1, 1) | |
| frame1_edge = SeperateSobelLayer(device)(frame1) | |
| warp_frame21_edge = SeperateSobelLayer(device)(warp_frame21) | |
| loss = nn.L1Loss()(frame1_edge * confMask * gradient_mask, warp_frame21_edge * confMask * gradient_mask) | |
| return loss | |