Merge pull request #18 from lotress/master

4.5.4

package.json

@@ -1,6 +1,6 @@
 {
   "name": "moephoto",
-  "version": "4.4.2",
+  "version": "4.5.4",
   "description": "MoePhoto Image Toolbox萌图工具箱   一个基于深度学习的AI图像修复软件（更新中...）   感兴趣的加QQ群320785467讨论   ### 基本功能 * 图像降噪 * 超分辨率 * 去雾 * 涂抹及风格化",
   "private": true,
   "directories": {

python/FIFOcache.py

@@ -0,0 +1,34 @@
+from collections import deque
+
+Null = lambda *_: None
+
+class Cache():
+  def __init__(self, size, default=None, onExtinct=Null):
+    self.cache = {}
+    self.size = size
+    self.queue = deque()
+    self.default = default
+    self.extinct = onExtinct
+
+  def put(self, key, item):
+    if len(self.queue) == self.size:
+      while len(self.queue):
+        oldKey = self.queue.popleft()
+        if oldKey in self.cache:
+          oldItem = self.cache[oldKey]
+          del self.cache[oldKey]
+          self.extinct(oldKey, oldItem)
+          break
+    self.cache[key] = item
+    self.queue.append(key)
+
+  def pop(self, key):
+    if key in self.cache:
+      res = self.cache[key]
+      del self.cache[key]
+      return res
+    else:
+      return self.default
+
+  def peek(self, key):
+    return key in self.cache

python/MoePhoto.py

@@ -42,11 +42,10 @@ def imageEnhance(size, *args):
     return begin(imNode, nodes, True if trace else -1).bindFunc(process)(size, name=name)
 
   return getMM(), {
-    'lock': lock,
+    'lockInterface': lock,
     'image_enhance': enhance(imageEnhance),
+    'batch': enhance(imageEnhance),
     'video_enhance': enhance(SR_vid),
-    'ednoise_enhance': enhance(imageEnhance),
-    'image_dehaze': enhance(imageEnhance),
     'systemInfo': enhance(config.system)
   }
 
@@ -60,12 +59,12 @@ def imageEnhance(size, *args):
   mp.Process(target=worker, args=(main, taskInReceiver, taskOutSender, notifier, stopEvent), daemon=True).start()
   from server import runserver
   from defaultConfig import defaultConfig
-  run = runserver(taskInSender, taskOutReceiver, noter, stopEvent, notifier, getMM())
+  run = runserver(taskInSender, taskOutReceiver, noter, stopEvent, getMM())
   host = '127.0.0.1'
   port = defaultConfig['port'][0]
   if len(sys.argv) > 1:
     if '-g' in sys.argv:
-      host = ''
+      host = '0.0.0.0'
     run(host, port)
   else:
     from webbrowser import open as startBrowser

python/defaultConfig.py

@@ -15,6 +15,8 @@
   'outDir': ('download', '输出目录'),
   'uploadDir': ('upload', '上传目录'),
   'logPath': ('.user/log.txt', '日志文件路径'),
-  'sharedMemSize': (100 * 2 ** 20, '前后台共享的内存文件交换区字节大小，要能装下png格式的一张输入或输出图片'),
+  'videoPreview': ('jpeg', '视频预览格式'),
+  'maxResultsKept': (1 << 10, '最多缓存几个运行结果'),
+  'sharedMemSize': (100 * 2 ** 20, '前后台共享的内存文件交换区字节大小，要能装下一张输入或输出图片'),
   'port': (2333, '监听端口')
 }

python/imageProcess.py

@@ -14,6 +14,7 @@
 
 deviceCPU = torch.device('cpu')
 outDir = defaultConfig['outDir'][0]
+previewFormat = defaultConfig['videoPreview'][0]
 log = logging.getLogger('Moe')
 genNameByTime = lambda: '{}/output_{}.png'.format(outDir, int(time.time()))
 padImageReflect = torch.nn.ReflectionPad2d
@@ -259,7 +260,7 @@ def initModel(model, weights=None):
 trans = [transpose, flip, flip2, combine(flip, transpose), combine(transpose, flip), combine(transpose, flip, transpose), combine(flip2, transpose)]
 transInv = [transpose, flip, flip2, trans[4], trans[3], trans[5], trans[6]]
 ensemble = lambda x, es, kwargs: reduce((lambda v, t: v + t[2](doCrop(x=t[1](x), **kwargs))), zip(range(es), trans, transInv), doCrop(x=x, **kwargs))
-previewPath = defaultConfig['outDir'][0] + '/.preview.png'
+previewPath = defaultConfig['outDir'][0] + '/.preview.{}'.format(previewFormat if previewFormat else '')
 
 def toInt(o, keys):
   for key in keys:
@@ -283,7 +284,7 @@ def appendFuncs(f, node, funcs, wrap=True):
   toOutput8,
   (lambda im: im.astype(np.uint8)),
   0,
-  lambda im: writeFile(im, context.shared, 'PNG'),
+  lambda im: writeFile(im, context.shared, previewFormat),
   lambda *_: context.root.trace(0, preview=previewPath, fileSize=context.shared.tell())]
 funcPreview = lambda im: reduce(applyNonNull, fPreview, im)
 
@@ -335,7 +336,7 @@ def procDehaze(opt, out, *_):
 
 def procResize(opt, out, nodes):
   node = Node(dict(op='resize', mode=opt['method']), 1, name=opt['name'] if 'name' in opt else None)
-  return [node.bindFunc(resize(opt, out, len(nodes), nodes))], [node], out
+  return [node.bindFunc(NonNullWrap(resize(opt, out, len(nodes), nodes)))], [node], out
 
 def procOutput(opt, out, *_):
   load = out['load']
@@ -348,10 +349,13 @@ def procOutput(opt, out, *_):
   if out['source']:
     fPreview[0] = restrictSize(2048)
     fs1 = [node0.bindFunc(toFloat), fOutput]
-    def o(im):
-      res = reduce(applyNonNull, fs1, im)
-      funcPreview(im)
-      return [res]
+    if previewFormat:
+      def o(im):
+        res = reduce(applyNonNull, fs1, im)
+        funcPreview(im)
+        return [res]
+    else:
+      o = lambda im: [reduce(applyNonNull, fs1, im)]
     fs = [o]
     if out['channel']:
       fPreview[4] = BGR2RGB
@@ -380,7 +384,11 @@ def genProcess(steps, root=True, outType=None):
       toInt(opt, ['scale'])
       opt['opt'] = runSR.getOpt(opt['scale'], opt['model'], config.ensembleSR)
     for opt in filter((lambda opt: opt['op'] == 'resize'), steps):
-      toInt(opt, ['scaleW', 'scaleH', 'width', 'height'])
+      toInt(opt, ['width', 'height'])
+      if 'scaleW' in opt:
+        opt['scaleW'] = float(opt['scaleW'])
+      if 'scaleH' in opt:
+        opt['scaleH'] = float(opt['scaleH'])
     for opt in filter((lambda opt: opt['op'] == 'DN'), steps):
       opt['opt'] = runDN.getOpt(opt['model'])
     for opt in filter((lambda opt: opt['op'] == 'dehaze'), steps):

python/server.py

@@ -7,6 +7,7 @@
 from flask import Flask, render_template, request, jsonify, send_from_directory, make_response, Response, send_file
 from gevent import pywsgi, idle, spawn
 from defaultConfig import defaultConfig
+from FIFOcache import Cache
 
 staticMaxAge = 86400
 app = Flask(__name__, root_path='.')
@@ -21,18 +22,22 @@ def current():pass
 E403 = ('Not authorized.', 403)
 E404 = ('Not Found', 404)
 OK = ('', 200)
+cache = Cache(defaultConfig['maxResultsKept'][0], OK, lambda *args: print('abandoned', *args))
 busy = lambda: (jsonify(result='Busy', eta=current.eta), 503)
 cwd = os.getcwd()
 outDir = defaultConfig['outDir'][0]
 uploadDir = defaultConfig['uploadDir'][0]
 logPath = os.path.abspath(defaultConfig['logPath'][0])
+previewFormat = defaultConfig['videoPreview'][0]
 downDir = os.path.join(app.root_path, outDir)
 if not os.path.exists(outDir):
   os.mkdir(outDir)
 with open('static/manifest.json') as manifest:
   assetMapping = json.load(manifest)
 vendorsJs = assetMapping['vendors.js']
 commonJs = assetMapping['common.js'] if 'common.js' in assetMapping else None
+getKey = lambda session, request: request.values['path'] + str(session) if 'path' in request.values else ''
+toResponse = lambda obj: (json.dumps(obj, ensure_ascii=False, separators=(',',':')), 200)
 
 def acquireSession(request):
   if current.session:
@@ -41,10 +46,11 @@ def acquireSession(request):
     noter.recv()
   current.session = request.values['session']
   current.path = request.path
+  current.key = current.path + str(current.session)
   current.eta = 10
   return False if current.session else E403
 
-def controlPoint(path, fMatch, fUnmatch, resNoCurrent, check=lambda *_: True):
+def controlPoint(path, fMatch, fUnmatch, fNoCurrent, check=lambda *_: True):
   def f():
     try:
       session = request.values['session']
@@ -55,7 +61,7 @@ def f():
     if current.session:
       return spawn(fMatch).get() if current.session == session and check(request) else fUnmatch()
     else:
-      return resNoCurrent
+      return fNoCurrent(session, request)
   app.route(path, methods=['GET', 'POST'], endpoint=path)(f)
 
 def stopCurrent():
@@ -70,19 +76,30 @@ def checkMsgMatch(request):
   return path == current.path
 
 def onConnect():
-  while current.session and not noter.poll():
+  while current.key and not (noter.poll() or cache.peek(current.key)):
     idle()
-  if current.session and noter.poll():
-    while noter.poll():
-      res = noter.recv()
-    current.eta = res['eta']
+  res = None
+  while current.key and noter.poll():
+    res = noter.recv()
+  if res:
+    if 'eta' in res:
+      current.eta = res['eta']
     if 'fileSize' in res:
       current.fileSize = res['fileSize']
       del res['fileSize']
-    return (json.dumps(res, ensure_ascii=False), 200)
+    return toResponse(res)
+  if cache.peek(current.key):
+    res = cache.pop(current.key)
+    return res
   else:
     return OK
 
+def endSession(result=None):
+  cache.put(current.key, result)
+  current.key = ''
+  current.session = None
+  return OK
+
 def makeHandler(name, prepare, final, methods=['POST']):
   def f():
     c = acquireSession(request)
@@ -91,9 +108,7 @@ def f():
     sender.send((name, *prepare(request)))
     while not receiver.poll():
       idle()
-    result = receiver.recv()
-    current.session = None
-    return final(result)
+    return endSession(final(receiver.recv()))
   app.route('/' + name, methods=methods, endpoint=name)(f)
 
 def renderPage(item, header=None, footer=None):
@@ -176,12 +191,11 @@ def getDynamicInfo(_):
   ('/', 'index.html', '主页', None),
   ('/video', 'video.html', 'AI视频', None),
   ('/batch', 'batch.html', '批量放大', None),
-  ('/ednoise', 'ednoise.html', '风格化', None),
-  ('/dehaze', 'dehaze.html', '去雾', None),
   ('/document', 'document.html', None, None),
   ('/about', 'about.html', None, about_updater, ['log']),
   ('/system', 'system.html', None, getDynamicInfo, ['disk_free', 'mem_free', 'session', 'path'], getSystemInfo()),
-  ('/gallery', 'gallery.html', None, gallery, ['var'])
+  ('/gallery', 'gallery.html', None, gallery, ['var']),
+  ('/lock', 'lock.html', None, None)
 ]
 
 for item in routes:
@@ -195,20 +209,21 @@ def getDynamicInfo(_):
 
 identity = lambda x: x
 readOpt = lambda req: json.loads(req.values['steps'])
-controlPoint('/stop', stopCurrent, lambda: E403, E404)
-controlPoint('/msg', onConnect, busy, OK, checkMsgMatch)
+onRequestCache = lambda session, request: cache.pop(getKey(session, request))
+controlPoint('/stop', stopCurrent, lambda: E403, lambda *_: E404)
+controlPoint('/msg', onConnect, busy, onRequestCache, checkMsgMatch)
 app.route('/log', endpoint='log')(lambda: send_file(logPath, add_etags=False))
 app.route('/favicon.ico', endpoint='favicon')(lambda: send_from_directory(app.root_path, 'logo3.ico'))
-app.route("/{}/.preview.png".format(outDir), endpoint='preview')(lambda: Response(current.getPreview(), mimetype='image/png'))
+if previewFormat:
+  app.route("/{}/.preview.{}".format(outDir, previewFormat), endpoint='preview')(
+    lambda: Response(current.getPreview(), mimetype='image/{}'.format(previewFormat)))
 sendFromDownDir = lambda filename: send_from_directory(downDir, filename, as_attachment=True)
 app.route("/{}/<path:filename>".format(outDir), endpoint='download')(sendFromDownDir)
-lockFinal = lambda result: (jsonify(result='Interrupted', remain=result), 499) if result > 0 else OK
-makeHandler('lock', (lambda req: [int(req.values['duration'])]), lockFinal, ['GET', 'POST'])
+lockFinal = lambda result: (jsonify(result='Interrupted', remain=result), 200) if result > 0 else (jsonify(result='Idle'), 200)
+makeHandler('lockInterface', (lambda req: [int(float(readOpt(req)[0]['duration']))]), lockFinal, ['GET', 'POST'])
 makeHandler('systemInfo', (lambda _: []), identity, ['GET', 'POST'])
 imageEnhancePrep = lambda req: (current.writeFile(req.files['file']), *readOpt(req))
 makeHandler('image_enhance', imageEnhancePrep, identity)
-makeHandler('ednoise_enhance', imageEnhancePrep, identity)
-makeHandler('image_dehaze', lambda req: (current.writeFile(req.files['file']), {'op': 'dehaze'}), identity)
 
 def videoEnhancePrep(req):
   vidfile = req.files['file']
@@ -233,43 +248,42 @@ def batchEnhance():
     os.makedirs(output_path)
   opt = readOpt(request)
   total = len(fileList)
-  print(total)
-  current.notifier.send({
+  print('batch total: {}'.format(total))
+  current.result = {
     'eta': total,
     'gone': 0,
     'total': total
-  })
+  }
+  opt[-1]['trace'] = False
   for image in fileList:
     if current.stopFlag.is_set():
       result = 'Interrupted'
       break
     name = output_path + os.path.basename(image.filename)
     start = time.time()
-    sender.send(('image_enhance', current.writeFile(image), *opt, { 'op': 'output', 'file': name, 'trace': False }))
+    opt[-1]['file'] = name
+    sender.send(('batch', current.writeFile(image), *opt))
     while not receiver.poll():
       idle()
-    if receiver.poll():
-      output = receiver.recv()
-      count += 1
-      note = {
-        'eta': (total - count) * (time.time() - start),
-        'gone': count,
-        'total': total
-      }
-      if output[1] == 200:
-        note['preview'] = name
-      else:
-        fail += 1
-      current.notifier.send(note)
-  current.session = None
-  return json.dumps({'result': (result, count, fail, output_path)}, ensure_ascii=False)
+    output = receiver.recv()
+    count += 1
+    note = {
+      'eta': (total - count) * (time.time() - start),
+      'gone': count,
+      'total': total
+    }
+    if output[1] == 200:
+      note['preview'] = name
+    else:
+      fail += 1
+    cache.put(current.key, toResponse(note))
+  return endSession(toResponse({'result': (result, count, fail, output_path)}))
 
-def runserver(taskInSender, taskOutReceiver, noteReceiver, stopEvent, notifier, mm):
+def runserver(taskInSender, taskOutReceiver, noteReceiver, stopEvent, mm):
   global sender, receiver, noter
   sender = taskInSender
   receiver = taskOutReceiver
   noter = noteReceiver
-  current.notifier = notifier
   current.stopFlag = stopEvent
   def preview():
     mm.seek(0)
@@ -282,7 +296,7 @@ def writeFile(file):
   current.writeFile = writeFile
   def f(host, port):
     app.debug = False
-    app.config['SERVER_NAME'] = '{}:{}'.format(host, port)
+    app.config['SERVER_NAME'] = None
     server = pywsgi.WSGIServer((host, port), app)
     print('Current working directory: {}'.format(cwd))
     print('Server starts to listen on http://{}:{}/, press Ctrl+C to exit.'.format(host, port))