def main():
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = Model()
model.load_model('../train_log/HDv2', -1)
model.eval()
model.device()
path = '../../data/UCF101/ucf101_interp_ours/'
dirs = os.listdir(path)
psnr_list = []
ssim_list = []
time_list = []
# print('=========>Start Calculate PSNR and SSIM')
for d in tqdm(dirs):
img0 = (path + d + '/frame_00.png')
img1 = (path + d + '/frame_02.png')
gt = (path + d + '/frame_01_gt.png')
img0 = (torch.tensor(cv2.imread(img0).transpose(2, 0, 1) /
255.)).to(device).float().unsqueeze(0)
img1 = (torch.tensor(cv2.imread(img1).transpose(2, 0, 1) /
255.)).to(device).float().unsqueeze(0)
gt = (torch.tensor(cv2.imread(gt).transpose(2, 0, 1) /
255.)).to(device).float().unsqueeze(0)
# Avg PSNR: 35.243666269214145 SSIM: 0.9683315742368154 Time: 0.133457749911717
# inference
pred = model.inference(img0, img1)[0]
# Calculate indicators
out = pred.detach().cpu().numpy().transpose(1, 2, 0)
out = np.round(out * 255) / 255.
gt = gt[0].cpu().numpy().transpose(1, 2, 0)
psnr = compute_psnr(gt, out)
ssim = compute_ssim(gt, out)
psnr_list.append(psnr)
ssim_list.append(ssim)
# print("Avg PSNR: {} SSIM: {}".format(np.mean(psnr_list), np.mean(ssim_list)))
# print('=========>Start Calculate Inference Time')
# inference time
for i in range(100):
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
pred = model.inference(img0, img1)[0]
end.record()
torch.cuda.synchronize()
time_list.append(start.elapsed_time(end))
time_list.remove(min(time_list))
time_list.remove(max(time_list))
print("Avg PSNR: {} SSIM: {} Time: {}".format(np.mean(psnr_list),
np.mean(ssim_list),
np.mean(time_list) / 100))
default=8,
type=int,
help='limit max number of bisectional cycles')
parser.add_argument('--model',
dest='modelDir',
type=str,
default='train_log',
help='directory with trained model files')
args = parser.parse_args()
try:
try:
from model.RIFE_HDv2 import Model
model = Model()
model.load_model(args.modelDir, -1)
print("Loaded v2.x HD model.")
except:
from train_log.RIFE_HDv3 import Model
model = Model()
model.load_model(args.modelDir, -1)
print("Loaded v3.x HD model.")
except:
from model.RIFE_HD import Model
model = Model()
model.load_model(args.modelDir, -1)
print("Loaded v1.x HD model")
model.eval()
model.device()
if args.img[0].endswith('.exr') and args.img[1].endswith('.exr'):
parser.add_argument('--img', dest='img', type=str, default=None)
parser.add_argument('--montage', dest='montage', action='store_true', help='montage origin video')
parser.add_argument('--UHD', dest='UHD', action='store_true', help='support 4k video')
parser.add_argument('--skip', dest='skip', action='store_true', help='whether to remove static frames before processing')
parser.add_argument('--fps', dest='fps', type=int, default=None)
parser.add_argument('--png', dest='png', action='store_true', help='whether to vid_out png format vid_outs')
parser.add_argument('--ext', dest='ext', type=str, default='mp4', help='vid_out video extension')
parser.add_argument('--exp', dest='exp', type=int, default=1)
args = parser.parse_args()
assert (not args.video is None or not args.img is None)
if not args.img is None:
args.png = True
from model.RIFE_HDv2 import Model
model = Model()
model.load_model(os.path.join(os.path.dirname(os.path.realpath(__file__)), 'train_log'), -1)
model.eval()
model.device()
if not args.video is None:
videoCapture = cv2.VideoCapture(args.video)
fps = videoCapture.get(cv2.CAP_PROP_FPS)
tot_frame = videoCapture.get(cv2.CAP_PROP_FRAME_COUNT)
videoCapture.release()
if args.fps is None:
fpsNotAssigned = True
args.fps = fps * (2 ** args.exp)
else:
fpsNotAssigned = False
videogen = skvideo.io.vreader(args.video)
lastframe = next(videogen)