瀏覽代碼

add api support for image colorization

jqueguiner 6 年之前
父節點
當前提交
77b923d6df
共有 3 個文件被更改,包括 124 次插入3 次删除
  1. 42 0
      Dockerfile-api
  2. 27 3
      README.md
  3. 55 0
      app.py

+ 42 - 0
Dockerfile-api

@@ -0,0 +1,42 @@
+From nvcr.io/nvidia/pytorch:19.04-py3
+
+RUN apt-get -y update
+
+RUN apt-get install -y python3-pip software-properties-common wget
+
+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 /data/models/ColorizeArtistic_gen.pth https://www.dropbox.com/s/zkehq1uwahhbc2o/ColorizeArtistic_gen.pth?dl=0 
+
+ADD . /data/
+
+WORKDIR /data
+
+RUN pip install -r requirements.txt
+
+RUN pip install  Flask
+
+RUN cd /data/test_images && git lfs pull
+
+EXPOSE 5000
+
+#ENTRYPOINT ["python3", "app.py"]
+

+ 27 - 3
README.md

@@ -215,9 +215,9 @@ jupyter lab
 
 
 From there you can start running the notebooks in Jupyter Lab, via the url they provide you in the console.  
 From there you can start running the notebooks in Jupyter Lab, via the url they provide you in the console.  
 
 
-#### Docker
+#### Docker for jupyter
 
 
-You can build and run the docker using the foloowing process:
+You can build and run the docker using the following process:
 
 
 Cloning
 Cloning
 ```console
 ```console
@@ -226,7 +226,7 @@ git clone https://github.com/jantic/DeOldify.git DeOldify
 
 
 Building Docker
 Building Docker
 ```console
 ```console
-cd DeOldify && docker build -t deoldify .
+cd DeOldify && docker build -t deoldify -f Dockerfile .
 ```
 ```
 
 
 Running Docker
 Running Docker
@@ -234,6 +234,30 @@ Running Docker
 echo "http://$(curl ifconfig.io):8888" && nvidia-docker run --ipc=host --env NOTEBOOK_PASSWORD="pass123" -p 8888:8888 -it deoldify
 echo "http://$(curl ifconfig.io):8888" && nvidia-docker run --ipc=host --env NOTEBOOK_PASSWORD="pass123" -p 8888:8888 -it deoldify
 ```
 ```
 
 
+#### Docker for api
+You can build and run the docker using the following process:
+
+Cloning
+```console
+git clone https://github.com/jantic/DeOldify.git DeOldify
+```
+
+Building Docker
+```console
+cd DeOldify && docker build -t deoldify_api -f Dockerfile-api .
+```
+
+Running Docker
+```console
+echo "http://$(curl ifconfig.io):5000" && nvidia-docker run --ipc=host -p 5000:5000 -it deoldify_api
+```
+
+Calling the api for colorization
+```console
+curl -X POST "http:/MY_API_IP:5000/process" -H "accept: image/png" -H "Content-Type: application/json" -d "{\"source_url\":\"http://www.afrikanheritage.com/wp-content/uploads/2015/08/slave-family-P.jpeg\", \"render_factor\":35}" --output colorized_image.png
+```
+#### Note Regarding Docker
+
 If you don't have Nvidia Docker here the installation guide :
 If you don't have Nvidia Docker here the installation guide :
 https://github.com/nvidia/nvidia-docker/wiki/Installation-(version-2.0)#installing-version-20
 https://github.com/nvidia/nvidia-docker/wiki/Installation-(version-2.0)#installing-version-20
 
 

+ 55 - 0
app.py

@@ -0,0 +1,55 @@
+# 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 uuid import uuid4
+
+from os import path
+import torch
+
+import fastai
+from fasterai.visualize import *
+from pathlib import Path
+
+
+torch.backends.cudnn.benchmark=True
+
+colorizer = get_image_colorizer(artistic=True)
+
+os.environ['CUDA_VISIBLE_DEVICES']='0'
+
+app = Flask(__name__)
+
+# define a predict function as an endpoint
+@app.route("/process", methods=["POST"])
+def process():
+    source_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)
+
+    random_filename = str(uuid4()) + '.png'
+
+    colorizer.plot_transformed_image_from_url(url=source_url, path=os.path.join(upload_directory, random_filename), 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')
+
+    os.remove(os.path.join("result_images", random_filename))
+    os.remove(os.path.join("upload", random_filename))
+
+    return callback
+
+
+if __name__ == '__main__':
+    port = 5000
+    host = '0.0.0.0'
+    app.run(host=host, port=port, threaded=True)