def getOptS(modelPath, modules, ramCoef):
opt = Option(modelPath)
weights = getStateDict(modelPath)
opt.modules = modules
opt.ramOffset = config.getRunType() * len(modules)
for i, key in enumerate(modules):
m = modules[key]
wKey = m['weight']
constructor = m.get('f', 0)
rc = m['ramCoef'][config.getRunType()] if 'ramCoef' in m else ramCoef[opt.ramOffset + i]
o = dict((k, m[k]) for k in ('align', 'padding', 'scale') if k in m)
model = initModel(opt, weights[wKey], key, constructor)
if 'outShape' in m:
opt.__dict__[key] = newOpt(model, rc, **o)
else:
model.ramCoef = rc
opt.__dict__[key] = model
return opt
def getOpt(optDe):
model = optDe.get('model', 'dehaze')
opt = Option()
modelPath, opt.modelDef, ramCoef, opt.padding, opt.align, opt.prepare = mode_switch[
model]
opt.ramCoef = ramCoef[config.getRunType()]
opt.model = modelPath
opt.modelCached = initModel(opt, modelPath, model)
return opt
def newOpt(f, ramType): opt = Option() opt.modelCached = lambda x: (f(x),) opt.ramCoef = ramCoef[config.getRunType() * 2 + ramType] opt.align = 32 opt.padding = 45 opt.squeeze = identity opt.unsqueeze = identity return opt
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]
opt.modelCached = initModel(opt, opt.model, 'DN' + model)
return opt
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]
opt.modelCached = initModel(getModel(opt),
load(opt.model, map_location='cpu'))
return opt
def getOpt(optDN):
model = optDN['model']
if not model in mode_switch:
return
opt = Option(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]
opt.modelCached = initModel(opt, opt.model, 'DN' + model)
opt.squeeze = lambda x: x.squeeze(1)
opt.unsqueeze = lambda x: x.unsqueeze(1)
opt.padding = 5
return opt
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]
if opt.cropsize:
print('当前denoise切块大小：', opt.cropsize)
opt.modelCached = getModel(opt, False)
return opt
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][0]
opt.modelDef = mode_switch[nmode][1]
opt.scale = scale
opt.ensemble = ensemble
opt.ramCoef = mode_switch[nmode][2][config.getRunType()]
opt.cropsize = config.getConfig()[0]
opt.modelCached = initModel(opt, opt.model, 'SR' + nmode)
return opt
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][0]
opt.modelDef = mode_switch[nmode][1]
opt.scale = scale
opt.ensemble = ensemble
opt.ramCoef = mode_switch[nmode][2][config.getRunType()]
opt.cropsize = config.getConfig()[0]
if opt.cropsize:
print('当前SR切块大小：', opt.cropsize)
opt.modelCached = initModel(getModel(opt), load(opt.model, map_location='cpu'))
return opt
def getOpt(optSR):
opt = Option()
opt.mode = optSR['model']
opt.scale = optSR['scale']
nmode = opt.mode + str(opt.scale)
if not nmode in mode_switch:
return
if opt.mode[:3] != 'gan':
opt.squeeze = lambda x: x.squeeze(1)
opt.unsqueeze = lambda x: x.unsqueeze(1)
opt.padding = 9 if opt.scale == 3 else 5
opt.model = mode_switch[nmode][0]
opt.modelDef = mode_switch[nmode][1]
opt.ensemble = optSR['ensemble'] if 'ensemble' in optSR and (
0 <= optSR['ensemble'] <= 7) else config.ensembleSR
opt.ramCoef = mode_switch[nmode][2][config.getRunType()]
opt.cropsize = config.getConfig()[0]
opt.modelCached = initModel(opt, opt.model, 'SR' + nmode)
return opt
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
opt.ensemble = ensemble
modelType = (scale - 2) * 3
if mode[:3] == 'gan':
modelType += 9
opt.ramCoef = ramCoef[config.getRunType() + modelType]
opt.cropsize = config.getConfig()[0]
if opt.cropsize:
print('当前SR切块大小：', opt.cropsize)
opt.modelCached = getModel(opt, False)
return opt
def getBatchSize(option): return max(1, int((config.calcFreeMem() / option['load']) * ramCoef[config.getRunType()]))
