NKU-AMT/AMT

1'''2    This code is partially borrowed from IFRNet (https://github.com/ltkong218/IFRNet). 3'''4import re5import sys6import torch7import random8import numpy as np9from PIL import ImageFile10import torch.nn.functional as F11from imageio import imread, imwrite12ImageFile.LOAD_TRUNCATED_IMAGES = True13 14class InputPadder:15    """ Pads images such that dimensions are divisible by divisor """16    def __init__(self, dims, divisor=16):17        self.ht, self.wd = dims[-2:]18        pad_ht = (((self.ht // divisor) + 1) * divisor - self.ht) % divisor19        pad_wd = (((self.wd // divisor) + 1) * divisor - self.wd) % divisor20        self._pad = [pad_wd//2, pad_wd - pad_wd//2, pad_ht//2, pad_ht - pad_ht//2]21 22    def pad(self, *inputs):23        if len(inputs) == 1:24            return F.pad(inputs[0], self._pad, mode='replicate')25        else:26            return [F.pad(x, self._pad, mode='replicate') for x in inputs]27 28    def unpad(self, *inputs):29        if len(inputs) == 1:30            return self._unpad(inputs[0])31        else:32            return [self._unpad(x) for x in inputs]33    34    def _unpad(self, x):35        ht, wd = x.shape[-2:]36        c = [self._pad[2], ht-self._pad[3], self._pad[0], wd-self._pad[1]]37        return x[..., c[0]:c[1], c[2]:c[3]]38 39def img2tensor(img):40    return torch.tensor(img).permute(2, 0, 1).unsqueeze(0) / 255.041 42def tensor2img(img_t):43    return (img_t * 255.).detach(44                        ).squeeze(0).permute(1, 2, 0).cpu().numpy(45                        ).clip(0, 255).astype(np.uint8)46 47 48def read(file):49    if file.endswith('.float3'): return readFloat(file)50    elif file.endswith('.flo'): return readFlow(file)51    elif file.endswith('.ppm'): return readImage(file)52    elif file.endswith('.pgm'): return readImage(file)53    elif file.endswith('.png'): return readImage(file)54    elif file.endswith('.jpg'): return readImage(file)55    elif file.endswith('.pfm'): return readPFM(file)[0]56    else: raise Exception('don\'t know how to read %s' % file)57 58def write(file, data):59    if file.endswith('.float3'): return writeFloat(file, data)60    elif file.endswith('.flo'): return writeFlow(file, data)61    elif file.endswith('.ppm'): return writeImage(file, data)62    elif file.endswith('.pgm'): return writeImage(file, data)63    elif file.endswith('.png'): return writeImage(file, data)64    elif file.endswith('.jpg'): return writeImage(file, data)65    elif file.endswith('.pfm'): return writePFM(file, data)66    else: raise Exception('don\'t know how to write %s' % file)67 68def readPFM(file):69    file = open(file, 'rb')70 71    color = None72    width = None73    height = None74    scale = None75    endian = None76 77    header = file.readline().rstrip()78    if header.decode("ascii") == 'PF':79        color = True80    elif header.decode("ascii") == 'Pf':81        color = False82    else:83        raise Exception('Not a PFM file.')84 85    dim_match = re.match(r'^(\d+)\s(\d+)\s$', file.readline().decode("ascii"))86    if dim_match:87        width, height = list(map(int, dim_match.groups()))88    else:89        raise Exception('Malformed PFM header.')90 91    scale = float(file.readline().decode("ascii").rstrip())92    if scale < 0:93        endian = '<'94        scale = -scale95    else:96        endian = '>'97 98    data = np.fromfile(file, endian + 'f')99    shape = (height, width, 3) if color else (height, width)100 101    data = np.reshape(data, shape)102    data = np.flipud(data)103    return data, scale104 105def writePFM(file, image, scale=1):106    file = open(file, 'wb')107 108    color = None109 110    if image.dtype.name != 'float32':111        raise Exception('Image dtype must be float32.')112 113    image = np.flipud(image)114 115    if len(image.shape) == 3 and image.shape[2] == 3:116        color = True117    elif len(image.shape) == 2 or len(image.shape) == 3 and image.shape[2] == 1:118        color = False119    else:120        raise Exception('Image must have H x W x 3, H x W x 1 or H x W dimensions.')121 122    file.write('PF\n' if color else 'Pf\n'.encode())123    file.write('%d %d\n'.encode() % (image.shape[1], image.shape[0]))124 125    endian = image.dtype.byteorder126 127    if endian == '<' or endian == '=' and sys.byteorder == 'little':128        scale = -scale129 130    file.write('%f\n'.encode() % scale)131 132    image.tofile(file)133 134def readFlow(name):135    if name.endswith('.pfm') or name.endswith('.PFM'):136        return readPFM(name)[0][:,:,0:2]137 138    f = open(name, 'rb')139 140    header = f.read(4)141    if header.decode("utf-8") != 'PIEH':142        raise Exception('Flow file header does not contain PIEH')143 144    width = np.fromfile(f, np.int32, 1).squeeze()145    height = np.fromfile(f, np.int32, 1).squeeze()146 147    flow = np.fromfile(f, np.float32, width * height * 2).reshape((height, width, 2))148 149    return flow.astype(np.float32)150 151def readImage(name):152    if name.endswith('.pfm') or name.endswith('.PFM'):153        data = readPFM(name)[0]154        if len(data.shape)==3:155            return data[:,:,0:3]156        else:157            return data158    return imread(name)159 160def writeImage(name, data):161    if name.endswith('.pfm') or name.endswith('.PFM'):162        return writePFM(name, data, 1)163    return imwrite(name, data)164 165def writeFlow(name, flow):166    f = open(name, 'wb')167    f.write('PIEH'.encode('utf-8'))168    np.array([flow.shape[1], flow.shape[0]], dtype=np.int32).tofile(f)169    flow = flow.astype(np.float32)170    flow.tofile(f)171 172def readFloat(name):173    f = open(name, 'rb')174 175    if(f.readline().decode("utf-8"))  != 'float\n':176        raise Exception('float file %s did not contain <float> keyword' % name)177 178    dim = int(f.readline())179 180    dims = []181    count = 1182    for i in range(0, dim):183        d = int(f.readline())184        dims.append(d)185        count *= d186 187    dims = list(reversed(dims))188 189    data = np.fromfile(f, np.float32, count).reshape(dims)190    if dim > 2:191        data = np.transpose(data, (2, 1, 0))192        data = np.transpose(data, (1, 0, 2))193 194    return data195 196def writeFloat(name, data):197    f = open(name, 'wb')198 199    dim=len(data.shape)200    if dim>3:201        raise Exception('bad float file dimension: %d' % dim)202 203    f.write(('float\n').encode('ascii'))204    f.write(('%d\n' % dim).encode('ascii'))205 206    if dim == 1:207        f.write(('%d\n' % data.shape[0]).encode('ascii'))208    else:209        f.write(('%d\n' % data.shape[1]).encode('ascii'))210        f.write(('%d\n' % data.shape[0]).encode('ascii'))211        for i in range(2, dim):212            f.write(('%d\n' % data.shape[i]).encode('ascii'))213 214    data = data.astype(np.float32)215    if dim==2:216        data.tofile(f)217 218    else:219        np.transpose(data, (2, 0, 1)).tofile(f)220 221def warp(img, flow):222    B, _, H, W = flow.shape223    xx = torch.linspace(-1.0, 1.0, W).view(1, 1, 1, W).expand(B, -1, H, -1)224    yy = torch.linspace(-1.0, 1.0, H).view(1, 1, H, 1).expand(B, -1, -1, W)225    grid = torch.cat([xx, yy], 1).to(img)226    flow_ = torch.cat([flow[:, 0:1, :, :] / ((W - 1.0) / 2.0), flow[:, 1:2, :, :] / ((H - 1.0) / 2.0)], 1)227    grid_ = (grid + flow_).permute(0, 2, 3, 1)228    output = F.grid_sample(input=img, grid=grid_, mode='bilinear', padding_mode='border', align_corners=True)229    return output230 231def check_dim_and_resize(tensor_list):232    shape_list = []233    for t in tensor_list:234        shape_list.append(t.shape[2:])235 236    if len(set(shape_list)) > 1:237        desired_shape = shape_list[0]238        print(f'Inconsistent size of input video frames. All frames will be resized to {desired_shape}')239        240        resize_tensor_list = []241        for t in tensor_list:242            resize_tensor_list.append(torch.nn.functional.interpolate(t, size=tuple(desired_shape), mode='bilinear'))243 244        tensor_list = resize_tensor_list245 246    return tensor_list247 248