Update app.py

app.py

@@ -122,6 +122,45 @@ def pil_to_binary_mask(pil_image, threshold=0):
     output_mask = Image.fromarray(mask)
     return output_mask
 
+@spaces.GPU
+def start_tryon_full_body(tops_image: Image.Image, bottoms_image: Image.Image, model_parse_tops: Image.Image, model_parse_bottoms: Image.Image, keypoints_tops: dict, keypoints_bottoms: dict):
+    """
+    Combines tops and bottoms images into a single output image after processing with get_mask_location.
+    """
+
+    # Get mask for the tops (upper body)
+    mask_tops, _ = get_mask_location('hd', "upper_body", model_parse_tops, keypoints_tops)
+
+    # Get mask for the bottoms (lower body)
+    mask_bottoms, _ = get_mask_location('hd', "lower_body", model_parse_bottoms, keypoints_bottoms)
+
+    # Convert masks to NumPy arrays
+    mask_tops_np = np.array(mask_tops)
+    mask_bottoms_np = np.array(mask_bottoms)
+
+    # Convert tops and bottoms images to NumPy arrays
+    tops_np = np.array(tops_image)
+    bottoms_np = np.array(bottoms_image)
+
+    # Ensure that tops and bottoms images have the same dimensions as their masks
+    tops_resized = cv2.resize(tops_np, (mask_tops_np.shape[1], mask_tops_np.shape[0]))
+    bottoms_resized = cv2.resize(bottoms_np, (mask_bottoms_np.shape[1], mask_bottoms_np.shape[0]))
+
+    # Create a blank canvas for the final output image
+    combined_image = np.zeros_like(tops_resized)
+
+    # Apply the tops mask to the combined image
+    combined_image[mask_tops_np > 0] = tops_resized[mask_tops_np > 0]
+
+    # Apply the bottoms mask to the combined image
+    combined_image[mask_bottoms_np > 0] = bottoms_resized[mask_bottoms_np > 0]
+
+    # Convert the final combined image back to a PIL image
+    combined_image_pil = Image.fromarray(combined_image)
+
+    return combined_image_pil
+
+
 @spaces.GPU
 def start_tryon(dict, garm_img, garment_des, is_checked, is_checked_crop, denoise_steps, seed):
     device = "cuda"
@@ -270,6 +309,37 @@ def tryon():
         'mask_image': mask_base64
     })
 
+@app.route('/tryon-full', methods=['POST'])
+def tryon_full():
+    data = request.json
+
+    # Decode input images
+    tops_image = decode_image_from_base64(data['tops_image'])
+    bottoms_image = decode_image_from_base64(data['bottoms_image'])
+    model_parse_tops = decode_image_from_base64(data['model_parse_tops'])
+    model_parse_bottoms = decode_image_from_base64(data['model_parse_bottoms'])
+
+    # Decode keypoints
+    keypoints_tops = data.get('keypoints_tops', {})
+    keypoints_bottoms = data.get('keypoints_bottoms', {})
+
+    # Call the start_tryon function
+    output_image = start_tryon_full_body(
+        tops_image,
+        bottoms_image,
+        model_parse_tops,
+        model_parse_bottoms,
+        keypoints_tops,
+        keypoints_bottoms
+    )
+
+    # Convert output image to base64
+    output_base64 = encode_image_to_base64(output_image)
+
+    return jsonify({
+        'output_image': output_base64
+    })
+
 if __name__ == "__main__":
     app.run(debug=True, host="0.0.0.0", port=7860)