Update main_code_script.py

main_code_script.py

@@ -54,9 +54,9 @@ def segment_clothing(image, results): #Added result
     mask_img = Image.fromarray(binary_mask).convert("L")
     return mask_img
 # --- 3. Image Inpainting (Replacing Clothing - using Stable Diffusion Inpainting) ---
-def inpaint_clothing(image, mask_img, garment_image_path, device="cuda" if torch.cuda.is_available() else "cpu"): # Changed input
+def inpaint_clothing(image, mask_img, clothing_prompt, device="cuda" if torch.cuda.is_available() else "cpu"):
     """
-    Replaces the clothing region in the image with the uploaded garment image,
+    Replaces the clothing region in the image with new clothing based on a text prompt,
     using Stable Diffusion Inpainting.
     """
     pipe = StableDiffusionInpaintPipeline.from_pretrained(
@@ -64,23 +64,18 @@ def inpaint_clothing(image, mask_img, garment_image_path, device="cuda" if torch
         torch_dtype=torch.float16
     )
     pipe = pipe.to(device)
-    # Resize the image and mask to the same size.  Important for inpainting.
-    image = image.resize((512, 512))
-    mask_img = mask_img.resize((512, 512))
-    
-    # Load the garment image
-    garment_image = Image.open(garment_image_path).convert("RGB")
-    garment_image = garment_image.resize((512,512)) # Resize if necessary
-    
-    # Inpaint using the garment image as a guide (This part might need further refinement)
-    # A simple approach is to use the garment image in the prompt.  
-    # More advanced techniques might involve using the garment image as 
-    # a style reference or directly manipulating the latent space.
-    prompt = f"A photo of a person wearing the uploaded garment" 
+    # Resize the image and mask to a smaller size for faster inpainting
+    width, height = image.size
+    inpainted_size = (256, 256)  # Smaller size for faster inpainting
+    image = image.resize(inpainted_size)
+    mask_img = mask_img.resize(inpainted_size)
+    prompt = f"A photo of a person wearing {clothing_prompt}" #Add style or detail
     image = pipe(prompt=prompt, image=image, mask_image=mask_img).images[0]
+    # Resize back to the original size (or a desired output size)
+    image = image.resize((width, height)) # Or resize to a target output size
     return image
 # --- 4. Main Function (Putting it all together) ---
-def change_clothing(image_path, garment_image_path): # Changed input
+def change_clothing(image_path, clothing_prompt):
     """
     Main function to change the clothing in an image.
     """
@@ -94,13 +89,13 @@ def change_clothing(image_path, garment_image_path): # Changed input
     # 3. Segment the clothing
     mask_img = segment_clothing(image, results)
     # 4. Inpaint the clothing
-    modified_image = inpaint_clothing(image, mask_img, garment_image_path) # Changed input
+    modified_image = inpaint_clothing(image, mask_img, clothing_prompt)
     return modified_image
 # --- Example Usage ---
 if __name__ == "__main__":
     input_image_path = "person.jpg"  # Replace with your image
-    garment_image_path = "garment.jpg" # Replace with your garment image
-    modified_image = change_clothing(input_image_path, garment_image_path)
+    clothing_description = "a red leather jacket"  # Replace with desired clothing
+    modified_image = change_clothing(input_image_path, clothing_description)
     if modified_image:
         modified_image.save("modified_image.jpg")
         print("Clothing changed and saved to modified_image.jpg")