split apis

Dockerfile-api

@@ -4,32 +4,12 @@ RUN apt-get -y update
 
 RUN apt-get install -y python3-pip software-properties-common wget ffmpeg
 
-RUN add-apt-repository ppa:git-core/ppa
-
 RUN apt-get -y update
 
-RUN curl -s https://packagecloud.io/install/repositories/github/git-lfs/script.deb.sh | bash
-
-RUN apt-get install -y git-lfs --allow-unauthenticated
-
-RUN git lfs install
-
-ENV GIT_WORK_TREE=/data
-
 RUN mkdir -p /root/.torch/models
 
 RUN mkdir -p /data/models
 
-RUN wget -O /root/.torch/models/vgg16_bn-6c64b313.pth https://download.pytorch.org/models/vgg16_bn-6c64b313.pth
-
-RUN wget -O /root/.torch/models/resnet34-333f7ec4.pth https://download.pytorch.org/models/resnet34-333f7ec4.pth
-
-RUN wget -O /root/.torch/models/resnet101-5d3b4d8f.pth https://download.pytorch.org/models/resnet101-5d3b4d8f.pth
-
-RUN wget -O /data/models/ColorizeArtistic_gen.pth https://www.dropbox.com/s/zkehq1uwahhbc2o/ColorizeArtistic_gen.pth?dl=0 
-
-RUN wget -O /data/models/ColorizeVideo_gen.pth https://www.dropbox.com/s/336vn9y4qwyg9yz/ColorizeVideo_gen.pth?dl=0
-
 ADD . /data/
 
 WORKDIR /data
@@ -38,9 +18,15 @@ RUN pip install -r requirements.txt
 
 RUN pip install  Flask
 
-RUN cd /data/test_images && git lfs pull
+RUN pip install Pillow
+
+RUN pip install scikit-image
+
+RUN pip install requests
 
 EXPOSE 5000
 
-ENTRYPOINT ["python3", "app.py"]
+ENTRYPOINT ["python3"]
+
+CMD ["app.py"]
 

app-video.py

@@ -0,0 +1,86 @@
+# import the necessary packages
+import os
+import sys
+import requests
+import ssl
+from flask import Flask
+from flask import request
+from flask import jsonify
+from flask import send_file
+
+
+from app_utils import download
+from app_utils import generate_random_filename
+from app_utils import clean_me
+from app_utils import clean_all
+from app_utils import create_directory
+from app_utils import get_model_bin
+from app_utils import convertToJPG
+
+from os import path
+import torch
+
+import fastai
+from fasterai.visualize import *
+from pathlib import Path
+import traceback
+
+
+torch.backends.cudnn.benchmark=True
+
+
+os.environ['CUDA_VISIBLE_DEVICES']='0'
+
+app = Flask(__name__)
+
+
+
+# define a predict function as an endpoint
+@app.route("/process", methods=["POST"])
+def process_video():
+
+    input_path = generate_random_filename(upload_directory,"mp4")
+    output_path = os.path.join(results_video_directory, os.path.basename(input_path))
+
+    try:
+        url = request.json["source_url"]
+        render_factor = int(request.json["render_factor"])
+
+        video_path = video_colorizer.colorize_from_url(source_url=url, file_name=input_path, render_factor=render_factor)
+        callback = send_file(output_path, mimetype='application/octet-stream')
+
+        return callback, 200
+
+    except:
+        traceback.print_exc()
+        return {'message': 'input error'}, 400
+
+    finally:
+        clean_all([
+            input_path,
+            output_path
+            ])
+
+if __name__ == '__main__':
+    global upload_directory
+    global results_video_directory
+    global video_colorizer
+
+    upload_directory = '/data/upload/'
+    create_directory(upload_directory)
+
+    results_video_directory = '/data/video/result/'
+    create_directory(results_video_directory)
+
+    model_directory = '/data/models/'
+    create_directory(model_directory)
+    
+    video_model_url = 'https://www.dropbox.com/s/336vn9y4qwyg9yz/ColorizeVideo_gen.pth?dl=0'
+    get_model_bin(video_model_url, os.path.join(model_directory, 'ColorizeVideo_gen.pth'))
+
+    video_colorizer = get_video_colorizer()
+    
+    port = 5000
+    host = '0.0.0.0'
+
+    app.run(host=host, port=port, threaded=False)

app.py

@@ -8,7 +8,14 @@ from flask import request
 from flask import jsonify
 from flask import send_file
 
-from uuid import uuid4
+
+from app_utils import download
+from app_utils import generate_random_filename
+from app_utils import clean_me
+from app_utils import clean_all
+from app_utils import create_directory
+from app_utils import get_model_bin
+from app_utils import convertToJPG
 
 from os import path
 import torch
@@ -21,68 +28,70 @@ import traceback
 
 torch.backends.cudnn.benchmark=True
 
-image_colorizer = get_image_colorizer(artistic=True)
-video_colorizer = get_video_colorizer()
 
 os.environ['CUDA_VISIBLE_DEVICES']='0'
 
 app = Flask(__name__)
 
-# define a predict function as an endpoint
 
-@app.route("/process_image", methods=["POST"])
+
+# define a predict function as an endpoint
+@app.route("/process", methods=["POST"])
 def process_image():
+
+    input_path = generate_random_filename(upload_directory,"jpeg")
+    output_path = os.path.join(results_img_directory, os.path.basename(input_path))
+
     try:
-        source_url = request.json["source_url"]
+        url = request.json["source_url"]
         render_factor = int(request.json["render_factor"])
 
-        upload_directory = 'upload'
-        if not os.path.exists(upload_directory):
-            os.mkdir(upload_directory)
+        download(url, input_path)
 
-        random_filename = str(uuid4()) + '.png'
-    
-        image_colorizer.plot_transformed_image_from_url(url=source_url, path=os.path.join(upload_directory, random_filename), figsize=(20,20),
+        try:
+            image_colorizer.plot_transformed_image(path=input_path, figsize=(20,20),
+                render_factor=render_factor, display_render_factor=True, compare=False)
+        except:
+            convertToJPG(input_path)
+            image_colorizer.plot_transformed_image(path=input_path, figsize=(20,20),
             render_factor=render_factor, display_render_factor=True, compare=False)
 
-        callback = send_file(os.path.join("result_images", random_filename), mimetype='image/jpeg')
+        callback = send_file(output_path, mimetype='image/jpeg')
         
-        return callback
+        return callback, 200
 
     except:
         traceback.print_exc()
         return {'message': 'input error'}, 400
 
     finally:
-        os.remove(os.path.join("result_images", random_filename))
-        os.remove(os.path.join("upload", random_filename))
-
-
-@app.route("/process_video", methods=["POST"])
-def process_video():
-    try:
-        source_url = request.json["source_url"]
-        render_factor = int(request.json["render_factor"])
+        pass
+        clean_all([
+            input_path,
+            output_path
+            ])
 
-        upload_directory = 'upload'
-        if not os.path.exists(upload_directory):
-            os.mkdir(upload_directory)
-
-        random_filename = str(uuid4()) + '.mp4'
+if __name__ == '__main__':
+    global upload_directory
+    global results_img_directory
+    global image_colorizer
 
-        video_path = video_colorizer.colorize_from_url(source_url, random_filename, render_factor)
-        callback = send_file(os.path.join("video/result/", random_filename), mimetype='application/octet-stream')
+    upload_directory = '/data/upload/'
+    create_directory(upload_directory)
 
-        return callback
+    results_img_directory = '/data/result_images/'
+    create_directory(results_img_directory)
 
-    except:
-        traceback.print_exc()
-        return {'message': 'input error'}, 400
+    model_directory = '/data/models/'
+    create_directory(model_directory)
+    
+    artistic_model_url = 'https://www.dropbox.com/s/zkehq1uwahhbc2o/ColorizeArtistic_gen.pth?dl=0'
+    get_model_bin(artistic_model_url, os.path.join(model_directory, 'ColorizeArtistic_gen.pth'))
 
-    finally:
-        os.remove(os.path.join("video/result/", random_filename))
 
-if __name__ == '__main__':
+    image_colorizer = get_image_colorizer(artistic=True)
+    
     port = 5000
     host = '0.0.0.0'
-    app.run(host=host, port=port, threaded=True)
+
+    app.run(host=host, port=port, threaded=False)

app_utils.py

@@ -0,0 +1,126 @@
+import os
+import requests
+import random
+import _thread as thread
+from uuid import uuid4
+
+import numpy as np
+import skimage
+from skimage.filters import gaussian
+from PIL import Image
+
+def compress_image(image, path_original):
+    size = 1920, 1080
+    width = 1920
+    height = 1080
+
+    name = os.path.basename(path_original).split('.')
+    first_name = os.path.join(os.path.dirname(path_original), name[0] + '.jpg')
+
+    if image.size[0] > width and image.size[1] > height:
+        image.thumbnail(size, Image.ANTIALIAS)
+        image.save(first_name, quality=85)
+    elif image.size[0] > width:
+        wpercent = (width/float(image.size[0]))
+        height = int((float(image.size[1])*float(wpercent)))
+        image = image.resize((width,height), PIL.Image.ANTIALIAS)
+        image.save(first_name,quality=85)
+    elif image.size[1] > height:
+        wpercent = (height/float(image.size[1]))
+        width = int((float(image.size[0])*float(wpercent)))
+        image = image.resize((width,height), Image.ANTIALIAS)
+        image.save(first_name, quality=85)
+    else:
+        image.save(first_name, quality=85)
+
+
+def convertToJPG(path_original):
+    img = Image.open(path_original)
+    name = os.path.basename(path_original).split('.')
+    first_name = os.path.join(os.path.dirname(path_original), name[0] + '.jpg')
+
+    if img.format == "JPEG":
+        image = img.convert('RGB')
+        compress_image(image, path_original)
+        img.close()
+
+    elif img.format == "GIF":
+        i = img.convert("RGBA")
+        bg = Image.new("RGBA", i.size)
+        image = Image.composite(i, bg, i)
+        compress_image(image, path_original)
+        img.close()
+
+    elif img.format == "PNG":
+        try:
+            image = Image.new("RGB", img.size, (255,255,255))
+            image.paste(img,img)
+            compress_image(image, path_original)
+        except ValueError:
+            image = img.convert('RGB')
+            compress_image(image, path_original)
+        
+        img.close()
+
+    elif img.format == "BMP":
+        image = img.convert('RGB')
+        compress_image(image, path_original)
+        img.close()
+
+
+
+def blur(image, x0, x1, y0, y1, sigma=1, multichannel=True):
+    y0, y1 = min(y0, y1), max(y0, y1)
+    x0, x1 = min(x0, x1), max(x0, x1)
+    im = image.copy()
+    sub_im = im[y0:y1,x0:x1].copy()
+    blur_sub_im = gaussian(sub_im, sigma=sigma, multichannel=multichannel)
+    blur_sub_im = np.round(255 * blur_sub_im)
+    im[y0:y1,x0:x1] = blur_sub_im
+    return im
+
+
+
+def download(url, filename):
+    data = requests.get(url).content
+    with open(filename, 'wb') as handler:
+        handler.write(data)
+
+    return filename
+
+
+def generate_random_filename(upload_directory, extension):
+    filename = str(uuid4())
+    filename = os.path.join(upload_directory, filename + "." + extension)
+    return filename
+
+
+def clean_me(filename):
+    if os.path.exists(filename):
+        os.remove(filename)
+
+
+def clean_all(files):
+    for me in files:
+        clean_me(me)
+
+
+def create_directory(path):
+    os.system("mkdir -p %s" % os.path.dirname(path))
+
+
+def get_model_bin(url, output_path):
+    if not os.path.exists(output_path):
+        create_directory(output_path)
+        cmd = "wget -O %s %s" % (output_path, url)
+        print(cmd)
+        os.system(cmd)
+
+    return output_path
+
+
+#model_list = [(url, output_path), (url, output_path)]
+def get_multi_model_bin(model_list):
+    for m in model_list:
+        thread.start_new_thread(get_model_bin, m)
+

requirements.txt

@@ -5,3 +5,4 @@ ffmpeg-python==0.1.17
 youtube-dl>=2019.4.17
 jupyterlab
 opencv-python>=3.3.0.10
+pillow