Update SegCloth.py

SegCloth.py

@@ -1,26 +1,21 @@
-import torch
 from transformers import pipeline
-from diffusers import StableDiffusionPipeline
 from PIL import Image
 import numpy as np
 from io import BytesIO
 import base64
+import torch
+from torchvision.transforms.functional import to_pil_image
 
-# Initialize segmentation pipeline
+# Initialize segmentation and super-resolution pipelines
 segmenter = pipeline(model="mattmdjaga/segformer_b2_clothes")
-
-# Initialize Stable Diffusion pipeline
-model_id = "CompVis/stable-diffusion-v1-4"
-device = "cuda" if torch.cuda.is_available() else "cpu"
-pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16 if device == "cuda" else torch.float32)
-pipe = pipe.to(device)
+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_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"]):
+def segment_and_enhance_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dress", "Belt", "Left-shoe", "Right-shoe", "Scarf"]):
     # Segment image
     segments = segmenter(img)
 
@@ -41,12 +36,9 @@ def segment_clothing(img, clothes=["Hat", "Upper-clothes", "Skirt", "Pants", "Dr
             final_mask = Image.fromarray(current_mask)
             resized_img.putalpha(final_mask)
 
-            # Generate prompt for Stable Diffusion
-            prompt = f"a clear photo of {s['label']}"
-
-            # Enhance image using Stable Diffusion
-            enhanced_img = pipe(prompt).images[0]
-
+            # Enhance image using super-resolution
+            enhanced_img = super_res(resized_img)
+            
             # Convert the final image to base64
             imageBase64 = encode_image_to_base64(enhanced_img)
             result_images.append((s['label'], imageBase64))