from flask import Flask, request, jsonify,send_file, Response from flask_cors import CORS import logging import gc import os from threading import Thread from flask_sse import sse from io import BytesIO from pathlib import Path import sys import torch from PIL import Image, ImageOps import numpy as np from run.utils_ootd import get_mask_location from run.cloths_db import cloths_map, modeL_db from preprocess.openpose.run_openpose import OpenPose from preprocess.humanparsing.run_parsing import Parsing from ootd.inference_ootd_dc import OOTDiffusionDC PROJECT_ROOT = Path(__file__).absolute().parents[1].absolute() sys.path.insert(0, str(PROJECT_ROOT)) from queue import Queue #run python garbage collector and nvidia cuda clear memory gc.collect() torch.cuda.empty_cache() # Setup Flask server app = Flask(__name__) CORS(app, origins="*") # Enable CORS for the entire app app.config["REDIS_URL"] = "redis://localhost:6379" app.register_blueprint(sse, url_prefix='/stream') logger = logging.getLogger() openpose_model = OpenPose(0) parsing_model_dc = Parsing(0) ootd_model_dc = OOTDiffusionDC(0) example_path = os.path.join(os.path.dirname(__file__), 'examples') garment_path = os.path.join(os.path.dirname(__file__), 'examples','garment') openpose_model.preprocessor.body_estimation.model.to('cuda') ootd_model_dc.pipe.to('cuda') ootd_model_dc.image_encoder.to('cuda') ootd_model_dc.text_encoder.to('cuda') category_dict = ['upperbody', 'lowerbody', 'dress'] category_dict_utils = ['upper_body', 'lower_body', 'dresses'] # Ensure this directory exists UPLOAD_FOLDER = 'temp_images' if not os.path.exists(UPLOAD_FOLDER): os.makedirs(UPLOAD_FOLDER) # progress_queue = Queue() # def progress_callback(step, total_steps): # if total_steps is not None and total_steps > 0: # progress = int((step + 1) / total_steps * 100) # progress_queue.put(progress) # else: # progress_queue.put(step + 1) def progress_callback(step, total_steps): if total_steps is not None and total_steps > 0: progress = int((step + 1) / total_steps * 100) sse.publish({"progress": progress}, type='progress') else: sse.publish({"step": step + 1}, type='progress') def process_dc(vton_img, garm_img, category): model_type = 'dc' if category == 'Upper-body': category = 0 elif category == 'Lower-body': category = 1 else: category = 2 with torch.no_grad(): # openpose_model.preprocessor.body_estimation.model.to('cuda') # ootd_model_dc.pipe.to('cuda') # ootd_model_dc.image_encoder.to('cuda') # ootd_model_dc.text_encoder.to('cuda') garm_img = Image.open(garm_img).resize((768, 1024)) vton_img = Image.open(vton_img).resize((768, 1024)) keypoints = openpose_model(vton_img.resize((384, 512))) print(len(keypoints["pose_keypoints_2d"])) print(keypoints["pose_keypoints_2d"]) left_point = keypoints["pose_keypoints_2d"][2] right_point = keypoints["pose_keypoints_2d"][5] neck_point = keypoints["pose_keypoints_2d"][1] hip_point = keypoints["pose_keypoints_2d"][8] print(f'left shoulder - {left_point}') print(f'right shoulder - {right_point}') # #find disctance using Euclidian distance shoulder_width_pixels = round(np.sqrt( np.power((right_point[0]-left_point[0]),2) + np.power((right_point[1]-left_point[1]),2)),2) height_pixels = round(np.sqrt( np.power((neck_point[0]-hip_point[0]),2) + np.power((neck_point[1]-hip_point[1]),2)),2) *2 # # Assuming an average human height average_height_cm = 172.72 *1.5 # Conversion factor from pixels to cm conversion_factor = average_height_cm / height_pixels # Convert shoulder width to real-world units shoulder_width_cm = shoulder_width_pixels * conversion_factor print(f'Shoulder width (in pixels): {shoulder_width_pixels}') print(f'Estimated height (in pixels): {height_pixels}') print(f'Conversion factor (pixels to cm): {conversion_factor}') print(f'Shoulder width (in cm): {shoulder_width_cm}') print(f'Shoulder width (in INCH): {round(shoulder_width_cm/2.54,1)}') model_parse,_ = parsing_model_dc(vton_img.resize((384, 512))) mask, mask_gray = get_mask_location(model_type, category_dict_utils[category], model_parse, keypoints) mask = mask.resize((768, 1024), Image.NEAREST) mask_gray = mask_gray.resize((768, 1024), Image.NEAREST) # Save the resized masks # mask.save("mask_resized.png") # mask_gray.save("mask_gray_resized.png") masked_vton_img = Image.composite(mask_gray, vton_img, mask) # masked_vton_img.save("masked_vton_img.png") print(f'category is {category}') # images = ootd_model_dc( # model_type=model_type, # category=category_dict[category], # image_garm=garm_img, # image_vton=masked_vton_img, # mask=mask, # image_ori=vton_img, # num_samples=3, # num_steps=20, # image_scale= 2.0, # seed=-1, # ) images = ootd_model_dc( model_type=model_type, category=category_dict[category], image_garm=garm_img, image_vton=masked_vton_img, mask=mask, image_ori=vton_img, num_samples=2, num_steps=10, image_scale=2.0, seed=42, progress_callback=progress_callback, progress_interval=1, # Update progress every step ) return images @app.route('/') def root(): try: response_data = {"message": "This is VTR API v1.0"} return jsonify(response_data) except Exception as e: logger.error(f"Root endpoint error: {str(e)}") response_data = {"message": "Internal server Error"} return jsonify(response_data), 500 @app.route('/stream') def stream(): return Response(sse.stream(), content_type='text/event-stream') #write Flask api name "generate" with POST method that will input 2 images and return 1 image @app.route('/generate', methods=['POST']) def generate(): """ A Flask route that handles a POST request to the '/generate' endpoint. It expects two files, 'garm_img' and 'vton_img', to be included in the request. The function calls the 'process_dc' function with the provided files and the category 'Upper-body'. It then sends the processed image as a file with the mimetype 'image/png' and returns it to the client. If any exception occurs, the function logs the error and returns a JSON response with a status code of 500. Parameters: None Returns: A Flask response object with the processed image as a file. Raises: None """ # if category == 'Upper-body': # category = 0 # elif category == 'Lower-body': # category = 1 # else: # category = 2 try: cloths_type = ["Upper-body", "Lower-body", "Dress"] garm_img = request.files['garm_img'] vton_img = request.files['vton_img'] cat = request.form['category'] print(f'category is {cat}') category =cloths_type[int(cat)] # Default to Upper-body if not specified # Save the uploaded files garm_path = os.path.join(UPLOAD_FOLDER, 'garm_input.png') vton_path = os.path.join(UPLOAD_FOLDER, 'vton_input.png') garm_img.save(garm_path) vton_img.save(vton_path) # Convert file objects to bytes IO objects # garm_img = BytesIO(garm_img.read()) # vton_img = BytesIO(vton_img.read()) output_images = process_dc(garm_img=garm_img, vton_img=vton_img, category=category) if not output_images: return Response("No output image generated", status=500) output_image = output_images[0] # Get the first image # Convert PIL Image to bytes img_byte_arr = BytesIO() output_image.save(img_byte_arr, format='PNG') img_byte_arr = img_byte_arr.getvalue() # Send the final "complete" event via SSE sse.publish({"message": "Processing complete"}, type='complete') return Response(img_byte_arr, mimetype='image/png') except Exception as e: print(f"Error: {str(e)}") # Log the error return Response(str(e), status=500) if __name__ == '__main__': app.run(debug=False, host='0.0.0.0', port=5009) # nohup gunicorn -b 0.0.0.0:5003 sentiment_api:app &