opensource/DeOldify
3
0
Fork 0
mirror of https://github.com/jantic/DeOldify.git synced 2026-08-30 18:02:24 +08:00
Code Issues Packages Projects Releases Wiki Activity

a few wrapper to simplify onboarding

Browse Source
This commit is contained in:
jqueguiner
2020-11-17 09:38:31 +01:00
parent a1296f58a7
commit b52d7868ff
8 changed files with 211 additions and 54 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Dockerfile
+47 -19
View File
@@ -1,20 +1,10 @@
From nvcr.io/nvidia/pytorch:19.04-py3
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 apt-get -y update && apt-get install -y \
python3-pip \
software-properties-common \
wget \
ffmpeg
RUN mkdir -p /root/.torch/models
@@ -24,17 +14,55 @@ 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 /data/models/ColorizeArtistic_gen.pth https://www.dropbox.com/s/zkehq1uwahhbc2o/ColorizeArtistic_gen.pth?dl=0
# if you want to avoid image building with downloading put your .pth file in root folder
COPY Dockerfile ColorizeArtistic_gen.* /data/models/
COPY Dockerfile ColorizeVideo_gen.* /data/models/
RUN pip install --upgrade pip
RUN pip install versioneer==0.18
RUN pip install tensorboardX==1.6
RUN pip install Flask==1.1.1
RUN pip install pillow==6.1
RUN pip install numpy==1.16
RUN pip install scikit-image==0.15.0
RUN pip install requests==2.21.0
RUN pip install ffmpeg==1.4
RUN pip install ffmpeg-python==0.1.17
RUN pip install youtube-dl>=2019.4.17
RUN pip install jupyterlab==1.2.4
RUN pip install opencv-python>=3.3.0.10
RUN pip install fastai==1.0.51
ADD . /data/
WORKDIR /data
RUN pip install -r requirements.txt
# force download of file if not provided by local cache
RUN [[ ! -f /data/models/ColorizeArtistic_gen.pth ]] && wget -O /data/models/ColorizeArtistic_gen.pth https://data.deepai.org/deoldify/ColorizeArtistic_gen.pth
RUN [[ ! -f /data/models/ColorizeVideo_gen.pth ]] && wget -O /data/models/ColorizeVideo_gen.pth https://data.deepai.org/deoldify/ColorizeVideo_gen.pth
RUN cd /data/test_images && git lfs pull
COPY run_notebook.sh /usr/local/bin/run_notebook
COPY run_image_api.sh /usr/local/bin/run_image_api
COPY run_video_api.sh /usr/local/bin/run_video_api
RUN chmod +x /usr/local/bin/run_notebook
RUN chmod +x /usr/local/bin/run_image_api
RUN chmod +x /usr/local/bin/run_video_api
EXPOSE 8888
EXPOSE 5000
ENTRYPOINT ["sh", "/data/run_notebook.sh"]
# run notebook
# ENTRYPOINT ["sh", "run_notebook"]
# run image api
# ENTRYPOINT ["sh", "run_image_api"]
# run image api
# ENTRYPOINT ["sh", "run_video_api"]
api_cmd_example.sh
Executable
+40
View File
@@ -0,0 +1,40 @@
#!/bin/bash
echo '
_____ ____ _ _ _ __
| __ \ / __ \| | | (_)/ _|
| | | | ___| | | | | __| |_| |_ _ _
| | | |/ _ \ | | | |/ _` | | _| | | |
| |__| | __/ |__| | | (_| | | | | |_| |
|_____/ \___|\____/|_|\__,_|_|_| \__, |
__/ |
|___/
'
echo "usage : $0 image -- to test image api"
echo "usage : $0 video -- to test video api"
echo ''
echo 'you can add non mandatory arguments'
echo "usage : $0 image $port $host -- for custom port or host"
echo ''
echo ''
if [ "$2" == "" ]; then
port=5000
else
port=$2
fi
if [ "$3" == "" ]; then
host="127.0.0.1"
else
host="$3"
fi
if [ "$1" == "video" ]; then
echo "testing deOldify Video API on $host:$port"
curl -X POST "http://$host:$port/process" -H "accept: application/octet-stream" -H "Content-Type: application/json" -d "{\"url\":\"https://v.redd.it/d1ku57kvuf421/HLSPlaylist.m3u8\", \"render_factor\":35}" --output colorized_video.mp4
elif [ "$1" == "image" ]; then
echo "testing deOldify Image API on $host:$port"
curl -X POST "http://$host:$port/process" -H "accept: image/png" -H "Content-Type: application/json" -d "{\"url\":\"http://www.afrikanheritage.com/wp-content/uploads/2015/08/slave-family-P.jpeg\", \"render_factor\":35}" --output colorized_image.png
fi
app-video.py
+50 -21
View File
@@ -26,62 +26,91 @@ from pathlib import Path
import traceback
torch.backends.cudnn.benchmark=True
os.environ['CUDA_VISIBLE_DEVICES']='0'
# Handle switch between GPU and CPU
if torch.cuda.is_available():
torch.backends.cudnn.benchmark = True
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
else:
del os.environ["CUDA_VISIBLE_DEVICES"]
app = Flask(__name__)
def allowed_file(filename):
return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
# define a predict function as an endpoint
@app.route("/process", methods=["POST"])
def process_video():
input_path = generate_random_filename(upload_directory,"mp4")
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"])
if 'file' in request.files:
file = request.files['file']
if allowed_file(file.filename):
file.save(input_path)
try:
render_factor = request.form.getlist('render_factor')[0]
except:
render_factor = 30
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')
else:
url = request.json["url"]
download(url, input_path)
try:
render_factor = request.json["render_factor"]
except:
render_factor = 30
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
return {"message": "input error"}, 400
finally:
clean_all([
input_path,
output_path
])
clean_all([input_path, output_path])
if __name__ == '__main__':
global upload_directory
global results_video_directory
global video_colorizer
global ALLOWED_EXTENSIONS
ALLOWED_EXTENSIONS = set(['mp4'])
upload_directory = '/data/upload/'
upload_directory = "/data/upload/"
create_directory(upload_directory)
results_video_directory = '/data/video/result/'
results_video_directory = "/data/video/result/"
create_directory(results_video_directory)
model_directory = '/data/models/'
model_directory = "/data/models/"
create_directory(model_directory)
video_model_url = (
"https://data.deepai.org/deoldify/ColorizeVideo_gen.pth"
)
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'))
get_model_bin(
video_model_url, os.path.join(model_directory, "ColorizeVideo_gen.pth")
)
video_colorizer = get_video_colorizer()
video_colorizer.result_folder = Path(results_video_directory)
port = 5000
host = '0.0.0.0'
host = "0.0.0.0"
app.run(host=host, port=port, threaded=False)
app.py
+36 -14
View File
@@ -26,14 +26,19 @@ from pathlib import Path
import traceback
torch.backends.cudnn.benchmark=True
# Handle switch between GPU and CPU
if torch.cuda.is_available():
torch.backends.cudnn.benchmark = True
os.environ["CUDA_VISIBLE_DEVICES"] = "0"
else:
del os.environ["CUDA_VISIBLE_DEVICES"]
os.environ['CUDA_VISIBLE_DEVICES']='0'
app = Flask(__name__)
def allowed_file(filename):
return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS
# define a predict function as an endpoint
@app.route("/process", methods=["POST"])
@@ -43,18 +48,31 @@ def process_image():
output_path = os.path.join(results_img_directory, os.path.basename(input_path))
try:
url = request.json["source_url"]
render_factor = int(request.json["render_factor"])
if 'file' in request.files:
file = request.files['file']
if allowed_file(file.filename):
file.save(input_path)
try:
render_factor = request.form.getlist('render_factor')[0]
except:
render_factor = 30
else:
url = request.json["url"]
download(url, input_path)
download(url, input_path)
try:
render_factor = request.json["render_factor"]
except:
render_factor = 30
try:
image_colorizer.plot_transformed_image(path=input_path, figsize=(20,20),
render_factor=render_factor, display_render_factor=True, compare=False)
render_factor=int(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)
render_factor=int(render_factor), display_render_factor=True, compare=False)
callback = send_file(output_path, mimetype='image/jpeg')
@@ -75,6 +93,8 @@ if __name__ == '__main__':
global upload_directory
global results_img_directory
global image_colorizer
global ALLOWED_EXTENSIONS
ALLOWED_EXTENSIONS = set(['png', 'jpg', 'jpeg'])
upload_directory = '/data/upload/'
create_directory(upload_directory)
@@ -84,15 +104,17 @@ if __name__ == '__main__':
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'))
artistic_model_url = "https://data.deepai.org/deoldify/ColorizeArtistic_gen.pth"
# only get the model binay if it not present in /data/models
get_model_bin(
artistic_model_url, os.path.join(model_directory, "ColorizeArtistic_gen.pth")
)
image_colorizer = get_image_colorizer(artistic=True)
image_colorizer.results_dir = Path(results_img_directory)
port = 5000
host = '0.0.0.0'
host = "0.0.0.0"
app.run(host=host, port=port, threaded=False)
quick_start.sh
Executable
+33
View File
@@ -0,0 +1,33 @@
#!/bin/bash
function usage {
echo '
_____ ____ _ _ _ __
| __ \ / __ \| | | (_)/ _|
| | | | ___| | | | | __| |_| |_ _ _
| | | |/ _ \ | | | |/ _` | | _| | | |
| |__| | __/ |__| | | (_| | | | | |_| |
|_____/ \___|\____/|_|\__,_|_|_| \__, |
__/ |
|___/
'
echo "usage : $0 notebook password -- to start the notebook with password"
echo " leave empty for no password (not recommended)"
echo "usage : $0 image_api -- to start image api"
echo "usage : $0 video_api -- to start video api"
echo ''
echo 'you can add non mandatory arguments'
echo "usage : $0 image_api $port $host -- for custom port or host"
echo ''
echo ''
}
NOTEBOOK_PASSWORD=$2
if [ "$1" == "" ]; then
echo "missing first argument"
usage
else
docker run -d -p 8888:8888 -p 5000:5000 -e NOTEBOOK_PASSWORD=$NOTEBOOK_PASSWORD deoldify run_$1 || docker build -t deoldify -f Dockerfile . && docker run -it -p 8888:8888 -p 5000:5000 -e NOTEBOOK_PASSWORD=$NOTEBOOK_PASSWORD deoldify run_$1
fi
run_image_api.sh
Executable
+2
View File
@@ -0,0 +1,2 @@
#!/bin/bash
python3 app.py
run_notebook.sh
+1
View File
@@ -1 +1,2 @@
#!/bin/bash
jupyter notebook --port=8888 --no-browser --allow-root --ip=0.0.0.0 --NotebookApp.token="" --NotebookApp.password="$(./set_password.py $NOTEBOOK_PASSWORD)"
run_video_api.sh
Executable
+2
View File
@@ -0,0 +1,2 @@
#!/bin/bash
python3 app-video.py