Merge pull request #106 from jqueguiner/jqueguiner-add-Docker-video

[Docker][API][Video] add docker video

Dockerfile-api

@@ -2,7 +2,7 @@ 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 apt-get install -y python3-pip software-properties-common wget ffmpeg
 
 RUN add-apt-repository ppa:git-core/ppa
 
@@ -24,8 +24,12 @@ RUN wget -O /root/.torch/models/vgg16_bn-6c64b313.pth https://download.pytorch.o
 
 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

README.md

@@ -253,11 +253,15 @@ Running Docker
 echo "http://$(curl ifconfig.io):5000" && nvidia-docker run --ipc=host -p 5000:5000 -it deoldify_api
 ```
 
-Calling the API
+Calling the API for image processing
 ```console
-curl -X POST "http://MY_SUPER_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
+curl -X POST "http://MY_SUPER_API_IP:5000/process_image" -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
 ```
 
+Calling the API for video processing
+```console
+curl -X POST "http://MY_SUPER_API_IP:5000/process_video" -H "accept: application/octet-stream" -H "Content-Type: application/json" -d "{\"source_url\":\"https://v.redd.it/d1ku57kvuf421/HLSPlaylist.m3u8\", \"render_factor\":35}" --output colorized_video.mp4
+```
 #### Note Regarding Docker
 If you don't have Nvidia Docker, here is the installation guide :
 https://github.com/nvidia/nvidia-docker/wiki/Installation-(version-2.0)#installing-version-20

app.py

@@ -20,15 +20,17 @@ from pathlib import Path
 
 torch.backends.cudnn.benchmark=True
 
-colorizer = get_image_colorizer(artistic=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", methods=["POST"])
-def process():
+
+@app.route("/process_image", methods=["POST"])
+def process_image():
     source_url = request.json["source_url"]
     render_factor = int(request.json["render_factor"])
 
@@ -37,8 +39,8 @@ def process():
            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),
+    
+    image_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')
@@ -48,6 +50,23 @@ def process():
 
     return callback
 
+@app.route("/process_video", methods=["POST"])
+def process_video():
+    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()) + '.mp4'
+
+    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')
+
+    os.remove(os.path.join("video/result/", random_filename))
+
+    return callback
 
 if __name__ == '__main__':
     port = 5000