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