diff options
Diffstat (limited to 'old/server/app/main/tasks.py')
| -rw-r--r-- | old/server/app/main/tasks.py | 374 |
1 files changed, 0 insertions, 374 deletions
diff --git a/old/server/app/main/tasks.py b/old/server/app/main/tasks.py deleted file mode 100644 index 970e6988..00000000 --- a/old/server/app/main/tasks.py +++ /dev/null @@ -1,374 +0,0 @@ -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 - } -} - |