Spaces:
Running
Running
File size: 2,325 Bytes
464bbb5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 |
import sys
sys.path.append("..") #To get access to the zombie module
sys.path.append("../pix2pixHD") #Give the zombie module access to pix2pixHD
sys.path.append("../../face-parsing.PyTorch")
from io import BytesIO
from flask import Flask
from flask import render_template,request,send_file
import zombie
import uuid
from PIL import Image
import onnxruntime
from flask_cors import CORS
import base64
import gc
import inferface
ort_session = onnxruntime.InferenceSession("../models/witch.onnx")
facenet = inferface.get_net("../../face-parsing.PyTorch/res/cp/79999_iter.pth").eval().cpu()
app = Flask(__name__)
CORS(app)
@app.route('/transform', methods=['POST'])
def serve_zombie():
gc.collect()
uploaded_file = request.files['image']
assert(uploaded_file.filename != '')
#suffix = uploaded_file.filename.split(".")[-1]
img_input_io = BytesIO()
uploaded_file.save(img_input_io)
human_image = zombie.ImageOpen(img_input_io).convert("RGB")
print(f"Input size: {human_image.size}")
#zombie_image = zombie.human_2_zombie(human_image,model,output_size=1024)
zombie_image = zombie.transition_onnx(human_image,ort_session)
if zombie_image is None:
return "No face found"
img_output_io = BytesIO()
zombie_image.save(img_output_io, 'JPEG', quality=75)
img_output_io.seek(0)
#return send_file(img_output_io, mimetype='image/jpeg')
return base64.b64encode(img_output_io.read())
@app.route('/transform_in_place', methods=['POST'])
def transform_in_place():
gc.collect()
uploaded_file = request.files['image']
assert(uploaded_file.filename != '')
img_input_io = BytesIO()
uploaded_file.save(img_input_io)
human_image = zombie.ImageOpen(img_input_io).convert("RGB")
print(f"Input size: {human_image.size}")
zombie_image = zombie.make_faces_zombie_from_array(human_image,facenet,ort_session)
img_output_io = BytesIO()
zombie_image.save(img_output_io, 'JPEG', quality=75)
img_output_io.seek(0)
return base64.b64encode(img_output_io.read())
@app.errorhandler(404)
def page_not_found(e):
print(request.url)
# note that we set the 404 status explicitly
return "endpoint not found", 404
if __name__ == '__main__':
app.run(host='0.0.0.0', debug=True, threaded=True) |