Update SegCloth.py

SegCloth.py

@@ -1,52 +1,42 @@
-from transformers import pipeline
 from PIL import Image
 import numpy as np
 from io import BytesIO
+import io
 import base64
-import torch
-from torchvision.transforms.functional import to_pil_image
 
-# Initialize segmentation and super-resolution pipelines
+
+
+# Initialize segmentation pipeline
 segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes")
-super_res = pipeline("image-super-resolution", model="CompVis/stable-diffusion-v-1-4")
+
 
 def encode_image_to_base64(image):
     buffered = BytesIO()
     image.save(buffered, format="PNG")
     return base64.b64encode(buffered.getvalue()).decode('utf-8')
 
-def segment_and_enhance_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"]):
+def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "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)
 
-            # Enhance image using super-resolution
-            enhanced_img = super_res(resized_img)
-            
+
+
+
             # Convert the final image to base64
-            imageBase64 = encode_image_to_base64(enhanced_img)
+            imageBase64 = encode_image_to_base64(resized_img)
             result_images.append((s['label'], imageBase64))
 
-    return result_images
+    return result_images  
+ 
 
-# Example usage
-# img = Image.open('your_image_path_here.png')
-# result_images = segment_and_enhance_clothing(img)
-# for clothing_type, image_base64 in result_images:
-#     print(clothing_type, image_base64)
+