virtual_tryon_advanced.py

import streamlit as st
import google.generativeai as genai
from dotenv import load_dotenv
import os
from PIL import Image, ImageEnhance, ImageFilter
import io
import base64
import time
import cv2
import numpy as np
import requests
from typing import Tuple, Optional

# Load environment variables
load_dotenv()

# Configure page
st.set_page_config(
    page_title="Advanced Virtual Try-On App",
    page_icon="👗",
    layout="wide",
    initial_sidebar_state="collapsed"
)

def load_css():
    """Load custom CSS for liquid glass design"""
    st.markdown("""
    <style>
    /* Import Google Fonts */
    @import url('https://fonts.googleapis.com/css2?family=Inter:wght@300;400;500;600;700&display=swap');
    
    /* Global Styles */
    .stApp {
        background: linear-gradient(135deg, 
            rgba(74, 144, 226, 0.1) 0%, 
            rgba(255, 107, 107, 0.1) 35%, 
            rgba(142, 68, 173, 0.1) 100%);
        font-family: 'Inter', sans-serif;
    }
    
    /* Liquid Glass Card Effect */
    .glass-card {
        background: rgba(255, 255, 255, 0.1);
        backdrop-filter: blur(20px);
        -webkit-backdrop-filter: blur(20px);
        border-radius: 20px;
        border: 1px solid rgba(255, 255, 255, 0.2);
        box-shadow: 0 8px 32px 0 rgba(31, 38, 135, 0.37);
        padding: 2rem;
        margin: 1rem 0;
        transition: all 0.3s ease;
    }
    
    .glass-card:hover {
        transform: translateY(-5px);
        box-shadow: 0 15px 40px 0 rgba(31, 38, 135, 0.5);
    }
    
    /* Header Styles */
    .main-header {
        text-align: center;
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        -webkit-background-clip: text;
        -webkit-text-fill-color: transparent;
        background-clip: text;
        font-size: 3.5rem;
        font-weight: 700;
        margin: 2rem 0;
        text-shadow: 0 0 30px rgba(102, 126, 234, 0.5);
    }
    
    .sub-header {
        text-align: center;
        color: rgba(255, 255, 255, 0.9);
        font-size: 1.3rem;
        font-weight: 300;
        margin-bottom: 3rem;
    }
    
    /* Button Styles */
    .stButton > button {
        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
        color: white;
        border: none;
        border-radius: 50px;
        padding: 1rem 2.5rem;
        font-weight: 600;
        font-size: 1.1rem;
        transition: all 0.3s ease;
        box-shadow: 0 4px 15px 0 rgba(102, 126, 234, 0.4);
    }
    
    .stButton > button:hover {
        transform: translateY(-3px);
        box-shadow: 0 12px 30px 0 rgba(102, 126, 234, 0.7);
    }
    
    /* Result container */
    .result-container {
        background: rgba(255, 255, 255, 0.08);
        backdrop-filter: blur(15px);
        border-radius: 20px;
        padding: 2rem;
        border: 1px solid rgba(255, 255, 255, 0.1);
        box-shadow: 0 8px 32px 0 rgba(31, 38, 135, 0.37);
        margin: 2rem 0;
    }
    
    /* Hide Streamlit branding */
    #MainMenu {visibility: hidden;}
    footer {visibility: hidden;}
    header {visibility: hidden;}
    </style>
    """, unsafe_allow_html=True)

def initialize_gemini():
    """Initialize Gemini AI with API key"""
    api_key = os.getenv("GEMINI_API_KEY")
    if not api_key or api_key == "your_gemini_api_key_here":
        return None
    
    try:
        genai.configure(api_key=api_key)
        return genai.GenerativeModel('gemini-2.0-flash-exp')
    except Exception as e:
        st.error(f"Error initializing Gemini AI: {str(e)}")
        return None

def analyze_images_for_tryon(person_image, clothing_image, model):
    """Analyze images using Gemini for virtual try-on guidance"""
    try:
        prompt = """
        You are an expert AI fashion designer and virtual try-on specialist. Analyze these two images:

        Image 1: Person's photo
        Image 2: Clothing item

        Please provide a detailed technical analysis for virtual try-on:

        1. **Person Analysis:**
           - Body type and proportions
           - Pose and angle
           - Lighting conditions and direction
           - Skin tone and hair color
           - Current clothing that needs to be replaced

        2. **Clothing Analysis:**
           - Type of garment (dress, shirt, pants, etc.)
           - Style and cut details
           - Color and pattern
           - Fabric type and texture
           - How it should fit on different body types

        3. **Virtual Try-On Instructions:**
           - Specific placement points on the body
           - How to adjust sizing for this person
           - Lighting adjustments needed
           - Shadow and highlight placement
           - Areas that need special attention

        4. **Technical Implementation:**
           - Key anchor points for placement
           - Perspective adjustments needed  
           - Color correction requirements
           - Blending techniques to use

        Provide specific, actionable guidance for creating a realistic virtual try-on.
        """

        # Convert PIL images to bytes
        person_img_bytes = io.BytesIO()
        person_image.save(person_img_bytes, format='PNG')
        person_img_bytes = person_img_bytes.getvalue()
        
        clothing_img_bytes = io.BytesIO()
        clothing_image.save(clothing_img_bytes, format='PNG')
        clothing_img_bytes = clothing_img_bytes.getvalue()
        
        # Generate analysis
        response = model.generate_content([
            prompt,
            {
                "mime_type": "image/png",
                "data": person_img_bytes
            },
            {
                "mime_type": "image/png", 
                "data": clothing_img_bytes
            }
        ])
        
        return response.text, True
        
    except Exception as e:
        return f"Error analyzing images: {str(e)}", False

def create_basic_tryon_composite(person_image, clothing_image):
    """Create a basic virtual try-on composite using image processing"""
    try:
        # Convert PIL to OpenCV format
        person_cv = cv2.cvtColor(np.array(person_image), cv2.COLOR_RGB2BGR)
        clothing_cv = cv2.cvtColor(np.array(clothing_image), cv2.COLOR_RGB2BGR)
        
        # Get image dimensions
        person_h, person_w = person_cv.shape[:2]
        clothing_h, clothing_w = clothing_cv.shape[:2]
        
        # Resize clothing to fit person (basic scaling)
        scale_factor = min(person_w * 0.6 / clothing_w, person_h * 0.7 / clothing_h)
        new_w = int(clothing_w * scale_factor)
        new_h = int(clothing_h * scale_factor)
        
        clothing_resized = cv2.resize(clothing_cv, (new_w, new_h))
        
        # Position clothing on person (center upper body)
        y_offset = int(person_h * 0.2)  # Start from 20% down
        x_offset = int((person_w - new_w) / 2)  # Center horizontally
        
        # Create a copy of the person image
        result = person_cv.copy()
        
        # Ensure we don't go out of bounds
        if y_offset + new_h <= person_h and x_offset + new_w <= person_w:
            # Simple overlay (this is very basic - real implementation would need advanced blending)
            result[y_offset:y_offset+new_h, x_offset:x_offset+new_w] = clothing_resized
        
        # Convert back to PIL
        result_rgb = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
        result_pil = Image.fromarray(result_rgb)
        
        return result_pil, True
        
    except Exception as e:
        return None, False

def validate_image(image):
    """Validate uploaded image"""
    if image is None:
        return False, "No image uploaded"
    
    # Check file size (max 10MB)
    if len(image.getvalue()) > 10 * 1024 * 1024:
        return False, "Image size too large. Please upload an image smaller than 10MB."
    
    try:
        # Try to open the image
        img = Image.open(image)
        
        # Check image dimensions
        if img.width < 200 or img.height < 200:
            return False, "Image too small. Please upload an image with minimum dimensions of 200x200 pixels."
        
        if img.width > 4000 or img.height > 4000:
            return False, "Image too large. Please upload an image with maximum dimensions of 4000x4000 pixels."
        
        return True, "Image is valid"
        
    except Exception as e:
        return False, f"Invalid image format: {str(e)}"

def main():
    load_css()
    
    # Header
    st.markdown('<h1 class="main-header">🚀 Advanced Virtual Try-On</h1>', unsafe_allow_html=True)
    st.markdown('<p class="sub-header">Experience the future of fashion with AI-powered virtual try-on technology</p>', unsafe_allow_html=True)
    
    # Feature explanation
    st.markdown("""
    <div class="glass-card">
        <h3 style="color: #667eea; text-align: center;">⚡ Advanced Features</h3>
        <div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(250px, 1fr)); gap: 1rem; margin: 1rem 0;">
            <div style="text-align: center; padding: 1rem;">
                <div style="font-size: 2rem; margin-bottom: 0.5rem;">🧠</div>
                <h4>AI Analysis</h4>
                <p>Gemini 2.0 Flash analyzes your body type and clothing compatibility</p>
            </div>
            <div style="text-align: center; padding: 1rem;">
                <div style="font-size: 2rem; margin-bottom: 0.5rem;">🎨</div>
                <h4>Visual Composite</h4>
                <p>Basic image processing creates a visual representation</p>
            </div>
            <div style="text-align: center; padding: 1rem;">
                <div style="font-size: 2rem; margin-bottom: 0.5rem;">📊</div>
                <h4>Detailed Feedback</h4>
                <p>Comprehensive styling and fit analysis</p>
            </div>
        </div>
    </div>
    """, unsafe_allow_html=True)
    
    # Initialize Gemini
    model = initialize_gemini()
    
    if model is None:
        st.markdown("""
        <div class="glass-card" style="text-align: center; color: #ff6b6b;">
            <h3>⚠️ API Key Required</h3>
            <p>Please add your Gemini API key to the .env file to use this application.</p>
            <p>Get your API key from <a href="https://makersuite.google.com/app/apikey" target="_blank" style="color: #667eea;">Google AI Studio</a></p>
        </div>
        """, unsafe_allow_html=True)
        return
    
    # Upload section
    col1, col2 = st.columns(2)
    
    with col1:
        st.markdown("""
        <div class="glass-card">
            <h3 style="text-align: center; color: #667eea; margin-bottom: 1rem;">📸 Your Photo</h3>
        </div>
        """, unsafe_allow_html=True)
        
        person_image = st.file_uploader(
            "Upload a clear, full-body photo",
            type=['png', 'jpg', 'jpeg'],
            key="person_upload",
            help="Best results with well-lit, front-facing photos"
        )
        
        if person_image:
            is_valid, message = validate_image(person_image)
            if is_valid:
                person_img = Image.open(person_image)
                st.image(person_img, caption="Your Photo", use_column_width=True)
                st.success("✅ Photo ready for analysis!")
            else:
                st.error(f"❌ {message}")
                person_image = None
    
    with col2:
        st.markdown("""
        <div class="glass-card">
            <h3 style="text-align: center; color: #764ba2; margin-bottom: 1rem;">👗 Clothing Item</h3>
        </div>
        """, unsafe_allow_html=True)
        
        clothing_image = st.file_uploader(
            "Upload the clothing item",
            type=['png', 'jpg', 'jpeg'],
            key="clothing_upload",
            help="Clear photos of individual clothing items work best"
        )
        
        if clothing_image:
            is_valid, message = validate_image(clothing_image)
            if is_valid:
                clothing_img = Image.open(clothing_image)
                st.image(clothing_img, caption="Clothing Item", use_column_width=True)
                st.success("✅ Clothing ready for try-on!")
            else:
                st.error(f"❌ {message}")
                clothing_image = None
    
    # Processing section
    if person_image and clothing_image:
        st.markdown("<br>", unsafe_allow_html=True)
        
        col1, col2, col3 = st.columns([1, 2, 1])
        with col2:
            if st.button("🚀 Create Advanced Virtual Try-On", key="advanced_tryon_button"):
                
                person_img = Image.open(person_image)
                clothing_img = Image.open(clothing_image)
                
                # Step 1: AI Analysis
                with st.spinner("🧠 Analyzing with Gemini 2.0 Flash..."):
                    progress = st.progress(0)
                    for i in range(0, 40):
                        time.sleep(0.03)
                        progress.progress(i)
                    
                    analysis_result, analysis_success = analyze_images_for_tryon(person_img, clothing_img, model)
                    progress.progress(40)
                
                # Step 2: Create Visual Composite
                with st.spinner("🎨 Creating visual composite..."):
                    for i in range(40, 80):
                        time.sleep(0.02)
                        progress.progress(i)
                    
                    composite_result, composite_success = create_basic_tryon_composite(person_img, clothing_img)
                    progress.progress(80)
                
                # Step 3: Finalize
                with st.spinner("✨ Finalizing results..."):
                    for i in range(80, 100):
                        time.sleep(0.01)
                        progress.progress(i)
                    progress.progress(100)
                    time.sleep(0.5)
                
                # Display results
                if analysis_success:
                    st.success("🎉 Advanced virtual try-on completed!")
                    
                    # Show composite if available
                    if composite_success and composite_result:
                        st.markdown("""
                        <div class="result-container">
                            <h3 style="text-align: center; color: #667eea;">📱 Visual Composite Result</h3>
                        </div>
                        """, unsafe_allow_html=True)
                        
                        col1, col2, col3 = st.columns([1, 2, 1])
                        with col2:
                            st.image(composite_result, caption="Basic Virtual Try-On Composite", use_column_width=True)
                        
                        st.warning("""
                        **Note**: This is a basic composite for demonstration. Professional virtual try-on requires 
                        specialized models like TryOnGAN, VITON, or commercial APIs like Zyler.
                        """)
                    
                    # Show AI analysis
                    st.markdown("""
                    <div class="result-container">
                        <h3 style="text-align: center; color: #764ba2;">🧠 AI Fashion Analysis</h3>
                    </div>
                    """, unsafe_allow_html=True)
                    
                    st.markdown(analysis_result)
                    
                    # Next steps
                    st.markdown("---")
                    st.markdown("""
                    <div class="glass-card">
                        <h4 style="color: #667eea;">🚀 For Professional Virtual Try-On</h4>
                        <p>For production-quality virtual try-on, consider integrating:</p>
                        <ul>
                            <li><strong>Zyler API:</strong> Professional virtual fitting technology</li>
                            <li><strong>TryNow:</strong> Real-time virtual try-on solutions</li>
                            <li><strong>VITON:</strong> Virtual try-on network for academic use</li>
                            <li><strong>StyleGAN + ControlNet:</strong> Custom AI model pipeline</li>
                        </ul>
                    </div>
                    """, unsafe_allow_html=True)
                
                else:
                    st.error(f"❌ Analysis failed: {analysis_result}")
    
    else:
        # Instructions
        st.markdown("""
        <div class="glass-card">
            <h3 style="color: #667eea; text-align: center;">🎯 How to Use Advanced Virtual Try-On</h3>
            <div style="display: grid; grid-template-columns: repeat(auto-fit, minmax(200px, 1fr)); gap: 2rem; margin: 2rem 0;">
                <div style="text-align: center;">
                    <div style="font-size: 3rem; margin-bottom: 1rem;">📸</div>
                    <h4>Upload Your Photo</h4>
                    <p>Clear, well-lit, front-facing photo works best</p>
                </div>
                <div style="text-align: center;">
                    <div style="font-size: 3rem; margin-bottom: 1rem;">👗</div>
                    <h4>Choose Clothing</h4>
                    <p>Upload the item you want to try on</p>
                </div>
                <div style="text-align: center;">
                    <div style="font-size: 3rem; margin-bottom: 1rem;">🚀</div>
                    <h4>Get Results</h4>
                    <p>AI analysis + visual composite</p>
                </div>
            </div>
        </div>
        """, unsafe_allow_html=True)
    
    # Sidebar
    with st.sidebar:
        st.markdown("""
        <div class="glass-card">
            <h3 style="color: #667eea;">💡 Pro Tips</h3>
            <ul style="list-style-type: none; padding: 0;">
                <li style="margin: 1rem 0;">📷 Use high-resolution images (min 200x200)</li>
                <li style="margin: 1rem 0;">💡 Good lighting is crucial</li>
                <li style="margin: 1rem 0;">🧍 Full-body shots work better</li>
                <li style="margin: 1rem 0;">👕 Single clothing items are easier to process</li>
                <li style="margin: 1rem 0;">🎯 Neutral backgrounds preferred</li>
                <li style="margin: 1rem 0;">⚡ Powered by Gemini 2.0 Flash</li>
            </ul>
        </div>
        """, unsafe_allow_html=True)

if __name__ == "__main__":
    main()