Update app.py

app.py

@@ -1,6 +1,6 @@
 import os
 from flask import Flask, request, jsonify
-from PIL import Image , ImageChops
+from PIL import Image
 from io import BytesIO
 import torch
 import base64
@@ -283,72 +283,36 @@ def tryon():
         'mask_image': mask_base64
     })
 
-def combine_masks(mask_upper, mask_lower):
-    if mask_upper.size != mask_lower.size:
-        mask_lower = mask_lower.resize(mask_upper.size)
-    
-    # Ensure masks are in 'L' mode
-    mask_upper = mask_upper.convert("L")
-    mask_lower = mask_lower.convert("L")
-
-    # Combine the two masks with logical OR
-    combined_mask = ImageChops.logical_or(mask_upper, mask_lower)
-    return combined_mask
-
-@spaces.GPU
+@spaces.GPU 
 def generate_mask(human_img, categorie='upper_body'):
-    device = "cuda"
-    openpose_model.preprocessor.body_estimation.model.to(device)
-    pipe.to(device)
-
     try:
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        
         # Redimensionner l'image pour le modèle
-        human_img_resized = human_img.convert("RGB").resize((384, 512))
-
-        if categorie == 'full_body':
-            # Générer les masques pour 'upper_body' et 'lower_body'
-            model_parse, _ = parsing_model(human_img_resized)
-            keypoints = openpose_model(human_img_resized)
-            
-            # Ajouter des logs pour vérifier la sortie
-            logging.info(f"model_parse type: {type(model_parse)}")
-            logging.info(f"keypoints type: {type(keypoints)}")
-            
-            # Vérifiez si `parsing_model` renvoie des données correctes
-            if not isinstance(model_parse, dict):
-                raise TypeError("parsing_model must return a dictionary.")
-            if not isinstance(keypoints, dict):
-                raise TypeError("openpose_model must return a dictionary.")
-            
-            mask_upper_body, _ = get_mask_location('hd', 'upper_body', model_parse, keypoints)
-            mask_lower_body, _ = get_mask_location('hd', 'lower_body', model_parse, keypoints)
-            
-            # Combiner les deux masques
-            combined_mask = combine_masks(mask_upper_body, mask_lower_body)
-        else:
-            # Générer le masque pour la catégorie spécifiée
-            keypoints = openpose_model(human_img_resized)
-            model_parse, _ = parsing_model(human_img_resized)
-
-            # Ajouter des logs pour vérifier la sortie
-            logging.info(f"model_parse type: {type(model_parse)}")
-            logging.info(f"keypoints type: {type(keypoints)}")
-            
-            # Vérifiez si `parsing_model` renvoie des données correctes
-            if not isinstance(model_parse, dict):
-                raise TypeError("parsing_model must return a dictionary.")
-            if not isinstance(keypoints, dict):
-                raise TypeError("openpose_model must return a dictionary.")
-            
-            combined_mask, _ = get_mask_location('hd', categorie, model_parse, keypoints)
-
+        human_img_resized = human_img.resize((384, 512))
+        
+        # Déplacer les modèles sur le bon appareil
+        openpose_model.to(device)
+        parsing_model.to(device)
+        
+        # Convertir l'image en tenseur et déplacer sur l'appareil
+        input_tensor = transforms.ToTensor()(human_img_resized).unsqueeze(0).to(device)
+        
+        # Génération des points clés et du masque
+        keypoints = openpose_model(input_tensor)
+        model_parse, _ = parsing_model(input_tensor)
+        mask, _ = get_mask_location('hd', categorie, model_parse, keypoints)
+        
+        # Déplacement du masque sur le CPU pour traitement avec PIL
+        mask = mask.cpu()
+        
         # Redimensionner le masque à la taille d'origine de l'image
-        mask_resized = combined_mask.resize(human_img.size)
+        mask_resized = transforms.ToPILImage()(mask.squeeze(0)).resize(human_img.size)
         
         return mask_resized
     except Exception as e:
         logging.error(f"Error generating mask: {e}")
-        raise e
+        raise e  # Renvoyer l'exception pour que l'API puisse la gérer
 
    
 @app.route('/generate_mask', methods=['POST'])
@@ -356,14 +320,11 @@ def generate_mask_api():
     try:
         # Récupérer les données de l'image à partir de la requête
         data = request.json
+        base64_image = data.get('human_image')
         categorie = data.get('categorie', 'upper_body')
         
-        valid_categories = ['upper_body', 'lower_body', 'dresses', 'full_body']
-        if categorie not in valid_categories:
-            raise ValueError(f"Invalid category '{categorie}'. Valid categories are: {', '.join(valid_categories)}")
-            
         # Décodage de l'image à partir de base64
-        human_img = decode_image_from_base64(data['human_image'])
+        human_img = decode_image_from_base64(base64_image).convert("RGB")
         
         # Appeler la fonction pour générer le masque
         mask_resized = generate_mask(human_img, categorie)
@@ -381,4 +342,3 @@ def generate_mask_api():
 
 if __name__ == "__main__":
     app.run(debug=True, host="0.0.0.0", port=7860)
-