mirror of
https://github.com/jantic/DeOldify.git
synced 2026-08-29 02:10:55 +08:00
Merge pull request #134 from jqueguiner/jqueguiner-split-apis
Jqueguiner split apis
This commit is contained in:
Jason Antic
+8
-22
@@ -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"]
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -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
|
||||
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"])
|
||||
|
||||
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')
|
||||
|
||||
return callback
|
||||
|
||||
except:
|
||||
traceback.print_exc()
|
||||
return {message: 'input error'}, 400
|
||||
|
||||
finally:
|
||||
os.remove(os.path.join("video/result/", random_filename))
|
||||
pass
|
||||
clean_all([
|
||||
input_path,
|
||||
output_path
|
||||
])
|
||||
|
||||
if __name__ == '__main__':
|
||||
global upload_directory
|
||||
global results_img_directory
|
||||
global image_colorizer
|
||||
|
||||
upload_directory = '/data/upload/'
|
||||
create_directory(upload_directory)
|
||||
|
||||
results_img_directory = '/data/result_images/'
|
||||
create_directory(results_img_directory)
|
||||
|
||||
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'))
|
||||
|
||||
|
||||
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)
|
||||
|
||||
+126
@@ -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)
|
||||
|
||||
@@ -5,3 +5,4 @@ ffmpeg-python==0.1.17
|
||||
youtube-dl>=2019.4.17
|
||||
jupyterlab
|
||||
opencv-python>=3.3.0.10
|
||||
pillow
|
||||
|
||||
Reference in New Issue
Block a user