From d1da6ed6b0a6911c3b24e012ea051c9253ce8479 Mon Sep 17 00:00:00 2001 From: Jules Laplace Date: Mon, 10 Dec 2018 22:43:10 +0100 Subject: pull in dulldream celery server --- server/app/main/tasks.py | 374 +++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 374 insertions(+) create mode 100644 server/app/main/tasks.py (limited to 'server/app/main/tasks.py') 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 + } +} + -- cgit v1.2.3-70-g09d2