Spaces:
Paused
Paused
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) | |
def index(): | |
return render_template('index.html') | |
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 | |
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 | |
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) |