from flask import Flask, request, send_file, jsonify
import cv2, io, random, os
import numpy as np

app = Flask(__name__)

model = cv2.dnn_superres.DnnSuperResImpl_create()
model.readModel('LapSRN_x8.pb')
model.setModel('lapsrn', 8)

@app.route("/upscale/", methods=["POST"])
def upscale_image():
    if 'file' not in request.files:
        return jsonify({"error": "No file part"}), 400

    file = request.files['file']
    if file.filename == '':
        return jsonify({"error": "No selected file"}), 400

    file_bytes = file.read()
    nparr = np.frombuffer(file_bytes, np.uint8)
    img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
    if img is None:
        return jsonify({"error": "Invalid image"}), 400

    result = model.upsample(img)

    _, encoded_img = cv2.imencode('.png', result)
    return send_file(
        io.BytesIO(encoded_img.tobytes()),
        mimetype='image/png',
        as_attachment=False,
        download_name=f'{random.randint(1, 10000000000000000000000000000000)}.png'
    )

if __name__ == "__main__":
    port = int(os.environ.get("PORT", 7860))
    app.run(host="0.0.0.0", port=port)