enable celery tasks

1 files changed, 187 insertions, 0 deletions

diff --git a/megapixels/app/server/api_task.py b/megapixels/app/server/api_task.py
new file mode 100644
index 00000000..acdb5b7d
--- /dev/null
+++ b/megapixels/app/server/api_task.py
@@ -0,0 +1,187 @@
+import os
+import re
+import uuid
+import time
+import dlib
+import simplejson as json
+import numpy as np
+from flask import Blueprint, request, jsonify
+from PIL import Image  # todo: try to remove PIL dependency
+
+# from app.models.sql_factory import load_sql_datasets, list_datasets, get_dataset, get_table
+# from app.utils.im_utils import pil2np
+
+from celery.result import AsyncResult
+from app.server.tasks import celery
+
+api_task = Blueprint('api_task', __name__)
+
+@api_task.route('/')
+def index():
+  """Dummy index"""
+  return jsonify({})
+
+# from flask import render_template, redirect, url_for, send_from_directory
+# from flask import request, make_response, jsonify
+# from . import main, utils
+from app.server.tasks import task_lookup, list_active_tasks
+# from PIL import Image, ImageOps
+# import cv2 as cv
+
+# import imutils
+
+@api_task.route('/<task_name>/<task_id>')
+def task_status(task_name, task_id):
+    """Return celery image processing status"""
+    if task_name in task_lookup:
+        task = task_lookup[task_name]['task'].AsyncResult(task_id)
+        # task = AsyncResult(task_id, app=celery)
+    else:
+        return jsonify({
+            'state': 'error',
+            'percent': 100,
+            'message': 'Unknown task'
+        })
+
+    # app.logger.info('task state: {}'.format(task.state))
+    if task.state == 'PENDING':
+        response = {
+            'state': task.state,
+            'percent': 0,
+            'message': 'Pending...'
+        }
+    elif task.state != 'FAILURE':
+        response = {
+            'state': task.state,
+            'percent': task.info.get('percent', 0),
+            'uuid': task.info.get('uuid', 0),
+            'message': task.info.get('message', '')
+        }
+        if 'result' in task.info:
+            response['result'] = task.info['result']
+    else:
+        # something went wrong in the background job
+        response = {
+            'state': task.state,
+            'percent': 100,
+            'message': str(task.info),  # this is the exception raised
+        }
+    return jsonify(response)
+
+@api_task.route('/upload/sleep', methods=['GET', 'POST'])
+def sleep_test():
+    async_task = task_lookup['sleep']['task'].apply_async(args=['sleep_test'])
+    task_url = '/task/{}/{}'.format('sleep', async_task.id)
+    return jsonify({
+        'result': True,
+        'task_url': task_url,
+    })
+
+# @api_task.route('/upload', methods=['POST'])
+# def upload():
+
+#     style = request.form['style']
+#     print('style',style)
+#     if style in task_lookup:
+#         task = task_lookup[style]['task']
+#         print('task',task)
+#     else:
+#         return jsonify({
+#             'result': False,
+#             'error': 'Unknown task',
+#         })
+
+#     file = request.files['user_image']
+#     agree = bool(request.form['agree'])
+#     ext = request.form['ext']
+#     if ext is None:
+#         ext = request.files['ext']
+
+#     uuid_name = str(uuid.uuid4())
+
+#     app.logger.info('[+] style: {}'.format(style))
+#     app.logger.info('[+] ext: {}'.format(ext))
+#     app.logger.info('[+] uuid_name: {}'.format(uuid_name))
+#     app.logger.info('[+] agreed: {}'.format(agree))
+
+#     # convert PNG to JPG
+#     print('[+] Resizing image')
+
+#     # LOL MaskRCNN needs to be run outside of the Celery Task    
+#     im = Image.open(file.stream).convert('RGB')
+#     im = ImageOps.fit(im,(512,512))
+#     if agree:
+#         upload_folder = app.config['UPLOADS']
+#     else:
+#         upload_folder = app.config['UPLOADS_PRIVATE']
+
+#     fpath = os.path.join(upload_folder, uuid_name + '.jpg')
+
+#     # Save image to disk
+#     print('[+] Save image to {}'.format(fpath))
+#     im.save(fpath, 'JPEG', quality=100)
+#     im_pil_256 = im.resize((256,256))
+
+#     print('[+] ensure_np...')
+#     im_np = imx.ensure_np(im_pil_256)
+#     #print('[+] resize np...')
+#     #im_np = imutils.resize(im_np,width=256)
+
+#     upload_dir, render_dir, json_dir, upload_uri, render_uri = get_paths(agree)
+
+#     print('[+] Run mrcnn...')
+#     try:
+#         result = mask_rcnn.create_segmentations(im_np,concat=True)
+#     except:
+#         print('[-] Error. Could not run mask_rcnn')
+#         result = []
+
+#     if len(result) > 0:
+#         result = result[0]
+
+#         # save data, then pass to celery task
+#         print('[+] Save masks')
+#         seg_mask = result['seg_mask']
+#         fpath_seg_mask = os.path.join(render_dir, uuid_name + '_seg_mask.jpg')
+#         #cv.imwrite(fpath_seg_mask,cv.cvtColor(seg_mask,cv.COLOR_BGR2RGB))
+#         #seg_mask = seg_mask[:,:,::-1]
+#         seg_mask_pil = imx.ensure_pil(seg_mask)
+#         seg_mask_pil.save(fpath_seg_mask, 'JPEG', quality=100)
+
+#         im_mask = result['im_mask']
+#         fpath_im_mask = os.path.join(render_dir, uuid_name + '_im_mask.jpg')
+#         #im_mask = im_mask[:,:,::-1]
+#         im_mask_pil = imx.ensure_pil(im_mask)
+#         im_mask_pil.save(fpath_im_mask, 'JPEG',quality=100)
+#         #cv.imwrite(fpath_im_mask,cv.cvtColor(im_mask,cv.COLOR_BGR2RGB))
+
+#         celery_result = {
+#             'score':str(result['score']),
+#             'name':str(result['name']),
+#             'class_index':str(result['class_index']),
+#             'color':str(result['color']),
+#             'fp_im_mask':fpath_im_mask,
+#             'fp_seg_mask':fpath_seg_mask,
+#             'valid':True
+#         }
+#     else:
+#         print('[-] no reults. process background only')
+#         celery_result = {
+#             'score':None,
+#             'name':None,
+#             'class_index':None,
+#             'color':None,
+#             'fp_im_mask':None,
+#             'fp_seg_mask':None,
+#             'valid':False
+#         }
+
+#     print('[+] Start celery')
+#     async_task = task.apply_async(args=[uuid_name, agree, celery_result])
+#     task_url = url_for('main.task_status', task_name=style, task_id=async_task.id)
+
+#     return jsonify({
+#         'result': True,
+#         'task_url': task_url,
+#         'uuid': uuid_name
+#     })