Upload code.py

code.py

@@ -0,0 +1,122 @@
+# Install necessary libraries (in your requirements.txt)
+# pillow opencv-python transformers mediapipe diffusers accelerate transformers
+# Example install command: pip install pillow opencv-python transformers mediapipe diffusers accelerate transformers
+
+from PIL import Image
+import cv2
+import mediapipe as mp
+import numpy as np
+from transformers import pipeline
+from diffusers import StableDiffusionInpaintPipeline
+import torch
+
+# --- 1. Pose Estimation (using Mediapipe) ---
+def estimate_pose(image_path):
+    """Detects the pose of a person in an image using Mediapipe.
+
+    Args:
+        image_path: Path to the input image.
+
+    Returns:
+        A list of landmarks (x, y, visibility) or None if no pose is detected.
+    """
+    mp_drawing = mp.solutions.drawing_utils
+    mp_pose = mp.solutions.pose
+
+    with mp_pose.Pose(
+        static_image_mode=True,
+        model_complexity=2,
+        enable_segmentation=True,
+        min_detection_confidence=0.5) as pose:
+
+        image = cv2.imread(image_path)
+        image_rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
+        results = pose.process(image_rgb)
+
+        if results.pose_landmarks:
+            # Example: Draw the pose landmarks on the image (for visualization)
+            annotated_image = image.copy()
+            mp_drawing.draw_landmarks(
+                annotated_image, results.pose_landmarks, mp_pose.POSE_CONNECTIONS)
+
+            #cv2.imwrite("pose_annotated.jpg", annotated_image)  # Save annotated image
+            #return results.pose_landmarks.landmark  # Return the landmarks
+            return results, image # Return the entire result
+
+        else:
+            return None, None # or raise an exception
+
+# --- 2. Clothing Segmentation (Example - using a placeholder function) ---
+def segment_clothing(image, results): #Added result
+    """Segments the clothing region in the image.
+       This is a simplified example. In reality, you would use a pre-trained
+       segmentation model.
+    """
+
+    # 1. Create a mask where the person is present.
+    segmentation_mask = results.segmentation_mask
+    threshold = 0.5  # Adjust this threshold as needed.
+
+    # Threshold the segmentation mask to create a binary mask.
+    binary_mask = (segmentation_mask > threshold).astype(np.uint8) * 255
+
+    # Convert binary mask to a PIL Image
+    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, clothing_prompt, device="cuda" if torch.cuda.is_available() else "cpu"):
+    """
+    Replaces the clothing region in the image with new clothing based on a text prompt,
+    using Stable Diffusion Inpainting.
+    """
+    pipe = StableDiffusionInpaintPipeline.from_pretrained(
+        "stabilityai/stable-diffusion-2-inpainting",
+        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))
+
+    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]
+    return image
+
+# --- 4. Main Function (Putting it all together) ---
+def change_clothing(image_path, clothing_prompt):
+    """
+    Main function to change the clothing in an image.
+    """
+    # 1. Load the image
+    image = Image.open(image_path).convert("RGB")
+
+    # 2. Estimate the pose
+    results, cv2_image = estimate_pose(image_path)
+    if results is None:
+        print("No pose detected.")
+        return None
+
+    # 3. Segment the clothing
+    mask_img = segment_clothing(image, results)
+
+    # 4. Inpaint the clothing
+    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
+    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")
+    else:
+        print("Failed to change clothing.")