Refactor: storage and db

CHANGELOG.md

@@ -11,6 +11,8 @@ The format is based on [Keep a Changelog](https://github.com/olivierlacan/keep-a
 
 ### Changed
 
+- DB schema and storage helper. (https://github.com/yushiang-demo/dula-net-webapp/pull/4)
+
 ### Fixed
 
 - Task response missing 2d coords. (https://github.com/yushiang-demo/dula-net-webapp/pull/3)

api/models/Task.py

@@ -19,14 +19,18 @@
 })
 
 Images = api.model('Images',{
-    'input': fields.String,
     'aligned': fields.String,
     'layout': fields.String,
 })
 
+Output = api.model('Output',{   
+    'images': fields.Nested(Images),
+    'layout': fields.Nested(Layout),
+})
+
 Task = api.model('Task', {
     'uuid':  fields.String(attribute='_id'),
     'status': fields.String,
-    'images': fields.Nested(Images),
-    'layout': fields.Nested(Layout),
+    'input': fields.String,
+    'output': fields.Nested(Output),
 })

api/resources/admin/task.py

@@ -9,9 +9,8 @@
 import base64
 from PIL import Image
 from io import BytesIO
+from helpers.storage import save_input, save_aligned, save_layout
 
-STORAGE_FOLDER = os.environ.get('STORAGE_FOLDER')
-STORAGE_BASE_URL = os.environ.get('STORAGE_BASE_URL')
 def base64_to_pil(base64_str):
     image_data = base64.b64decode(base64_str)
     img_buffer = BytesIO(image_data)
@@ -24,26 +23,22 @@ def put(self):
         callback endpoints for worker.
         '''
         data = request.json
-        input_img = base64_to_pil(data['images']['input'])
         aligned_img = base64_to_pil(data['images']['aligned'])
         layout_img = base64_to_pil(data['images']['layout'])
         layout = data['layout']
         id = data['id']
-
-        output = os.path.join(STORAGE_FOLDER,id)
-        os.makedirs(output, exist_ok=True)
-
-        input_img.save(os.path.join(output,"input.jpg"))
-        aligned_img.save(os.path.join(output,"aligned.jpg"))
-        layout_img.save(os.path.join(output,"layout.jpg"))
+        aligned_path = save_aligned(aligned_img, id)
+        layout_path = save_layout(layout_img, id)
 
         db.updateTask(id, {
-            'images': {
-                'input': f'{STORAGE_BASE_URL}/{id}/input.jpg',
-                'aligned': f'{STORAGE_BASE_URL}/{id}/aligned.jpg',
-                'layout': f'{STORAGE_BASE_URL}/{id}/layout.jpg',
-            },
-            'layout': layout
+            'output':{
+                'images': {
+                    'aligned': aligned_path,
+                    'layout': layout_path,
+                },
+                'layout': layout
+            }
+
         })
         if db.getTask(id):
             db.updateTask(id, {

api/resources/task.py

@@ -1,5 +1,7 @@
 import os
+import io
 import uuid
+from PIL import Image
 from worker_helper import inference_from_file
 from api.constant import TASK_STATUS
 
@@ -8,6 +10,7 @@
 from flask import request, jsonify
 from flask_restx import Resource, abort
 from api.models import Request, Task
+from helpers.storage import save_input
 
 postReq = Request()
 postReq.addFile(name='file',required=True)
@@ -66,11 +69,16 @@ def put(self):
                 'status': TASK_STATUS.PROCESSING,
             })
             if isReady:                
-                inference_from_file(file, id)
-                output = {
-                    '_id': id
-                }
-                return output, 200
+                pil_image = Image.open(io.BytesIO(file.read()))
+                image_path = save_input(pil_image, id)
+                db.updateTask(id, {
+                    'input': image_path,
+                    'output': None
+                })
+
+                inference_from_file(pil_image, id)
+
+                return db.getTask(id), 200
             else:
                 return abort(400, message=f'Task {id} is not ready.')
         else:
@@ -83,11 +91,15 @@ def post(self):
         file = args.file
         if file:
             id = db.createTask()
-            inference_from_file(file, id)
+
+            pil_image = Image.open(io.BytesIO(file.read()))
+            image_path = save_input(pil_image, id)
+            db.updateTask(id, {
+                'input': image_path,
+            })
 
-            output = {
-                '_id': id
-            }
-            return output, 200
+            inference_from_file(pil_image, id)
+
+            return db.getTask(id), 200
         else:
             return {}, 400

helpers/storage.py

@@ -0,0 +1,19 @@
+import os
+STORAGE_FOLDER = os.environ.get('STORAGE_FOLDER')
+STORAGE_BASE_URL = os.environ.get('STORAGE_BASE_URL')
+
+def save_pil_image(image, id, name):
+    output = os.path.join(STORAGE_FOLDER,id)
+    os.makedirs(output, exist_ok=True)
+    image.save(os.path.join(output, name))
+    return f'{STORAGE_BASE_URL}/{id}/{name}'
+
+def save_input(image, id):
+    return save_pil_image(image, id, 'input.jpg')
+
+def save_aligned(image, id):
+    return save_pil_image(image, id, 'aligned.jpg')
+
+def save_layout(image, id):
+    return save_pil_image(image, id, 'layout.jpg')
+

worker_helper.py

@@ -1,7 +1,6 @@
 import os
 import io
 import base64
-from PIL import Image
 from celery import Celery
 
 BROKER_URL = os.environ.get('WORKER_BROKER_URL')
@@ -11,8 +10,7 @@
 CALLBACK_URL = os.environ.get('WORKER_CALLBACK_URL')
 app = Celery(APP_NAME, broker=BROKER_URL, backend=BACKEND_URL)
 
-def inference_from_file(file, id, callback_url=CALLBACK_URL):
-    pil_image = Image.open(io.BytesIO(file.read()))
+def inference_from_file(pil_image, id, callback_url=CALLBACK_URL):
     img_bytes = io.BytesIO()
     pil_image.save(img_bytes, format='JPEG')
     img_base64 = base64.b64encode(img_bytes.getvalue()).decode()