import streamlit as st
import cv2
import numpy as np
import mediapipe as mp
from PIL import Image

st.title("👕 Virtual Try-On with Pose Estimation")

person_file = st.file_uploader("Upload Person Image", type=["jpg", "jpeg", "png"])
cloth_file = st.file_uploader("Upload Clothing Image (PNG)", type=["png"])

mp_pose = mp.solutions.pose
pose = mp_pose.Pose(static_image_mode=True)

def overlay_cloth(person_img, cloth_img):
    img_rgb = cv2.cvtColor(person_img, cv2.COLOR_BGR2RGB)
    result = pose.process(img_rgb)

    if not result.pose_landmarks:
        st.warning("Pose not detected.")
        return person_img

    h, w, _ = person_img.shape
    landmarks = result.pose_landmarks.landmark

    left_shoulder = landmarks[mp_pose.PoseLandmark.LEFT_SHOULDER]
    right_shoulder = landmarks[mp_pose.PoseLandmark.RIGHT_SHOULDER]
    left_hip = landmarks[mp_pose.PoseLandmark.LEFT_HIP]
    right_hip = landmarks[mp_pose.PoseLandmark.RIGHT_HIP]

    # Shoulder and hip coordinates
    x1, y1 = int(left_shoulder.x * w), int(left_shoulder.y * h)
    x2, y2 = int(right_shoulder.x * w), int(right_shoulder.y * h)
    x3, y3 = int(left_hip.x * w), int(left_hip.y * h)
    x4, y4 = int(right_hip.x * w), int(right_hip.y * h)

    # Width and height based on shoulder width and upper body height
    cloth_width = int(np.linalg.norm([x2 - x1, y2 - y1]) * 1.2)
    cloth_height = int(np.linalg.norm([((x3 + x4) // 2) - ((x1 + x2) // 2), ((y3 + y4) // 2) - ((y1 + y2) // 2)]) * 1.3)

    # Center the clothing at chest
    center_x = (x1 + x2) // 2
    center_y = (y1 + y2) // 2

    # Resize cloth
    cloth_img_resized = cloth_img.resize((cloth_width, cloth_height), Image.Resampling.LANCZOS)
    # Compute top-left corner for overlay
    paste_x = int(center_x - cloth_width // 2)
    paste_y = int(center_y - cloth_height // 3)

    # Convert person_img to PIL for blending
    person_pil = Image.fromarray(cv2.cvtColor(person_img, cv2.COLOR_BGR2RGB)).convert("RGBA")
    cloth_img_resized = cloth_img_resized.convert("RGBA")

    # Paste with transparency
    person_pil.paste(cloth_img_resized, (paste_x, paste_y), cloth_img_resized)

    return np.array(person_pil)

if person_file and cloth_file:
    person_img = np.array(Image.open(person_file).convert("RGB"))
    cloth_img = Image.open(cloth_file)

    result_img = overlay_cloth(person_img, cloth_img)
    st.image(result_img, caption="Result", use_column_width=True)