add optimization server test

Whitebox_style_transfer.py

@@ -144,6 +144,7 @@ def optimize(effect, preset, result_image_placeholder):
     content = st.session_state["Content_im"]
     style = st.session_state["Style_im"]
     result_image_placeholder.text("<- Custom content/style needs to be style transferred")
+    st.sidebar.info("Note: Optimizing takes up to 5 minutes.")
     optimize_button = st.sidebar.button("Optimize Style Transfer")
     if optimize_button:
         if HUGGING_FACE:

pages/test.py

@@ -0,0 +1,119 @@
+import base64
+import datetime
+import os
+import sys
+from io import BytesIO
+from pathlib import Path
+import numpy as np
+import requests
+import torch
+import torch.nn.functional as F
+from PIL import Image
+import json
+import time
+
+PACKAGE_PARENT = 'wise'
+SCRIPT_DIR = os.path.dirname(os.path.realpath(os.path.join(os.getcwd(), os.path.expanduser(__file__))))
+sys.path.append(os.path.normpath(os.path.join(SCRIPT_DIR, PACKAGE_PARENT)))
+
+import streamlit as st
+from streamlit.logger import get_logger
+from st_click_detector import click_detector
+import streamlit.components.v1 as components
+from streamlit.source_util import get_pages
+from streamlit_extras.switch_page_button import switch_page
+
+import helpers.session_state as session_state 
+from parameter_optimization.strotss_org import strotss, pil_resize_long_edge_to
+from helpers import torch_to_np, np_to_torch
+from effects import get_default_settings, MinimalPipelineEffect 
+
+st.set_page_config(layout="wide")
+BASE_URL = "https://ivpg.hpi3d.de/wise/wise-demo/images/"
+LOGGER = get_logger(__name__)
+
+
+def upload_form(imgtype):
+    with st.form(imgtype + "-form", clear_on_submit=True):
+        uploaded_im = st.file_uploader(f"Load {imgtype} image:", type=["png", "jpg"], )
+        upload_pressed = st.form_submit_button("Upload")
+
+        if upload_pressed and uploaded_im is not None:
+            img = Image.open(uploaded_im)
+            buffered = BytesIO()
+            img.save(buffered, format="JPEG")
+            encoded = base64.b64encode(buffered.getvalue()).decode()
+            # session_state.get(uploaded_im=img, content_render_src=f"data:image/jpeg;base64,{encoded}")
+            session_state.get(**{f"{imgtype}_im": img, f"{imgtype}_render_src": f"data:image/jpeg;base64,{encoded}",
+                                    f"{imgtype}_id": "uploaded"})
+            st.session_state["action"] = "uploaded"
+            st.write("uploaded.")
+
+upload_form("Content")
+upload_form("Style")
+content = st.session_state["Content_im"]
+style = st.session_state["Style_im"]
+base_url = "http://mr2632.byod.hpi.de:5000"
+
+if content is not None and style is not None:
+    optimize_button = st.sidebar.button("Optimize Style Transfer")
+    if optimize_button:
+        url = base_url + "/upload"
+        content_path=f"/tmp/content-wise-uploaded{str(datetime.datetime.timestamp(datetime.datetime.now()))}.jpg"
+        style_path=f"/tmp/content-wise-uploaded{str(datetime.datetime.timestamp(datetime.datetime.now()))}.jpg"
+        content = pil_resize_long_edge_to(content, 1024)
+        content.save(content_path)
+        style = pil_resize_long_edge_to(style, 1024)
+        style.save(style_path)
+        files = {'style-image': open(style_path, "rb"), "content-image": open(content_path, "rb")}
+        print("start-optimizing")
+        task_id_res = requests.post(url, files=files)
+        if task_id_res.status_code != 200:
+            st.error(task_id_res.content)
+            st.stop()
+        else:
+            task_id = task_id_res.json()['task_id']
+        
+        progress_bar = st.empty()
+        with st.spinner(text="Optimizing parameters.."):
+            started_time = time.time()
+            while True:
+                time.sleep(3)
+                status = requests.get(base_url+"/get_status", params={"task_id": task_id})
+                if status.status_code != 200:
+                    print("get_status got status_code", status.status_code)
+                    st.warning(status.content)
+                    continue
+                status = status.json()
+                print(status)
+                if status["status"] != "running" and status["status"] != "queued" :
+                    if status["msg"] != "":
+                        st.error(status["msg"])
+                    break
+                elif status["status"] == "queued":
+                    started_time = time.time()
+                    queue_length = requests.get(base_url+"/queue_length").json()
+                    progress_bar.write(f"There are {queue_length['length']} tasks in the queue")
+                elif status["progress"] == 0.0:
+                    progressed = min(0.5 * (time.time() - started_time) / 80.0, 0.5) #estimate 80s for strotts
+                    progress_bar.progress(progressed)
+                else:
+                    progress_bar.progress(min(0.5 + status["progress"] / 2.0, 1.0))
+            vp_res = requests.get(base_url+"/get_vp", params={"task_id": task_id})
+            if vp_res.status_code != 200:
+                st.warning("got status" + str(vp_res.status_code))
+                vp_res.raise_for_status()
+            else:
+                vp = np.load(BytesIO(vp_res.content))["vp"] 
+                print("received vp from server")
+                print("got numpy array", vp.shape)
+                vp = torch.from_numpy(vp).cuda()
+
+                effect, preset, param_set = get_default_settings("minimal_pipeline")
+                effect.enable_checkpoints()
+                effect.cuda()
+                content_cuda = np_to_torch(content).cuda()
+                with torch.no_grad():
+                    result_cuda = effect(content_cuda, vp)
+                img_res = Image.fromarray((torch_to_np(result_cuda) * 255.0).astype(np.uint8))
+                st.image(img_res)

worker/requirements.txt

@@ -0,0 +1,10 @@
+imageio
+imageio-ffmpeg
+scipy
+Pillow
+numpy
+--extra-index-url https://download.pytorch.org/whl/cu113
+torch
+torchvision
+Flask
+Flask-Reuploaded

worker/serve.py

@@ -0,0 +1,284 @@
+import datetime
+import os
+from pathlib import Path
+import sys
+from flask import Flask, jsonify, request, send_file, abort
+from flask_uploads import UploadSet, configure_uploads, IMAGES
+from werkzeug.exceptions import default_exceptions
+from werkzeug.exceptions import HTTPException, NotFound 
+import json
+import torch
+import time
+import threading
+import traceback
+from PIL import Image
+import numpy as np
+
+PACKAGE_PARENT = '..'
+WISE_DIR = '../wise/'
+SCRIPT_DIR = os.path.dirname(os.path.realpath(os.path.join(os.getcwd(), os.path.expanduser(__file__))))
+sys.path.append(os.path.normpath(os.path.join(SCRIPT_DIR, PACKAGE_PARENT)))
+sys.path.append(os.path.normpath(os.path.join(SCRIPT_DIR, WISE_DIR)))
+
+
+
+from parameter_optimization.parametric_styletransfer import single_optimize
+from parameter_optimization.parametric_styletransfer import CONFIG as ST_CONFIG
+from parameter_optimization.strotss_org import strotss, pil_resize_long_edge_to
+from helpers import torch_to_np, np_to_torch
+from effects import get_default_settings, MinimalPipelineEffect 
+
+class JSONExceptionHandler(object):
+
+    def __init__(self, app=None):
+        if app:
+            self.init_app(app)
+
+    def std_handler(self, error):
+        response = jsonify(message=error.message)
+        response.status_code = error.code if isinstance(error, HTTPException) else 500
+        return response
+
+
+    def init_app(self, app):
+        self.app = app
+        self.register(HTTPException)
+        for code, v in default_exceptions.items():
+            self.register(code)
+
+    def register(self, exception_or_code, handler=None):
+        self.app.errorhandler(exception_or_code)(handler or self.std_handler)
+
+
+
+app = Flask(__name__)
+handler = JSONExceptionHandler(app)
+
+image_folder = 'img_received'
+photos = UploadSet('photos', IMAGES)
+app.config['UPLOADED_PHOTOS_DEST'] = image_folder
+configure_uploads(app, photos)
+
+class Args(object):
+    def __init__(self, initial_data):
+        for key in initial_data:
+            setattr(self, key, initial_data[key])
+    def set_attributes(self, val_dict):
+        for key in val_dict:
+            setattr(self, key, val_dict[key])
+
+default_args = {
+        "output_image" : "output.jpg",
+        ## values always set by request ##
+        "content_image": "",
+        "style_image": "",
+        "output_vp": "",
+        "iters": 500
+}
+
+
+total_task_count = 0
+
+class NeuralOptimizer():
+    def __init__(self, args) -> None:
+        self.cur_iteration = 0
+        self.args = args
+
+    def optimize(self):
+        base_dir = f"result/{datetime.datetime.now().strftime(r'%Y-%m-%d %H.%Mh %Ss')}"
+        os.makedirs(base_dir)
+
+        content = Image.open(self.args.content_image)
+        style = Image.open(self.args.style_image)
+
+        def set_iter(iter):
+            self.cur_iteration=iter
+
+        effect, preset, _ = get_default_settings("minimal_pipeline")
+        effect.enable_checkpoints()
+
+        reference = strotss(pil_resize_long_edge_to(content, 1024),
+                            pil_resize_long_edge_to(style, 1024), content_weight=16.0,
+                            device=torch.device("cuda"), space="uniform")
+
+        ref_save_path = os.path.join(base_dir, "reference.jpg")
+        resize_to = 720
+        reference = pil_resize_long_edge_to(reference, resize_to)
+        reference.save(ref_save_path)
+        ST_CONFIG["n_iterations"] = self.args.iters
+        vp, content_img_cuda = single_optimize(effect, preset, "l1", self.args.content_image, str(ref_save_path),
+                                            write_video=False, base_dir=base_dir,
+                                            iter_callback=set_iter)
+
+        output = Image.fromarray(torch_to_np(content_img_cuda.detach().cpu() * 255.0).astype(np.uint8))
+        output.save(self.args.output_image)
+        # torch.save (vp.detach().clone(), self.args.output_vp)
+        # preset_tensor = effect.vpd.preset_tensor(preset, np_to_torch(np.array(content)).cuda(), add_local_dims=True)
+        np.savez_compressed(self.args.output_vp, vp=vp.detach().cpu().numpy())
+
+        
+
+class StyleTask:
+    def __init__(self, task_id, style_filename, content_filename):
+        self.content_filename=content_filename
+        self.style_filename=style_filename
+
+        self.status = "queued"
+        self.task_id = task_id
+        self.error_msg = ""
+        self.output_filename = content_filename.split(".")[0] + "_output.jpg"
+        self.vp_output_filename = content_filename.split(".")[0] + "_output.npz"
+
+        # global neural_optimizer
+        # if neural_optimizer is None:
+        #     neural_optimizer = NeuralOptimizer(Args(default_args))
+
+        self.neural_optimizer = NeuralOptimizer(Args(default_args))
+    
+    def start(self):
+        self.neural_optimizer.args.set_attributes(default_args)
+
+        self.neural_optimizer.args.style_image = os.path.join(image_folder, self.style_filename)
+        self.neural_optimizer.args.content_image = os.path.join(image_folder, self.content_filename)
+        self.neural_optimizer.args.output_image = os.path.join(image_folder, self.output_filename)
+        self.neural_optimizer.args.output_vp = os.path.join(image_folder, self.vp_output_filename)
+ 
+        thread = threading.Thread(target=self.run, args=())
+        thread.daemon = True                            # Daemonize thread
+        thread.start()                                  # Start the execution
+
+    def run(self):
+        self.status = "running"
+        try:
+            self.neural_optimizer.optimize()
+        except Exception as e:
+            print("Error in task %d :"%(self.task_id), str(e))
+            traceback.print_exc()
+
+            self.status = "error"
+            self.error_msg = str(e)
+            return
+
+        self.status = "finished"
+        print("finished styling task: " + str(self.task_id))
+
+class StylerQueue:
+    queued_tasks = []
+    finished_tasks = []
+    running_task = None
+
+    def __init__(self):
+        thread = threading.Thread(target=self.status_checker, args=())
+        thread.daemon = True                            # Daemonize thread
+        thread.start()                                  # Start the execution
+
+    def queue_task(self, *args):
+        global total_task_count
+        total_task_count += 1
+        task = StyleTask(total_task_count, *args)
+        self.queued_tasks.append(task)
+
+        return total_task_count
+
+    def get_task(self, task_id):
+        if self.running_task is not None and self.running_task.task_id == task_id:
+            return self.running_task
+        task = next((task for task in self.queued_tasks + self.finished_tasks if task.task_id == task_id), None)
+        return task
+
+    def status_checker(self):
+        while True:
+            time.sleep(0.3)
+
+            if self.running_task is None:
+                if len(self.queued_tasks) > 0:
+                    self.running_task = self.queued_tasks[0]
+                    self.running_task.start()
+                    self.queued_tasks = self.queued_tasks[1:]
+            elif self.running_task.status == "finished" or self.running_task.status == "error": 
+                self.finished_tasks.append(self.running_task)
+                if len(self.queued_tasks) > 0:
+                    self.running_task = self.queued_tasks[0]
+                    self.running_task.start()
+                    self.queued_tasks = self.queued_tasks[1:]
+                else:
+                    self.running_task = None
+
+styler_queue = StylerQueue()
+
+
+@app.route('/upload', methods=['POST'])
+def upload():
+    if 'style-image' in request.files and \
+        'content-image' in request.files:
+
+        style_filename = photos.save(request.files['style-image'])
+        content_filename = photos.save(request.files['content-image'])
+
+        job_id = styler_queue.queue_task(style_filename, content_filename)
+        print('added new stylization task', style_filename, content_filename)
+
+        return jsonify({"task_id": job_id})
+    abort(jsonify(message="request needs style, content image"), 400)
+
+@app.route('/get_status')
+def get_status():
+    task_id = int(request.args.get("task_id"))
+    task = styler_queue.get_task(task_id)
+
+    if task is None:
+        abort(jsonify(message="task with id %d not found"%task_id), 400) 
+
+    status = {
+        "status": task.status,
+        "msg": task.error_msg
+    }
+
+    if task.status == "running":
+        if isinstance(task, StyleTask):
+            status["progress"] = float(task.neural_optimizer.cur_iteration) / float(default_args["iters"]) 
+
+    return jsonify(status)
+
+@app.route('/queue_length')
+def get_queue_length():
+    tasks = len(styler_queue.queued_tasks)
+    if styler_queue.running_task is not None:
+        tasks += 1
+
+    status = {
+        "length": tasks
+    }
+
+    return jsonify(status)
+
+
+@app.route('/get_image')
+def get_image():
+    task_id = int(request.args.get("task_id"))
+    task = styler_queue.get_task(task_id)
+
+    if task is None:
+        abort(jsonify(message="task with id %d not found"%task_id), 400)
+
+    if task.status != "finished":
+        abort(jsonify(message="task with id %d not in finished state"%task_id), 400)
+
+    return send_file(os.path.join(image_folder, task.output_filename), mimetype='image/jpg')
+
+@app.route('/get_vp')
+def get_vp():
+    task_id = int(request.args.get("task_id"))
+    task = styler_queue.get_task(task_id)
+
+    if task is None:
+        abort(jsonify(message="task with id %d not found"%task_id), 400)
+
+    if task.status != "finished":
+        abort(jsonify(message="task with id %d not in finished state"%task_id), 400)
+
+    return send_file(os.path.join(image_folder, task.vp_output_filename), mimetype='application/zip')
+
+
+if __name__ == '__main__':
+    app.run(debug=False, host="0.0.0.0",port=5000)