api/app.py

from flask import Flask, render_template, send_from_directory,  send_file, request, jsonify
from flask_cors import CORS
import subprocess
import os
import io
from model.src.utils.arg_parser import eval_parse_args  # Nuovo import corretto
import sys
from PIL import Image
import base64
import traceback


from model.src import eval

app = Flask(__name__)
CORS(app)

@app.route('/')
def index():
    return render_template('index.html')

@app.route('/generate-design', methods=['GET','POST'])
def generate_design():

    #image_save_dir = "/api/model/assets/data/vitonhd/test/im_sketch"
    text_file_path = "/api/model/assets/data/vitonhd/test/test_paired.txt"

    try:
        print("### Inizio generate_design ###")  # Stampa per il debug
        
        # Getting json
        json_data_from_req = request.get_json()
        if not json_data_from_req:
            print("### JSON non ricevuto ###")  # Stampa per il debug
            return "Invalid or missing JSON data", 400

        # Getting Image
        if 'image' not in json_data_from_req:
            print("### Nessuna immagine ricevuta ###")  # Stampa per il debug
            return "No image file in request", 400

        print(json_data_from_req)

        encoded_image = json_data_from_req['image']

        try:
            image_data = base64.b64decode(encoded_image)
            image = Image.open(io.BytesIO(image_data))
            print("### Immagine aperta con successo ###")  # Stampa per il debug
        except Exception as e:
            print(f"### Errore nell'apertura dell'immagine: {str(e)} ###")  # Stampa per il debug
            return f"Failed to open the image: {str(e)}", 400

        if 'vitonhd' in json_data_from_req:
            dataset_name = 'vitonhd'
            image_suffix = '_00.jpeg'
            dataset_path = '/api/model/assets/data/vitonhd'
            image_save_dir = os.path.join(dataset_path, 'test/im_sketch')
        elif 'dresscode' in json_data_from_req:
            dataset_name = 'dresscode'
            image_suffix = '_1.jpeg'
            body_part = json_data_from_req['body_part']
            dataset_path = '/api/model/assets/data/dresscode'
            sketch_path = f'/api/model/assets/data/dresscode/{body_part}'
            image_save_dir = os.path.join(sketch_path, 'im_sketch')
        else:
            print("### Chiave 'vitonhd' o 'dresscode' non trovata nel JSON ###")  # Stampa per il debug
            return "Invalid JSON data, expected 'vitonhd' or 'dresscode'", 400

        
        # Saving sketch in filesystem
        model_key = list(json_data_from_req[dataset_name].keys())[0]  # Es: "03191"
        image_filename = f"{model_key}{image_suffix}"  # Nome file immagine es: "03191_00.jpg"
        print("###################################### NOME SKETCH: ", image_filename)
        print("###################################### DIRECTORY DOVE LO SALVO: ", image_save_dir)
        image_save_path = os.path.join(image_save_dir, image_filename)
        image.save(image_save_path, format='JPEG')

        # Elenca tutti i file nella directory
        files_in_directory = os.listdir(image_save_dir)

        # Stampa i file trovati
        for file_name in files_in_directory:
            print(file_name)
        
        #with open(text_file_path, "w") as text_file:
        #    text_file.write(f"{image_filename} {image_filename}\n")

        
        # Argomenti per eval.main
        sys.argv = [
            'eval.py',
            '--dataset_path', dataset_path,
            '--batch_size', '1',
            '--mixed_precision', 'fp16',
            '--num_workers_test', '4',
            '--sketch_cond_rate', '0.2',
            '--dataset', dataset_name,
            '--start_cond_rate', '0.0',
            '--test_order', 'paired'
        ]
        print("### Esecuzione eval.main ###")  # Stampa per il debug
        
        import traceback

        # Esecuzione del modello
        try:
            final_image = eval.main(json_data_from_req)
            print("### eval.main eseguito con successo ###")  # Stampa per il debug
        except AttributeError as e:
            print(f"### AttributeError: {str(e)} ###")  # Stampa per il debug
            return f"AttributeError: {traceback.format_exc()}", 500
        except TypeError as e:
            print(f"### TypeError: {str(e)} ###")  # Stampa per il debug
            return f"TypeError: {traceback.format_exc()}", 500
        except ValueError as e:
            print(f"### ValueError: {str(e)} ###")  # Stampa per il debug
            return f"ValueError: {traceback.format_exc()}", 500
        except IOError as e:
            print(f"### IOError: {str(e)} ###")  # Stampa per il debug
            return f"IOError: {traceback.format_exc()}", 500
        except Exception as e:
            print(f"### Unexpected error: {str(e)} ###")  # Stampa per il debug
            return f"Unexpected error: {traceback.format_exc()}", 500


        # Salvataggio immagine e invio come risposta
        img_io = io.BytesIO()
        final_image.save(img_io, 'JPEG')
        img_io.seek(0)
        print("### Immagine finale pronta per il download ###")  # Stampa per il debug

        return send_file(img_io, mimetype='image/jpeg', as_attachment=True, download_name='inmemory_image.jpg')

    except Exception as e:
        print(f"### Errore globale: {str(e)} ###")  # Stampa per il debug
        return str(e), 500
    
@app.route('/test_post', methods=['POST'])
def test_post():
    # Verifica se la richiesta è POST
    if request.method == 'POST':
        # Ottieni i dati dal corpo della richiesta
        data = request.get_json()
        
        # Se i dati non sono presenti o il formato non è corretto
        if not data:
            return jsonify({'message': 'Nessun dato ricevuto o dati non validi'}), 400
        
        # Restituisci una risposta con i dati ricevuti
        return jsonify({'message': 'POST funzionante', 'data_received': data}), 200

@app.route('/process-image', methods=['POST'])
def process_image():
    try:
        # Ottieni il JSON dalla richiesta
        json_data = request.get_json()
        if not json_data or 'image' not in json_data:
            return "Invalid or missing JSON data", 400

        # Decodifica l'immagine base64
        encoded_image = json_data['image']
        try:
            image_data = base64.b64decode(encoded_image)
            image = Image.open(io.BytesIO(image_data))
            print("### Immagine aperta con successo ###")  # Stampa per il debug
        except Exception as e:
            print(f"### Errore nell'apertura dell'immagine: {str(e)} ###")  # Stampa per il debug
            return f"Failed to open the image: {str(e)}", 400

        # Elaborazione dell'immagine con PIL (ad esempio, convertiamo l'immagine in scala di grigi)
        processed_image = image.convert("L")  # Converti l'immagine in scala di grigi
        
        # Salva l'immagine su un buffer
        img_io = io.BytesIO()
        processed_image.save(img_io, 'JPEG')
        img_io.seek(0)

        # Restituisci l'immagine come file scaricabile
        return send_file(img_io, mimetype='image/jpeg', as_attachment=True, download_name='processed_image.jpg')

    except Exception as e:
        print(f"### Errore durante l'elaborazione dell'immagine: {str(e)} ###")  # Stampa per il debug
        return str(e), 500


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