app.py

import os
import numpy as np
import pandas as pd
import gradio as gr
from PIL import Image
import cv2
import mediapipe as mp

# Initialize Mediapipe Pose
mp_pose = mp.solutions.pose
pose = mp_pose.Pose(static_image_mode=True, model_complexity=2)

# Load sample product data
product_data = pd.DataFrame({
    'Product': ['Dress 1', 'Dress 2', 'Dress 3'],
    'Size': ['S', 'M', 'L'],
    'Color': ['Red', 'Blue', 'Green'],
    'Image': ['sample_dress1.jpg', 'sample_dress2.jpg', 'sample_dress3.jpg']
})

def get_keypoints(image):
    results = pose.process(cv2.cvtColor(image, cv2.COLOR_BGR2RGB))
    if not results.pose_landmarks:
        return None
    keypoints = {}
    for idx, lm in enumerate(results.pose_landmarks.landmark):
        keypoints[idx] = (int(lm.x * image.shape[1]), int(lm.y * image.shape[0]))
    return keypoints

def process_image(image, product):
    # Convert the uploaded image to a PIL image
    customer_image = Image.fromarray(image)
    customer_image = customer_image.convert("RGBA")

    # Convert to OpenCV format
    customer_image_cv = cv2.cvtColor(np.array(customer_image), cv2.COLOR_RGBA2BGR)

    # Get keypoints
    keypoints = get_keypoints(customer_image_cv)
    
    # Fetch the garment image corresponding to the selected product
    garment_filename = product_data[product_data['Product'] == product]['Image'].values[0]
    garment_path = os.path.join(os.getcwd(), garment_filename)

    dress_image = Image.open(garment_path).convert("RGBA")

    # Resize dress image based on keypoints (for simplicity, just resizing to fit)
    if keypoints and 11 in keypoints and 12 in keypoints:  # Left and right hips
        shoulder_width = abs(keypoints[11][0] - keypoints[12][0])
        scaling_factor = shoulder_width / dress_image.width
        new_dress_height = int(dress_image.height * scaling_factor)
        resized_dress_image = dress_image.resize((shoulder_width, new_dress_height), resample=Image.BICUBIC)
        
        # Calculate position
        dress_position = (keypoints[11][0] - shoulder_width // 2, keypoints[11][1])
        
        # Create result image
        result_image = Image.new("RGBA", customer_image.size)
        result_image.paste(customer_image, (0, 0))
        result_image.paste(resized_dress_image, dress_position, resized_dress_image)
        
        # Convert to numpy array
        result_array = np.array(result_image)
    else:
        result_array = np.array(customer_image)
    
    # Fetch product details
    product_details = product_data[product_data['Product'] == product].iloc[0].to_dict()
    
    return result_array, product_details

# Gradio interface
iface = gr.Interface(
    fn=process_image,
    inputs=[
        gr.Image(type="numpy", label="Upload Your Image"),
        gr.Dropdown(choices=product_data['Product'].tolist(), label="Select Product")
    ],
    outputs=[
        gr.Image(type="numpy", label="Output Image"),
        gr.JSON(label="Product Details")
    ],
    title="Virtual Dress Fitting",
    description="Upload an image and select a product to see how it fits on you."
)

if __name__ == "__main__":
    iface.launch()