Update app.py

app.py

@@ -1,43 +1,47 @@
-from flask import Flask, request, jsonify
+from transformers import pipeline
 from PIL import Image
-import base64
+import numpy as np
 from io import BytesIO
-import io 
-from SegCloth import segment_clothing
-
-
-
-# Créer une instance FastAPI
-app = Flask(__name__)
-
-def segment_image(img, clothes):
-    img = decode_image_from_base64(img)
-    return segment_clothing(img, clothes)
+import base64
 
-# Fonction pour décoder une image encodée en base64 en objet PIL.Image.Image
-def decode_image_from_base64(image_data):
-    image_data = base64.b64decode(image_data)
-    image = Image.open(io.BytesIO(image_data))
-    return image
+# Initialize segmentation pipeline
+segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes")
 
 def encode_image_to_base64(image):
     buffered = BytesIO()
     image.save(buffered, format="PNG")
     return base64.b64encode(buffered.getvalue()).decode('utf-8')
 
-@app.get("/")
-def root():
-    return "Welcome to  Fashion Clothing Detectiony API!"
-
-# Route pour l'API REST
-@app.route('/api/classify', methods=['POST'])
-def classify():
-    data = request.json
-    print(data)
-    clothes = ["Hat", "Upper-clothes", "Skirt", "Pants", "Dress"]
-    image = data['image']
-    result = segment_image(image,clothes)
-    return jsonify({'result': result})
-
-if __name__ == "__main__":
-    app.run(debug=True, host="0.0.0.0", port=7860)
+def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"]):
+    # Segment image
+    segments = segmenter(img)
+
+    # Convert image to RGBA
+    img = img.convert("RGBA")
+
+    # Create list of masks
+    result_images = []
+    for s in segments:
+        if s['label'] in clothes:
+            # Extract mask and resize image to mask size
+            current_mask = np.array(s['mask'])
+            mask_size = current_mask.shape[::-1]  # Mask size is (width, height)
+            
+            # Resize the original image to match the mask size
+            resized_img = img.resize(mask_size)
+
+            # Apply mask to resized image
+            final_mask = Image.fromarray(current_mask)
+            resized_img.putalpha(final_mask)
+
+            # Convert the final image to base64
+            imageBase64 = encode_image_to_base64(resized_img)
+            result_images.append((s['label'], imageBase64))
+
+    return result_images
+
+# Example usage
+# img = Image.open('your_image_path_here.png')
+# result_images = segment_clothing(img)
+# for clothing_type, image_base64 in result_images:
+#     print(clothing_type, image_base64)