File size: 2,615 Bytes
ec2c4c5 |
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 75 76 77 78 79 80 81 82 83 84 85 |
from flask import Flask, jsonify, send_file, request
import base64
from PIL import Image, ImageOps
import io
import hydra
from omegaconf import DictConfig
from lama_predict import main as lama_predict
import os
import yaml
from omegaconf import OmegaConf
cwd = os.getcwd()
print(cwd)
config_path = os.path.join(cwd, "configs/prediction/default.yaml")
with open(config_path, 'r') as f:
config = OmegaConf.create(yaml.safe_load(f))
config.model.path = os.path.join(cwd, "big-lama")
config.indir = os.path.join(cwd, "web_server_input")
config.outdir = os.path.join(cwd, "web_server_output")
config.refine = False
app = Flask(__name__)
@app.route("/api/v2/image", methods=["GET", "POST"])
def echo_image():
# Get the image data from the request body
json_dict = request.get_json()
print(type(json_dict))
# Get the value of the "image" key, which is the base64 encoded image data
base64_image_data = json_dict["image"]
#print(base64_image_data[0:500])
image_bytes = base64.b64decode(base64_image_data)
image_stream = io.BytesIO(image_bytes)
image = Image.open(image_stream)
print(image.format_description)
if not os.path.exists("web_server_input"):
os.makedirs("web_server_input")
image.save("web_server_input/server.png")
base64_mask_data = json_dict["mask"]
image_bytes = base64.b64decode(base64_mask_data)
image_stream = io.BytesIO(image_bytes)
mask = Image.open(image_stream)
print(mask.format_description)
print(mask.format)
print(mask.size)
print(mask.mode)
if (mask.mode != "L"):
mask = mask.convert("L")
if not os.path.exists("web_server_input"):
os.makedirs("web_server_input")
mask.save("web_server_input/server_mask.png")
# Apply the mask to the image
# Create a new transparent image with the same size and mode as the image
transparent = Image.new(image.mode, image.size, (0, 0, 0, 0))
# Composite the image and the transparent image using the mask
masked_image = Image.composite(image, transparent, mask)
masked_image.save("server_masked_image.png")
# Convert the masked image to bytes and create a new stream
masked_image_stream = io.BytesIO()
masked_image.save(masked_image_stream, format='PNG')
masked_image_stream.seek(0)
lama_predict(config)
with open("web_server_output/server_mask.png", "rb") as image_file:
image_bytes = image_file.read()
image_inpainted_stream = io.BytesIO(image_bytes)
print(image.format_description)
image_inpainted_stream.seek(0)
return send_file(image_inpainted_stream, mimetype="image/png")
if __name__ == "__main__":
app.run(debug=True, port=9171)
|