Эх сурвалжийг харах

add api support for image colorization

jqueguiner 6 жил өмнө
parent
commit
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.  
 
-#### 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
 ```console
@@ -226,7 +226,7 @@ git clone https://github.com/jantic/DeOldify.git DeOldify
 
 Building Docker
 ```console
-cd DeOldify && docker build -t deoldify .
+cd DeOldify && docker build -t deoldify -f Dockerfile .
 ```
 
 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
 ```
 
+#### 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 :
 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)