server.py

from flask import Flask, request, send_file
from PIL import Image
import io

from src.data.datasets import CityscapesDataset

app = Flask(__name__)


@app.route('/process-image', methods=['POST'])
def process_image():
    if 'semantic_mask' not in request.files:
        return 'Missing image semantic mask!', 400
    #
    # if 'regeneration_mask' not in request.files:
    #     return 'Missing image regeneration mask!', 400

    if 'image' not in request.files:
        print("No source image found! Generating new image based on the regeneration and sematic masks!")
    else:
        image = Image.open(request.files['image'].stream)

    semantic_mask = Image.open(request.files['semantic_mask'].stream)
    # regeneration_mask = Image.open(request.files['regeneration_mask'].stream)

    print("Semantic mask shape: ", semantic_mask.size)
    print("Semantic mask mode: ", semantic_mask.mode)

    if semantic_mask.mode != 'RGB':
        semantic_mask = semantic_mask.convert('RGB')

    print("Semantic mask mode: ", semantic_mask.mode)

    output = CityscapesDataset.color_map_to_label_map(semantic_mask)
    reconstructed = CityscapesDataset.label_map_to_color_map(output)

    img_byte_arr = io.BytesIO()
    reconstructed.save(img_byte_arr, format='PNG')
    img_byte_arr = img_byte_arr.getvalue()

    return send_file(
        io.BytesIO(img_byte_arr),
        mimetype='image/png'
    )


if __name__ == '__main__':
    app.run(debug=True, host='0.0.0.0', port=5000)