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
opt.sf = option['sf']
opt.firstTime = 1
opt.notLast = 1
if opt.sf < 2:
raise RuntimeError('Error: --sf/slomo factor has to be at least 2')
return opt
def getOpt(option):
def opt():pass
# Initialize model
opt.model = modelPath
dict1 = getStateDict(modelPath)
opt.flowComp = initModel(opt, dict1['state_dictFC'], 'flowComp', getFlowComp)
opt.ArbTimeFlowIntrp = initModel(opt, dict1['state_dictAT'], 'ArbTimeFlowIntrp', getFlowIntrp)
opt.sf = option['sf']
opt.firstTime = 1
opt.notLast = 1
opt.batchSize = 0
if opt.sf < 2:
raise RuntimeError('Error: --sf/slomo factor has to be at least 2')
return opt
def getOpt(option):
opt = Option(modelPath)
# Initialize model
dict1 = getStateDict(modelPath)
flowComp = initModel(opt, dict1['state_dictFC'], 'flowComp', getFlowComp)
ArbTimeFlowIntrp = initModel(opt, dict1['state_dictAT'], 'ArbTimeFlowIntrp', getFlowIntrp)
opt.sf = option['sf']
opt.firstTime = 1
opt.notLast = 1
opt.batchSize = 0
opt.flowBackWarp = None
opt.optFlow = newOpt(flowComp, 0)
opt.optArb = newOpt(ArbTimeFlowIntrp, 1)
if opt.sf < 2:
raise RuntimeError('Error: --sf/slomo factor has to be at least 2')
return opt
def getOpt(*_): def opt():pass opt.model = modelPath opt.modelDef = AODnet dict1 = load(modelPath, map_location='cpu') opt.modelCached = initModel(getModel(opt), dict1) return opt
def getOpt(optDe):
model = optDe.get('model', 'dehaze')
opt = Option()
modelPath, opt.modelDef, opt.ram, opt.padding, opt.align, opt.prepare = mode_switch[
model]
opt.model = modelPath
opt.modelCached = initModel(opt, modelPath, model)
return opt
def getOpt(*_):
def opt():
pass
opt.model = modelPath
opt.modelDef = AODnet
opt.modelCached = initModel(opt, modelPath, 'dehaze')
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(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 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(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 getFlowBack(opt, width, height):
if opt.flowBackWarp:
return opt.flowBackWarp
opt.flowBackWarp = initModel(backWarp(width, height, config.device(), config.dtype()))
return opt.flowBackWarp
modelName = 'moire_obj'
test = False
inputFolder = '../test-pics'
refFile = 0 #'test/1566005911.7879605_ci.png'
def context():
pass
opt = Option(
('test/{}.pth' if test else 'model/demoire/{}.pth').format(modelName))
opt.padding = 31
opt.ramCoef = 1 / 8000.
opt.align = 128
opt.modelCached = initModel(opt, weights=opt.model, f=lambda _: Net())
toTorch = lambda x: torch.from_numpy(np.array(x)).permute(2, 0, 1).to(
dtype=config.dtype(), device=config.device()) / 256
time = 0.0
for pic in os.listdir(inputFolder):
original = toTorch(readFile(context=context)(inputFolder + '/' + pic))
ref = toTorch(readFile(context=context)(refFile + '/' +
pic)) if refFile else original
start = perf_counter()
y = ensemble(opt)(original)
time += perf_counter() - start
print(pic, float(y.mean(dtype=torch.float)),
float((y - ref).abs().mean(dtype=torch.float)))
out = toOutput(8)(toFloat(y))
writeFile(
out,
def f(data, opt):
node.reset()
node.trace(0, p='slomo start')
_, oriHeight, oriWidth = data[0][0].size()
width = upTruncBy32(oriWidth)
height = upTruncBy32(oriHeight)
pad = ReflectionPad2d((0, width - oriWidth, 0, height - oriHeight))
unpad = lambda im: im[:, :oriHeight, :oriWidth]
opt.width = width
opt.height = height
flowBackWarp = initModel(opt, None, 'flowBackWarp', getFlowBack)
if not opt.batchSize:
opt.batchSize = getBatchSize({'load': width * height})
log.info('Slomo batch size={}'.format(opt.batchSize))
batchSize = len(data)
sf = opt.sf
tempOut = [0 for _ in range(batchSize * sf + 1)]
# Save reference frames
if opt.notLast or opt.firstTime:
tempOut[0] = func(data[0][0])
outStart = 0
else:
outStart = 1
for i, frames in enumerate(data):
tempOut[(i + 1) * sf] = frames[1]
# Load data
I0 = pad(torch.stack([frames[0] for frames in data]))
I1 = pad(torch.stack([frames[1] for frames in data]))
flowOut = opt.flowComp(torch.cat((I0, I1), dim=1))
F_0_1 = flowOut[:,:2,:,:]
F_1_0 = flowOut[:,2:,:,:]
node.trace()
# Generate intermediate frames
for intermediateIndex in range(1, sf):
t = intermediateIndex / sf
temp = -t * (1 - t)
fCoeff = (temp, t * t, (1 - t) * (1 - t), temp)
wCoeff = (1 - t, t)
F_t_0 = fCoeff[0] * F_0_1 + fCoeff[1] * F_1_0
F_t_1 = fCoeff[2] * F_0_1 + fCoeff[3] * F_1_0
g_I0_F_t_0 = flowBackWarp(I0, F_t_0)
g_I1_F_t_1 = flowBackWarp(I1, F_t_1)
intrpOut = opt.ArbTimeFlowIntrp(torch.cat((I0, I1, F_0_1, F_1_0, F_t_1, F_t_0, g_I1_F_t_1, g_I0_F_t_0), dim=1))
F_t_0_f = intrpOut[:, :2, :, :] + F_t_0
F_t_1_f = intrpOut[:, 2:4, :, :] + F_t_1
V_t_0 = torch.sigmoid(intrpOut[:, 4:5, :, :])
V_t_1 = 1 - V_t_0
g_I0_F_t_0_f = flowBackWarp(I0, F_t_0_f)
g_I1_F_t_1_f = flowBackWarp(I1, F_t_1_f)
Ft_p = (wCoeff[0] * V_t_0 * g_I0_F_t_0_f + wCoeff[1] * V_t_1 * g_I1_F_t_1_f) / (wCoeff[0] * V_t_0 + wCoeff[1] * V_t_1)
# Save intermediate frame
for i in range(batchSize):
tempOut[intermediateIndex + i * sf] = unpad(Ft_p[i].detach())
node.trace()
tempOut[intermediateIndex] = func(tempOut[intermediateIndex])
for i in range(sf, len(tempOut)):
tempOut[i] = func(tempOut[i])
res = []
for item in tempOut[outStart:]:
if type(item) == list:
res.extend(item)
elif type(item) != type(None):
res.append(item)
opt.firstTime = 0
return res
def f(data):
node.reset()
node.trace(0, p='slomo start')
if opt.flowBackWarp is None:
width, height, opt.pad, opt.unpad = getPadBy32(data[0][0], opt)
opt.width = width
opt.height = height
opt.flowBackWarp = initModel(opt, None, None, getFlowBack)
else:
width, height = opt.width, opt.height
flowBackWarp = opt.flowBackWarp
if not opt.batchSize:
opt.batchSize = getBatchSize({'load': width * height})
log.info('Slomo batch size={}'.format(opt.batchSize))
batchSize = len(data)
opt.optFlow.outShape = (batchSize, 4, height, width)
opt.optArb.outShape = (batchSize, 5, height, width)
sf = opt.sf
tempOut = [0 for _ in range(batchSize * sf + 1)]
# Save reference frames
if opt.notLast or opt.firstTime:
tempOut[0] = func(data[0][0])
outStart = 0
else:
outStart = 1
for i, frames in enumerate(data):
tempOut[(i + 1) * sf] = frames[1]
# Load data
I0 = opt.pad(torch.stack([frames[0] for frames in data]))
I1 = opt.pad(torch.stack([frames[1] for frames in data]))
flowOut = doCrop(opt.optFlow, torch.cat((I0, I1), dim=1))
F_0_1 = flowOut[:,:2,:,:]
F_1_0 = flowOut[:,2:,:,:]
node.trace()
# Generate intermediate frames
for intermediateIndex in range(1, sf):
t = intermediateIndex / sf
temp = -t * (1 - t)
fCoeff = (temp, t * t, (1 - t) * (1 - t), temp)
wCoeff = (1 - t, t)
F_t_0 = fCoeff[0] * F_0_1 + fCoeff[1] * F_1_0
F_t_1 = fCoeff[2] * F_0_1 + fCoeff[3] * F_1_0
g_I0_F_t_0 = flowBackWarp(I0, F_t_0)
g_I1_F_t_1 = flowBackWarp(I1, F_t_1)
intrpOut = doCrop(opt.optArb, torch.cat((I0, I1, F_0_1, F_1_0, F_t_1, F_t_0, g_I1_F_t_1, g_I0_F_t_0), dim=1))
F_t_0_f = intrpOut[:, :2, :, :] + F_t_0
F_t_1_f = intrpOut[:, 2:4, :, :] + F_t_1
V_t_0 = torch.sigmoid(intrpOut[:, 4:5, :, :])
V_t_1 = 1 - V_t_0
g_I0_F_t_0_f = flowBackWarp(I0, F_t_0_f)
g_I1_F_t_1_f = flowBackWarp(I1, F_t_1_f)
Ft_p = (wCoeff[0] * V_t_0 * g_I0_F_t_0_f + wCoeff[1] * V_t_1 * g_I1_F_t_1_f) / (wCoeff[0] * V_t_0 + wCoeff[1] * V_t_1)
# Save intermediate frame
for i in range(batchSize):
tempOut[intermediateIndex + i * sf] = opt.unpad(Ft_p[i].detach())
node.trace()
tempOut[intermediateIndex] = func(tempOut[intermediateIndex])
for i in range(sf, len(tempOut)):
tempOut[i] = func(tempOut[i])
res = []
for item in tempOut[outStart:]:
if type(item) == list:
res.extend(item)
elif not item is None:
res.append(item)
opt.firstTime = 0
return res
def f(data):
node.reset()
node.trace(0, p='slomo start')
batchSize = len(data)
if not batchSize or len(data[0]) < 2:
return
if opt.flowBackWarp is None:
width, height, opt.pad, opt.unpad = getPadBy32(data[0][0], opt)
opt.width = width
opt.height = height
opt.flowBackWarp = initModel(opt, None, None, getFlowBack)
setOutShape(opt, height, width)
opt.batchSize = opt.flowComp.outShape[0]
log.info('Slomo batch size={}'.format(opt.batchSize))
flowBackWarp = opt.flowBackWarp
opt.flowComp.outShape[0] = batchSize
opt.ArbTimeFlowIntrp.outShape[0] = batchSize
sf = opt.sf
tempOut = [0 for _ in range(batchSize * sf + 1)]
# Save reference frames
tempOut[0] = data[0][0]
for i, frames in enumerate(data):
tempOut[(i + 1) * sf] = frames[1]
# Load data
I0 = opt.pad(torch.stack([frames[0] for frames in data]))
I1 = opt.pad(torch.stack([frames[1] for frames in data]))
flowOut = doCrop(opt.flowComp, torch.cat((I0, I1), dim=1))
F_0_1 = flowOut[:,:2,:,:]
F_1_0 = flowOut[:,2:,:,:]
node.trace()
# Generate intermediate frames
for intermediateIndex in range(1, sf):
t = intermediateIndex / sf
temp = -t * (1 - t)
fCoeff = (temp, t * t, (1 - t) * (1 - t), temp)
wCoeff = (1 - t, t)
F_t_0 = fCoeff[0] * F_0_1 + fCoeff[1] * F_1_0
F_t_1 = fCoeff[2] * F_0_1 + fCoeff[3] * F_1_0
g_I0_F_t_0 = flowBackWarp(I0, F_t_0)
g_I1_F_t_1 = flowBackWarp(I1, F_t_1)
intrpOut = doCrop(opt.ArbTimeFlowIntrp, torch.cat((I0, I1, F_0_1, F_1_0, F_t_1, F_t_0, g_I1_F_t_1, g_I0_F_t_0), dim=1))
F_t_0_f = intrpOut[:, :2, :, :] + F_t_0
F_t_1_f = intrpOut[:, 2:4, :, :] + F_t_1
V_t_0 = torch.sigmoid(intrpOut[:, 4:5, :, :])
V_t_1 = 1 - V_t_0
g_I0_F_t_0_f = flowBackWarp(I0, F_t_0_f)
g_I1_F_t_1_f = flowBackWarp(I1, F_t_1_f)
Ft_p = (wCoeff[0] * V_t_0 * g_I0_F_t_0_f + wCoeff[1] * V_t_1 * g_I1_F_t_1_f) / (wCoeff[0] * V_t_0 + wCoeff[1] * V_t_1)
# Save intermediate frame
for i in range(batchSize):
tempOut[intermediateIndex + i * sf] = opt.unpad(Ft_p[i].detach())
node.trace()
if data is None and opt.outEnd:
tempOut = tempOut[:opt.outEnd]
opt.outEnd = 0
res = []
for item in tempOut[opt.outStart:]:
extendRes(res, func(item))
opt.outStart = max(0, opt.outStart - len(tempOut))
return res