328afde9c0
SPSR_model really isn't that different from SRGAN_model. Rather than continuing to re-implement everything I've done in SRGAN_model, port the new stuff from SPSR over. This really demonstrates the need to refactor SRGAN_model a bit to make it cleaner. It is quite the beast these days..
21 lines
689 B
Python
21 lines
689 B
Python
import logging
|
|
logger = logging.getLogger('base')
|
|
|
|
|
|
def create_model(opt):
|
|
model = opt['model']
|
|
# image restoration
|
|
if model == 'sr': # PSNR-oriented super resolution
|
|
from .SR_model import SRModel as M
|
|
elif model == 'srgan' or model == 'corruptgan' or model == 'spsrgan':
|
|
from .SRGAN_model import SRGANModel as M
|
|
elif model == 'feat':
|
|
from .feature_model import FeatureModel as M
|
|
elif model == 'spsr':
|
|
from .SPSR_model import SPSRModel as M
|
|
else:
|
|
raise NotImplementedError('Model [{:s}] not recognized.'.format(model))
|
|
m = M(opt)
|
|
logger.info('Model [{:s}] is created.'.format(m.__class__.__name__))
|
|
return m
|