mirror of
https://github.com/xinntao/ESRGAN.git
synced 2024-08-17 08:30:32 +03:00
add command argv
This commit is contained in:
parent
e3caa577be
commit
164beed5c7
6
test.py
6
test.py
@ -1,3 +1,4 @@
|
||||
import sys
|
||||
import os.path
|
||||
import glob
|
||||
import cv2
|
||||
@ -5,13 +6,12 @@ import numpy as np
|
||||
import torch
|
||||
import architecture as arch
|
||||
|
||||
mode = 'ESRGAN' # ESRGAN or RRDB_PSNR
|
||||
|
||||
mode = sys.argv[1] # ESRGAN or RRDB_PSNR
|
||||
|
||||
if mode == 'ESRGAN':
|
||||
model_path = './models/RRDB_ESRGAN_x4.pth'
|
||||
elif mode == 'RRDB_PSNR':
|
||||
model_path = './models/RRDB_PSNR_DF2K_x4.pth'
|
||||
model_path = './models/RRDB_PSNR_x4.pth'
|
||||
|
||||
test_img_folder = 'LR/*'
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user