mirror of
https://github.com/jantic/DeOldify.git
synced 2026-09-01 15:36:32 +08:00
Alexandre Vicenzi
547fb7e56f
It will be better to do:
from deoldify import visualize
Than:
from fasterai import visualize
The PyPI package will be called DeOldify, in this case
it makes more sense to have an import name that matches
the package name.
Also, fasterai resembles fastai library, and DeOldify
is not a library to use with fastai, it's built on
top of it.
98 lines
2.4 KiB
Python
98 lines
2.4 KiB
Python
# 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 deoldify.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_image():
|
|
|
|
input_path = generate_random_filename(upload_directory,"jpeg")
|
|
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"])
|
|
|
|
download(url, input_path)
|
|
|
|
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(output_path, mimetype='image/jpeg')
|
|
|
|
return callback, 200
|
|
|
|
except:
|
|
traceback.print_exc()
|
|
return {'message': 'input error'}, 400
|
|
|
|
finally:
|
|
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=False)
|