import torch import torch.nn as nn import torchvision.transforms as transforms import numpy as np import cv2 import os def get_module_list(module, n_modules): ml = nn.ModuleList() for _ in range(n_modules): ml.append(module()) return ml # def conv2d_padding_same(input_size, input_channels, output_channels, kernel_size, \ # stride=[1,1], bias=True, weight=None, padding_value=0): # """ # This function is following tensorflow padding style, indicating that it tries to # pad evenly left(top) and right(bottom), but if the amount of columns to be added is odd, # it will add the extra column to the right(bottom). # """ # if weight is not None: # weight = np.asarray(weight) # assert weight is None or list(weight.shape) == [output_channels, input_channels, *kernel_size] # height = float(input_size[2]) # width = float(input_size[3]) # out_size = np.ceil([height / stride[0], width / stride[1]]) # padding_vertical = (out_size[0] - 1) * stride[0] + kernel_size[0] - height # padding_horizontal = (out_size[1] - 1) * stride[1] + kernel_size[1] - width # padding_left = int(np.floor(padding_horizontal / 2)) # padding_right = int(np.ceil(padding_horizontal / 2)) # padding_top = int(np.floor(padding_vertical / 2)) # padding_bottom = int(np.ceil(padding_vertical / 2)) # assert padding_left + padding_right == padding_horizontal, "{}, {}, {}".format(padding_left, padding_right, padding_horizontal) # padding_layer = nn.ConstantPad2d((padding_left, padding_right, padding_top, padding_bottom), padding_value) # conv_layer = nn.Conv2d(input_channels, output_channels, kernel_size, stride=tuple(stride), bias=bias) # if weight is not None: # conv_layer.weight.data = torch.FloatTensor(weight) # return nn.Sequential(padding_layer, conv_layer) def im_tensor_to_numpy(x): transpose = transforms.ToPILImage() x = np.asarray(transpose(x)) return x def save_im_tensor(x, addr): x = x.cpu().float() transpose = transforms.ToPILImage() x = transpose(x[0]) x.save(addr) def save_flow_to_img(flow, h, w, c, name='_result.png'): hsv = np.zeros((h, w, c), dtype=np.uint8) hsv[:, :, 0] = 255 hsv[:, :, 2] = 255 mag, ang = cv2.cartToPolar(flow[..., 0], flow[..., 1]) hsv[..., 0] = ang * 180 / np.pi / 2 hsv[..., 1] = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX) rgb = cv2.cvtColor(hsv, cv2.COLOR_HSV2BGR) res_img_path = os.path.join('result', name) cv2.imwrite(res_img_path, rgb) def torch_where(cond, x_1, x_2): cond = cond.float() return (cond * x_1) + ((1-cond) * x_2) def meshgrid(height, width, n_repeat): # print(height, width, n_repeat) x_t = torch.matmul(torch.ones(height, 1), torch.transpose(torch.linspace(-1.0, 1.0, width)[:, None], 1, 0)) # print(x_t) y_t = torch.matmul(torch.linspace(-1.0, 1.0, height)[:, None], torch.ones(1, width)) # print(y_t) x_t_flat = x_t.view(1, -1) y_t_flat = y_t.view(1, -1) grid = torch.cat([x_t_flat, y_t_flat])[None, ...].cuda().view(-1) grid = grid.repeat(n_repeat) grid = grid.view(n_repeat, 2, -1) grid = grid.permute(0, 2, 1).contiguous().view(n_repeat, height, width, 2) # print(grid[0, :, :, 0], x_t) return grid