import torch
from torch.nn import ReflectionPad2d
#from torchvision.transforms import Normalize
from slomo import UNet, backWarp
from imageProcess import genGetModel, initModel
from config import config
log = logging.getLogger('Moe')
modelPath = './model/slomo/SuperSloMo.ckpt'
ramCoef = [.9 / x for x in (6418.7, 1393., 1156.3)]
#mean = [0.429, 0.431, 0.397]
#std = [1, 1, 1]
#negMean = [-x for x in mean]
#identity = lambda x, *_: x
upTruncBy32 = lambda x: (-x & 0xffffffe0 ^ 0xffffffff) + 1
getFlowComp = genGetModel(lambda *_: UNet(6, 4))
getFlowIntrp = genGetModel(lambda *_: UNet(20, 5))
def getFlowBack(opt, width, height):
if opt.flowBackWarp:
return opt.flowBackWarp
opt.flowBackWarp = initModel(backWarp(width, height, config.device(), config.dtype()))
return opt.flowBackWarp
def getOpt(option):
def opt():pass
# Initialize model
opt.model = modelPath
dict1 = torch.load(modelPath, map_location='cpu')
opt.flowComp = initModel(getFlowComp(opt), dict1['state_dictFC'])
opt.ArbTimeFlowIntrp = initModel(getFlowIntrp(opt), dict1['state_dictAT'])
opt.flowBackWarp = 0
def sr(x, opt):
sc = opt.scale
sum = ensemble(
x, opt.ensemble, {
'opt': opt,
'model': getModel(opt),
'padding': 2 if sc == 3 else 1,
'sc': sc
})
if opt.ensemble:
return sum / (opt.ensemble + 1)
else:
return sum
getModel = genGetModel(
lambda opt, *args: models[opt.mode[:3] == 'gan'][opt.scale - 2]())
##################################
def getOpt(scale, mode, ensemble):
def opt():
pass
nmode = mode + str(scale)
if not nmode in mode_switch:
return
opt.C2B = mode[:3] != 'gan'
opt.mode = mode
opt.model = mode_switch[nmode]
opt.scale = scale
from models import NetDN, SEDN
from config import config
ramCoef = .9 / np.array([[2700., 2400., 1253.4], [4106.9, 7405., 4304.2]])
mode_switch = {
'15': ('./model/l15/model_new.pth', SEDN, ramCoef[1]),
'25': ('./model/l25/model_new.pth', SEDN, ramCoef[1]),
'50': ('./model/l50/model_new.pth', SEDN, ramCoef[1]),
'lite5': ('./model/dn_lite5/model_new.pth', NetDN, ramCoef[0]),
'lite10': ('./model/dn_lite10/model_new.pth', NetDN, ramCoef[0]),
'lite15': ('./model/dn_lite15/model_new.pth', NetDN, ramCoef[0]),
}
dn = lambda x, opt: ensemble(x, 0, {'opt': opt, 'model': opt.modelCached})
getModel = genGetModel()
##################################
def getOpt(model):
def opt():
pass
if not model in mode_switch:
return
opt.model = mode_switch[model][0]
opt.modelDef = mode_switch[model][1]
opt.ramCoef = mode_switch[model][2][config.getRunType()]
opt.cropsize = config.getConfig()[1 if model[:4] == 'lite' else 2]
from models import NetDN, SEDN
from config import config
model_dict = {
'15': './model/l15/model_new.pth',
'25': './model/l25/model_new.pth',
'50': './model/l50/model_new.pth',
'lite5': './model/dn_lite5/model_new.pth',
'lite10': './model/dn_lite10/model_new.pth',
'lite15': './model/dn_lite15/model_new.pth'
}
ramCoef = .9 / np.array([2700., 4106.9, 2400., 7405., 1253.4, 4304.2])
dn = lambda x, opt: ensemble(x, 0, {'opt': opt, 'model': getModel(opt)})
getModel = genGetModel(lambda opt, *args: NetDN()
if 'dn_lite' in opt.model else SEDN())
##################################
def getOpt(model):
def opt():
pass
if not model in model_dict:
return
opt.model = model_dict[model]
modelType = 0 if model[:4] == 'lite' else 1
opt.ramCoef = ramCoef[config.getRunType() * 2 + modelType]
opt.cropsize = config.getConfig()[modelType + 1]