pull in dulldream celery server

1 files changed, 374 insertions, 0 deletions

diff --git a/server/app/main/tasks.py b/server/app/main/tasks.py
new file mode 100644
index 00000000..970e6988
--- /dev/null
+++ b/server/app/main/tasks.py
@@ -0,0 +1,374 @@
+import os
+import sys
+import time
+import datetime
+import json
+from PIL import Image, ImageFilter
+import cv2 as cv
+import numpy as np
+from . import main, utils
+from .. import basemodels
+from flask import current_app as app
+from .paths import get_paths
+celery = basemodels.celery
+from celery.utils.log import get_task_logger
+celery_logger = get_task_logger(__name__)
+import imutils
+
+
+# init image processors
+sys.path.append('/dulldream/src/')
+from .img_proc_config import ImgProcConfig
+from image_processors.mask_rcnn import MaskRCNN
+from image_processors.pix2pix import Pix2Pix
+from utils import imx
+from utils import fiox
+
+
+# initialize image processor
+img_proc_config = ImgProcConfig()
+p2p = Pix2Pix(img_proc_config.p2p_ckpts_dir,epoch=img_proc_config.p2p_epoch)
+p2p_objects = Pix2Pix(img_proc_config.p2p_bg_ckpts_dir,epoch=img_proc_config.p2p_epoch)
+
+mask_rcnn = MaskRCNN(img_proc_config.mask_rcnn_class_config,
+    model_path=img_proc_config.mask_rcnn_model)
+
+
+@celery.task(bind=True)
+def task_dull(self, uuid_name, agree, mask_rcnn_result):
+  """Process image and update during"""
+  celery_logger.debug('process_image_task, uuid: {}'.format(uuid_name))
+  
+  upload_dir, render_dir, json_dir, upload_uri, render_uri = get_paths(agree)
+
+  files = []
+  im = Image.open(os.path.join(upload_dir, uuid_name + '.jpg')).convert('RGB')
+  #im_np = cv.cvtColor(imx.ensure_np(im),cv.COLOR_RGB2BGR)
+  im_np = imx.ensure_np(im)
+  im_np = im_np[:,:,::-1]
+  im = im.resize((256,256))
+  im_np_256 = imutils.resize(im_np,width=256)
+
+  # Add original
+  fpath = os.path.join(render_dir, uuid_name + '_orig.jpg')
+  im.save(fpath, 'JPEG', quality=95)
+  files.append({
+      'title': 'Original',
+      'fn': render_uri + uuid_name + '_orig.jpg'
+  })
+  
+  if mask_rcnn_result['valid']:
+    # -----------------------------------------------
+    # Segment image (processed in views)
+    # seems to be an error with async celery processor?
+    # -----------------------------------------------
+
+    # parse mrcnn data
+    im_mask = cv.imread(mask_rcnn_result['fp_im_mask'])
+    seg_mask = cv.imread(mask_rcnn_result['fp_seg_mask'])
+    #score = mask_rcnn_result['score']
+    #name = mask_rcnn_result['name']
+    #color = mask_rcnn_result['color']
+    files.append({
+        'title': 'Semantic Segmentation',
+        'fn': render_uri + uuid_name + '_seg_mask.jpg'
+    })
+    files.append({
+        'title': 'Semantic Segmentation Isolate',
+        'fn': render_uri + uuid_name + '_im_mask.jpg'
+    })
+
+    
+    # -----------------------------------------------
+    # run rag generator
+    # -----------------------------------------------
+
+    self.update_state(
+    state = 'PROCESSING',
+      meta = {
+        'percent': 0.50,
+        'message': 'Applying Region Adjacency Graph',
+        'uuid': uuid_name
+    })
+
+    # save the regions adjancency graph
+    im_rag = imx.create_rag_mean(im_mask,compactness=30,n_segments=128)
+    fpath = os.path.join(render_dir, uuid_name + '_rgraph.jpg')
+    imx.save_np_as_pil(fpath,im_rag,quality=95)
+    files.append({
+        'title': 'Region Adjacency Graph',
+        'fn': render_uri + uuid_name + '_rgraph.jpg'
+    })
+
+
+    # -----------------------------------------------
+    # generate p2p fake
+    # -----------------------------------------------
+
+    self.update_state(
+      state = 'PROCESSING',
+      meta = {
+        'percent': 0.75,
+        'message': 'Running generative adversarial network...',
+        'uuid': uuid_name
+    })
+
+    
+    # convert segmentation to mask
+    seg_mask_gray = cv.cvtColor(seg_mask,cv.COLOR_BGR2GRAY)
+    seg_mask_gray[seg_mask_gray > 1] = 255
+
+    # find best P2P fit
+    ims_p2p = []
+    match_amts = []
+    iters = 15
+    for i in range(0,iters):
+      im_p2p = p2p.create_p2p(im_rag)
+      ims_p2p.append(im_p2p)
+      im_p2p_mask = cv.cvtColor(im_p2p,cv.COLOR_RGB2GRAY)
+      im_p2p_mask[im_p2p_mask > 1] = 255
+      # find where masks intersect
+      matches = np.bitwise_and(im_p2p_mask,seg_mask_gray)
+      amt = len(np.where(matches == 255)[0])
+      match_amts.append(amt)
+      self.update_state(
+      state = 'PROCESSING',
+      meta = {
+        'percent': 0.75,
+        'message': 'Generating ({}/{})'.format(i,iters),
+        'uuid': uuid_name
+      })
+      
+    best_idx = np.argmax(match_amts)
+    im_p2p = ims_p2p[best_idx]
+
+    fpath = os.path.join(render_dir, uuid_name + '_gan.jpg')
+    imx.save_np_as_pil(fpath,im_p2p,quality=95)
+    files.append({
+        'title': 'Generative Adversarial Network',
+        'fn': render_uri + uuid_name + '_gan.jpg'
+    })
+
+
+    # -----------------------------------------------
+    # generate p2p fake
+    # -----------------------------------------------
+
+    # announce to user
+    self.update_state(
+      state = 'PROCESSING',
+      meta = {
+        'percent': 0.90,
+        'message': 'Compositing images...',
+        'uuid': uuid_name
+    })
+
+
+    # apply masked cloning
+    im_p2p_gray = cv.cvtColor(im_p2p,cv.COLOR_BGR2GRAY)
+    im_clone_mask = np.zeros_like(im_p2p_gray,dtype=np.uint8)
+    im_clone_mask[im_p2p_gray > 1] = 255
+
+
+    # apply smoothed copy+paste clone
+    im_blur_mask = np.zeros(im_np_256.shape[:2],dtype=np.float64)
+    im_blur_mask[im_p2p_gray > 1] = 1.0
+    im_blur_mask = np.array([im_blur_mask,im_blur_mask,im_blur_mask]).transpose((1,2,0))
+    
+    # erode mask to remove black border
+    kernel = np.ones((3,3),np.uint8)
+    im_blur_mask = cv.erode(im_blur_mask,kernel,iterations = 3)
+
+    # feather mask
+    feather_amt = (3,3)
+    im_blur_mask = (cv.GaussianBlur(im_blur_mask,feather_amt, 0) > 0) * 1.0 #?
+    im_blur_mask = cv.GaussianBlur(im_blur_mask,feather_amt, 0)
+    im_blur_mask = np.clip(im_blur_mask,0.0,1.0)
+
+    # mask p2p fg --> photo bg
+    im_dull = im_np_256.astype(np.float64) * (1.0 - im_blur_mask) + im_p2p.astype(np.float64) * im_blur_mask
+    im_dull = im_dull.astype(np.uint8)
+    
+
+  else:
+    print('No person. Apply background P2P')
+    celery_logger.debug('No person. Apply background P2P, uuid: {}'.format(uuid_name))
+    im_bg_blur = cv.GaussianBlur(im_np_256,(31,31),0)
+    im_bg_rag = imx.create_rag_mean(im_bg_blur,compactness=30,n_segments=64)
+
+    # apply gan
+    im_dull = p2p_objects.create_p2p(im_bg_rag)
+
+  # resize back to full 512px
+  im_dull_512 = imutils.resize(im_dull,width=512)
+
+  # save dulldream image
+  fpath = os.path.join(render_dir, uuid_name + '_dull.jpg')
+  imx.save_np_as_pil(fpath,im_dull_512,quality=95)
+  files.append({
+      'title': 'Your DullDream',
+      'fn': render_uri + uuid_name + '_dull.jpg'
+  })
+  
+
+  # -----------------------------------------------
+  # Write data to disk
+  # -----------------------------------------------
+
+  data = {
+      'uuid': uuid_name,
+      'date': str(datetime.datetime.now()),
+      'files': files
+  }
+
+  json_path = os.path.join(json_dir, uuid_name + '.json')
+  with open(json_path, 'w') as json_file:
+      json.dump(data, json_file)
+ 
+  return {
+      'percent': 100,
+      'state': 'complete',
+      'uuid': uuid_name
+  }
+
+
+
+
+@celery.task(bind=True)
+def blur_task(self, uuid_name, agree, extra):
+    """Process image and update during"""
+    celery_logger.debug('process_image_task, uuid: {}'.format(uuid_name))
+
+    upload_dir, render_dir, json_dir, upload_uri, render_uri = get_paths(agree)
+
+    files = []
+
+    im = Image.open(os.path.join(upload_dir, uuid_name + '.jpg')).convert('RGB')
+    im = im.resize((256,256))
+    files.append({
+        'title': 'Original image',
+        'fn': upload_uri + uuid_name + '.jpg'
+    })
+
+    self.update_state(
+        state = 'PROCESSING',
+        meta = {
+            'percent': 0.25,
+            'message': 'Applying blur',
+            'uuid': uuid_name
+        })
+
+    im_np = utils.ensure_np(im)
+    im_blur = cv.blur(im_np, (5,5), 1.0)
+    im_blur_pil = utils.ensure_pil(im_blur)
+
+    fn = uuid_name + '_blur.jpg'
+    fpath = os.path.join(render_dir, fn)
+    im_blur_pil.save(fpath, 'JPEG', quality=95)
+
+    files.append({
+        'title': 'Blurred image',
+        'fn': render_uri + uuid_name + '_blur.jpg'
+    })
+
+    time.sleep(3)
+
+    self.update_state(
+        state = 'PROCESSING',
+        meta = {
+            'percent': 0.50,
+            'message': 'Sleeping for some reason',
+            'uuid': uuid_name
+        })
+    time.sleep(2)
+
+    self.update_state(
+        state = 'PROCESSING',
+        meta = {
+            'percent': 0.75,
+            'message': 'Sleeping some more',
+            'uuid': uuid_name
+        })
+    time.sleep(2)
+
+    data = {
+        'uuid': uuid_name,
+        'date': str(datetime.datetime.now()),
+        'files': files
+    }
+
+    json_path = os.path.join(json_dir, uuid_name + '.json')
+    with open(json_path, 'w') as json_file:
+        json.dump(data, json_file)
+
+    celery_logger.debug('ok')
+        
+    return {
+        'percent': 100,
+        'state': 'complete',
+        'uuid': uuid_name,
+    }
+
+@celery.task(bind=True)
+def sleep_task(self, uuid_name):
+    celery_logger.debug('sleep_task'.format(uuid_name))
+    msgs = [ 
+        {'msg':'Uploaded OK','time':.1},
+        {'msg':'Segmenting Image...','time':2},
+        {'msg':'Found: Person, Horse','time':1},
+        {'msg':'Creating Pix2Pix','time':2}
+    ]
+    for i,m in enumerate(msgs):
+        percent = int(float(i)/float(len(msgs))*100.0)
+        self.update_state(
+            state = 'PROCESSING',
+            meta = {
+                'percent': percent,
+                'message': m['msg'],
+                'uuid': uuid_name
+            })
+        celery_logger.debug(m['msg'])
+        time.sleep(m['time'])
+
+    return {
+        'percent': 100,
+        'state': 'complete',
+        'uuid': uuid_name
+    }
+
+def make_task_json():
+    dropdown = {}
+    for k,v in task_lookup.items():
+        if 'active' not in v or v['active'] is not False:
+            is_default = 'default' in v and v['default'] is True
+            task = {
+                'name': k,
+                'title': v['title'],
+                'selected': is_default,
+            }
+            dropdown[k] = task
+    return json.dumps(dropdown)
+
+# Add all valid tasks to this lookup.
+# Set 'active': False to disable a task
+# Set 'default': True to define the default task
+
+task_lookup = {
+    'sleep': {
+        'title': 'Sleep Test',
+        'task': sleep_task,
+        'active': False
+    },
+    'blur': {
+        'title': 'Blur',
+        'task': blur_task,
+        'active': False
+    },
+    'task_dull': {
+        'title': 'DullDream V2',
+        'task': task_dull,
+        'active': True,
+        'default': True
+    }
+}
+