File size: 1,388 Bytes
e34aada |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import cv2
import os
import imageio
def dilate(bin_img, ksize=5):
# bin_img
pad = (ksize-1)//2
bin_img = F.pad(bin_img, pad=[pad,pad,pad,pad], mode='reflect')
out = F.max_pool2d(bin_img, kernel_size=ksize, stride=1, padding=0)
return out
def erode(bin_img, ksize=5):
out = 1 - dilate(1-bin_img, ksize)
return out
def to8b(x):
return (255*np.clip(x, 0, 1)).astype(np.uint8)
def mse2psnr(x):
return -10. * torch.log(x) / torch.log(torch.Tensor([10.]))
def img2mse(x, y):
return torch.mean((x - y) ** 2)
def video2images(video_name, out_dir):
cap = cv2.VideoCapture(video_name)
frame_num = 0
while(True):
_, frame = cap.read()
if frame is None:
break
out_frame_name = os.path.join(out_dir, str(frame_num) + '.jpg')
cv2.imwrite(out_frame_name, frame)
frame_num += + 1
cap.release()
def load_image_as_uint8_tensor(fname):
"""
img: (H, W, 3) floatTensor
"""
img = torch.as_tensor(imageio.imread(fname))
return img
if __name__ =='__main__':
video2images("test_data/May_val/AD-NeRF.mp4", "test_data/May_val/AD-NeRF")
video2images("test_data/May_val/GeneFace.mp4", "test_data/May_val/GeneFace")
video2images("test_data/May_val/GT.mp4", "test_data/May_val/GT") |