Update main_code_script.py

main_code_script.py

@@ -1,7 +1,6 @@
 # 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
@@ -9,7 +8,6 @@ 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.
@@ -21,55 +19,44 @@ def estimate_pose(image_path):
     """
     mp_drawing = mp.solutions.drawing_utils
     mp_pose = mp.solutions.pose
-    try:
-        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)
-            if image is None:
-                raise FileNotFoundError(f"Could not open image: {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:
-                raise ValueError("No pose detected in the image.")
-    except Exception as e:
-        raise RuntimeError(f"Error processing image: {e}")
+    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
+    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"):
+def inpaint_clothing(image, mask_img, garment_image_path, device="cuda" if torch.cuda.is_available() else "cpu"): # Changed input
     """
-    Replaces the clothing region in the image with new clothing based on a text prompt,
+    Replaces the clothing region in the image with the uploaded garment image,
     using Stable Diffusion Inpainting.
     """
     pipe = StableDiffusionInpaintPipeline.from_pretrained(
@@ -77,45 +64,43 @@ def inpaint_clothing(image, mask_img, clothing_prompt, device="cuda" if torch.cu
         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
+    
+    # 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" 
     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):
+def change_clothing(image_path, garment_image_path): # Changed input
     """
     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)
-
+    modified_image = inpaint_clothing(image, mask_img, garment_image_path) # Changed input
     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)
-
+    garment_image_path = "garment.jpg" # Replace with your garment image
+    modified_image = change_clothing(input_image_path, garment_image_path)
     if modified_image:
         modified_image.save("modified_image.jpg")
         print("Clothing changed and saved to modified_image.jpg")