def doCrop(opt, model, x, padding=1, sc=1):
pad = padImageReflect(padding)
unpad = unpadImage(sc * padding)
squeeze = 1 if (not hasattr(opt, 'C2B')) or opt.C2B else 0
c = x.shape[0]
baseFlag = hasattr(opt, 'mode') and opt.mode == 'lite.old'
if baseFlag:
c = c >> 1
base = padImageReflect(sc * padding)(x[c:].unsqueeze(squeeze))
x = x[:c, :x.shape[1] >> 1, :x.shape[2] >> 1]
hOut = x.shape[1] * sc
wOut = x.shape[2] * sc
x = pad(x.unsqueeze(squeeze))
_, _, h, w = x.shape
tmp_image = torch.zeros([c, hOut, wOut]).to(x)
cropsize = opt.cropsize
if not cropsize:
try:
freeRam = config.calcFreeMem()
cropsize = int(np.sqrt(freeRam * opt.ramCoef / c))
except:
raise MemoryError()
if cropsize > 2048:
cropsize = 2048
if not cropsize > 32:
raise MemoryError()
size = cropsize - 2 * padding
for topS, leftS in itertools.product(cropIter(h, padding, size),
cropIter(w, padding, size)):
leftT = leftS * sc
topT = topS * sc
bottomS = topS + cropsize
rightS = leftS + cropsize
s = x[:, :, topS:bottomS, leftS:rightS]
r = model(s, base[:, :, topT:bottomS * sc,
leftT:rightS * sc]) if baseFlag else model(s)[-1]
tmp = unpad(r.squeeze(squeeze))
tmp_image[:, topT:topT + tmp.shape[1],
leftT:leftT + tmp.shape[2]] = tmp
return tmp_image
def prepareOpt(opt, shape):
sc, pad = opt.scale, opt.padding
padSc = pad * sc
if opt.iterClip is None:
try:
freeRam = config.calcFreeMem()
except Exception:
raise MemoryError('Can not calculate free memory.')
if opt.ensemble > 0:
opt2 = copy(opt)
opt2.iterClip, opt2.padImage, opt2.unpad, *_ = prepare(transposeShape(shape), freeRam, opt, pad, sc, opt.align, opt.cropsize)
opt.iterClip, opt.padImage, opt.unpad, outShape, opt.blend = prepare(shape, freeRam, opt, pad, sc, opt.align, opt.cropsize)
if (not hasattr(opt, 'outShape')) or opt.outShape is None:
opt.outShape = outShape
opt.outShape = list(opt.outShape)
if opt.ensemble > 0:
opt2.blend = opt.blend
opt2.outShape = transposeShape(opt.outShape)
opt.transposedOpt = opt2
return sc, padSc
def doCrop(opt, model, x, padding=1, sc=1):
pad = padImageReflect(padding)
c = x.shape[0]
hOut = x.shape[1] * sc
wOut = x.shape[2] * sc
squeeze = 1 if (not hasattr(opt, 'C2B')) or opt.C2B else 0
x = pad(x.unsqueeze(squeeze))
_, _, h, w = x.shape
tmp_image = torch.zeros([c, hOut, wOut]).to(x)
cropsize = opt.cropsize
if not cropsize:
try:
freeRam = config.calcFreeMem()
cropsize = int(np.sqrt(freeRam * opt.ramCoef / c))
except:
raise MemoryError()
if cropsize > 2048:
cropsize = 2048
if not cropsize > 32:
raise MemoryError()
size = cropsize - 2 * padding
print(cropsize, cropsize * cropsize * c, freeRam)
# sys.stdin.readline()
for topS, leftS in itertools.product(cropIter(h, padding, size),
cropIter(w, padding, size)):
leftT = leftS * sc
topT = topS * sc
s = x[:, :, topS:topS + cropsize, leftS:leftS + cropsize]
print('processing')
r = model(s)[-1]
tmp = r.squeeze(squeeze)[:, sc * padding:-sc * padding,
sc * padding:-sc * padding]
tmp_image[:, topT:topT + tmp.shape[1],
leftT:leftT + tmp.shape[2]] = tmp
return tmp_image
def getBatchSize(option): return max(1, int((config.calcFreeMem() / option['load']) * ramCoef[config.getRunType()]))
# pylint: disable=E1101
import logging
import torch
from slomo import UNet, backWarp
from imageProcess import initModel, getStateDict, getPadBy32, doCrop, identity, Option, extend
from config import config
log = logging.getLogger('Moe')
modelPath = './model/slomo/SuperSloMo.ckpt'
RefTime = 2
WindowSize = 2
ramCoef = [.95 / x for x in (8100., 2484., 8100., 2466., 4014., 1080.)]
getFlowComp = lambda *_: UNet(6, 4)
getFlowIntrp = lambda *_: UNet(20, 5)
getFlowBack = lambda opt: backWarp(opt.width, opt.height, config.device(), config.dtype())
getBatchSize = lambda load, ramCoef: max(1, int((config.calcFreeMem() / load) * ramCoef))
modules = dict(
flowComp={'weight': 'state_dictFC', 'f': getFlowComp, 'outShape': (1, 4, 1, 1)},
ArbTimeFlowIntrp={'weight': 'state_dictAT', 'f': getFlowIntrp, 'outShape': (1, 5, 1, 1)})
def newOpt(func, ramCoef, align=32, padding=45, scale=1, **_):
opt = Option()
opt.modelCached = func
opt.ramCoef = ramCoef
opt.align = align
opt.padding = padding
opt.scale = scale
opt.squeeze = identity
opt.unsqueeze = identity
return opt
