app.py

import cv2
import numpy as np
import argparse
import time
import sys
import os
from pathlib import Path
import streamlit as st
from PIL import Image, ImageTk
import tempfile
import io
import threading
import tkinter as tk
from tkinter import filedialog, Label, Button, Frame

# Constants
MODEL_DIR = "models"
TEMP_DIR = "temp"

def parse_args():
    parser = argparse.ArgumentParser(description='Advanced Virtual Try-On')
    parser.add_argument('--garment', type=str, help='Path to garment image')
    parser.add_argument('--webcam', type=int, default=0, help='Webcam index to use')
    parser.add_argument('--resolution', type=str, default='640x480', help='Camera resolution')
    parser.add_argument('--streamlit', action='store_true', help='Run in Streamlit mode')
    parser.add_argument('--tkinter', action='store_true', help='Run with Tkinter UI')
    return parser.parse_args()

class HumanPoseEstimator:
    """Human pose estimation using OpenPose or similar model"""
    
    def __init__(self):
        # Create model directory if it doesn't exist
        if not os.path.exists(MODEL_DIR):
            os.makedirs(MODEL_DIR)
        
        # Download pose model if not present (simplified here)
        self.download_models_if_needed()
        
        # Load COCO body model for OpenPose
        self.BODY_PARTS = {
            "Nose": 0, "Neck": 1, "RShoulder": 2, "RElbow": 3, "RWrist": 4,
            "LShoulder": 5, "LElbow": 6, "LWrist": 7, "RHip": 8, "RKnee": 9,
            "RAnkle": 10, "LHip": 11, "LKnee": 12, "LAnkle": 13, "REye": 14,
            "LEye": 15, "REar": 16, "LEar": 17, "Background": 18
        }
        
        self.POSE_PAIRS = [
            ["Neck", "RShoulder"], ["Neck", "LShoulder"], ["RShoulder", "RElbow"],
            ["RElbow", "RWrist"], ["LShoulder", "LElbow"], ["LElbow", "LWrist"],
            ["Neck", "RHip"], ["RHip", "RKnee"], ["RKnee", "RAnkle"], ["Neck", "LHip"],
            ["LHip", "LKnee"], ["LKnee", "LAnkle"], ["Neck", "Nose"], ["Nose", "REye"],
            ["REye", "REar"], ["Nose", "LEye"], ["LEye", "LEar"]
        ]
        
        # Load OpenPose network
        self.net = self.load_pose_model()
        
        print("Pose estimation model loaded successfully")
    
    def download_models_if_needed(self):
        """Download models if not present"""
        # Model paths
        pose_model_path = os.path.join(MODEL_DIR, "pose_iter_440000.caffemodel")
        pose_proto_path = os.path.join(MODEL_DIR, "pose_deploy_linevec.prototxt")
        
        # Check if models exist
        if not os.path.exists(pose_model_path) or not os.path.exists(pose_proto_path):
            print("Models not found. Downloading pose estimation models...")
            # Normally we'd download the models here using requests or urllib
            # For this example, we'll direct the user to download them manually
            print("Please download the OpenPose model:")
            print(f"1. Download pose_iter_440000.caffemodel to {MODEL_DIR}")
            print(f"2. Download pose_deploy_linevec.prototxt to {MODEL_DIR}")
            print("Models can be found at: https://github.com/CMU-Perceptual-Computing-Lab/openpose/tree/master/models")
            
            # Create directory for models
            Path(MODEL_DIR).mkdir(parents=True, exist_ok=True)
            
            # For demonstration, we'll create dummy files with instructions
            with open(pose_proto_path, 'w') as f:
                f.write("# Download the actual model file from OpenPose repository")
            with open(pose_model_path, 'w') as f:
                f.write("# Download the actual model file from OpenPose repository")
                
            print("Created placeholder files. Replace with actual model files before running.")
    
    def load_pose_model(self):
        """Load the pose detection model"""
        try:
            # Try to load the OpenPose model
            model_path = os.path.join(MODEL_DIR, "pose_iter_440000.caffemodel")
            config_path = os.path.join(MODEL_DIR, "pose_deploy_linevec.prototxt")
            
            if os.path.getsize(model_path) < 1000:  # Placeholder file
                print("Warning: Using placeholder model file. Results will be simulated.")
                # Fall back to a basic pose estimation
                return None
            
            net = cv2.dnn.readNetFromCaffe(config_path, model_path)
            
            # Try to use GPU if available - safely check for CUDA availability
            try:
                # Check if CUDA is available by testing if the cv2.cuda module exists
                cuda_available = False
                if hasattr(cv2, 'cuda'):
                    try:
                        cuda_available = cv2.cuda.getCudaEnabledDeviceCount() > 0
                    except:
                        cuda_available = False
                
                if cuda_available:
                    print("CUDA is available. Using GPU acceleration.")
                    net.setPreferableBackend(cv2.dnn.DNN_BACKEND_CUDA)
                    net.setPreferableTarget(cv2.dnn.DNN_TARGET_CUDA)
                else:
                    print("CUDA is not available. Using CPU.")
                    net.setPreferableBackend(cv2.dnn.DNN_BACKEND_DEFAULT)
                    net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)
            except Exception as cuda_err:
                print(f"Error checking CUDA availability: {cuda_err}. Using CPU instead.")
                net.setPreferableBackend(cv2.dnn.DNN_BACKEND_DEFAULT)
                net.setPreferableTarget(cv2.dnn.DNN_TARGET_CPU)
                
            return net
        except Exception as e:
            print(f"Error loading pose model: {e}")
            print("Falling back to simulation mode")
            return None
    
    def estimate_pose(self, frame):
        """Estimate human pose in the frame"""
        frame_height, frame_width = frame.shape[:2]
        
        # If we don't have the actual model, simulate pose detection
        if self.net is None:
            return self.simulate_pose(frame)
        
        # Prepare input for the network
        input_blob = cv2.dnn.blobFromImage(frame, 1.0 / 255, (368, 368), (0, 0, 0), swapRB=False, crop=False)
        self.net.setInput(input_blob)
        
        # Forward pass through the network
        output = self.net.forward()
        
        # Parse keypoints
        keypoints = []
        threshold = 0.1
        
        for i in range(len(self.BODY_PARTS) - 1):  # Exclude background
            # Get confidence map
            prob_map = output[0, i, :, :]
            prob_map = cv2.resize(prob_map, (frame_width, frame_height))
            
            # Find global maximum
            _, confidence, _, point = cv2.minMaxLoc(prob_map)
            
            if confidence > threshold:
                keypoints.append((point[0], point[1], confidence))
            else:
                keypoints.append(None)
        
        return keypoints
    
    def simulate_pose(self, frame):
        """Simulate pose detection when model isn't available"""
        # Use face detection to estimate body position
        face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
        gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
        faces = face_cascade.detectMultiScale(gray, 1.3, 5)
        
        # Initialize keypoints
        keypoints = [None] * (len(self.BODY_PARTS) - 1)
        
        if len(faces) > 0:
            # Get the largest face
            x, y, w, h = max(faces, key=lambda rect: rect[2] * rect[3])
            
            # Center of face
            face_center_x = x + w // 2
            face_center_y = y + h // 2
            
            # Estimate keypoints based on face position
            frame_height, frame_width = frame.shape[:2]
            
            # Nose (center of face)
            keypoints[self.BODY_PARTS["Nose"]] = (face_center_x, face_center_y, 0.9)
            
            # Neck (below face)
            neck_y = y + h + h // 4
            keypoints[self.BODY_PARTS["Neck"]] = (face_center_x, neck_y, 0.8)
            
            # Shoulders (on either side of neck)
            shoulder_y = neck_y + h // 8
            keypoints[self.BODY_PARTS["RShoulder"]] = (face_center_x - w, shoulder_y, 0.7)
            keypoints[self.BODY_PARTS["LShoulder"]] = (face_center_x + w, shoulder_y, 0.7)
            
            # Approximate other body parts
            keypoints[self.BODY_PARTS["RHip"]] = (face_center_x - w//2, frame_height - h*2, 0.5)
            keypoints[self.BODY_PARTS["LHip"]] = (face_center_x + w//2, frame_height - h*2, 0.5)
            
        return keypoints
    
    def draw_skeleton(self, frame, keypoints):
        """Draw skeleton on the frame for visualization"""
        # Draw keypoints
        for i, keypoint in enumerate(keypoints):
            if keypoint:
                cv2.circle(frame, (int(keypoint[0]), int(keypoint[1])), 8, (0, 255, 255), thickness=-1, lineType=cv2.FILLED)
        
        # Draw connections
        for pair in self.POSE_PAIRS:
            part_from = self.BODY_PARTS[pair[0]]
            part_to = self.BODY_PARTS[pair[1]]
            
            if keypoints[part_from] and keypoints[part_to]:
                cv2.line(frame, 
                         (int(keypoints[part_from][0]), int(keypoints[part_from][1])),
                         (int(keypoints[part_to][0]), int(keypoints[part_to][1])),
                         (0, 255, 0), 3)
        
        return frame

class GarmentProcessor:
    """Process garment images for virtual try-on"""
    
    def __init__(self):
        # Create temp directory for processed images
        if not os.path.exists(TEMP_DIR):
            os.makedirs(TEMP_DIR)
    
    def load_garment(self, path):
        """Load and preprocess a garment image"""
        # Load the image
        garment = cv2.imread(path, cv2.IMREAD_UNCHANGED)
        
        if garment is None:
            raise FileNotFoundError(f"Could not load garment image from {path}")
        
        # If garment doesn't have alpha channel, add one
        if garment.shape[2] == 3:
            garment = self.remove_background(garment)
        
        return garment
    
    def remove_background(self, img):
        """Remove background from garment image including black backgrounds"""
        # Convert to RGBA
        rgba = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)
        
        # Get image dimensions
        h, w = img.shape[:2]
        
        # Create an initial mask
        # Instead of simple thresholding which fails for black clothes,
        # we'll use a combination of techniques
        
        # 1. Start with an approximate mask using color detection
        # Convert to different color spaces
        hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        
        # Create masks for different color spaces
        # Detect very dark regions (potential black backgrounds)
        _, dark_mask = cv2.threshold(gray, 30, 255, cv2.THRESH_BINARY)
        
        # Detect edges - useful for finding garment boundaries
        edges = cv2.Canny(img, 50, 150)
        kernel = np.ones((5,5), np.uint8)
        dilated_edges = cv2.dilate(edges, kernel, iterations=2)
        
        # Create initial GrabCut mask
        # 0 = background, 1 = foreground, 2 = probable background, 3 = probable foreground
        gc_mask = np.zeros(img.shape[:2], np.uint8)
        
        # Mark the borders as likely background
        border_width = w // 10  # 10% of width
        gc_mask[:border_width, :] = 2
        gc_mask[-border_width:, :] = 2
        gc_mask[:, :border_width] = 2
        gc_mask[:, -border_width:] = 2
        
        # Mark the center as likely foreground
        center_rect = (border_width, border_width, w - 2*border_width, h - 2*border_width)
        cv2.rectangle(gc_mask, (center_rect[0], center_rect[1]), 
                     (center_rect[0] + center_rect[2], center_rect[1] + center_rect[3]), 3, -1)
        
        # Use edges to refine foreground
        gc_mask[dilated_edges > 0] = 1
        
        # Initialize GrabCut background and foreground models
        bgd_model = np.zeros((1, 65), np.float64)
        fgd_model = np.zeros((1, 65), np.float64)
        
        # Run GrabCut algorithm
        try:
            cv2.grabCut(img, gc_mask, center_rect, bgd_model, fgd_model, 5, cv2.GC_INIT_WITH_MASK)
        except Exception as e:
            print(f"GrabCut failed: {e}. Using fallback method.")
            # Fallback to simpler method if GrabCut fails
            # Create a simple mask based on color threshold
            _, mask = cv2.threshold(gray, 200, 255, cv2.THRESH_BINARY_INV)
            kernel = np.ones((5, 5), np.uint8)
            mask = cv2.morphologyEx(mask, cv2.MORPH_OPEN, kernel)
            mask = cv2.morphologyEx(mask, cv2.MORPH_CLOSE, kernel)
            rgba[:, :, 3] = mask
            return rgba
        
        # Create final mask where 1 and 3 are foreground
        final_mask = np.where((gc_mask == 1) | (gc_mask == 3), 255, 0).astype('uint8')
        
        # Clean up mask with morphological operations
        kernel = np.ones((5, 5), np.uint8)
        final_mask = cv2.morphologyEx(final_mask, cv2.MORPH_OPEN, kernel)
        final_mask = cv2.morphologyEx(final_mask, cv2.MORPH_CLOSE, kernel)
        
        # Dilate the mask slightly to include edge details
        final_mask = cv2.dilate(final_mask, kernel, iterations=1)
        
        # Apply mask to alpha channel
        rgba[:, :, 3] = final_mask
        
        print("Added transparency to garment image with advanced background removal")
        return rgba
    
    def warp_garment(self, garment, keypoints, frame_size, sizing_params=None):
        """Warp the garment to fit the detected pose"""
        frame_height, frame_width = frame_size
        
        # Set default sizing parameters if not provided
        if sizing_params is None:
            sizing_params = {
                'width_scale': 1.2,  # Default width scale - reduced for better fit
                'height_scale': 1.1   # Default height scale
            }
        
        # If no valid keypoints, return original garment
        if not keypoints or not keypoints[1]:  # Check if neck keypoint exists
            return garment
        
        # Get relevant keypoints for garment warping
        neck = keypoints[1]
        right_shoulder = keypoints[2]
        left_shoulder = keypoints[5]
        right_hip = keypoints[8]
        left_hip = keypoints[11]
        
        if not all([neck, right_shoulder, left_shoulder]):
            return garment  # Not enough keypoints
        
        # Calculate garment dimensions based on body
        if right_shoulder and left_shoulder:
            # Calculate Euclidean distance between shoulders
            shoulder_width = np.linalg.norm(
                [left_shoulder[0] - right_shoulder[0], left_shoulder[1] - right_shoulder[1]]
            )
            # Get angle between shoulders for rotation
            shoulder_angle = np.arctan2(
                left_shoulder[1] - right_shoulder[1], 
                left_shoulder[0] - right_shoulder[0]
            ) * 180 / np.pi
        else:
            shoulder_width = frame_width * 0.2  # Fallback
            shoulder_angle = 0
        
        # Calculate torso measurements for better proportions
        torso_height = 0
        if right_hip and left_hip and neck:
            # Distance from neck to hips
            hip_center_x = (left_hip[0] + right_hip[0]) / 2
            hip_center_y = (left_hip[1] + right_hip[1]) / 2
            torso_height = np.linalg.norm([hip_center_x - neck[0], hip_center_y - neck[1]])
        else:
            # Estimate torso height based on shoulder width and typical human proportions
            # Use a more conservative estimate for better fit
            torso_height = shoulder_width * 1.4  # Adjusted from 1.6 for better fit
        
        # Calculate body size estimate
        body_width = shoulder_width * 1.1  # Slightly wider than shoulders (reduced from 1.2)
        
        # Get garment original dimensions
        garment_height, garment_width = garment.shape[:2]
        
        # Calculate aspect ratio of the garment
        garment_aspect = garment_width / float(garment_height) if garment_height > 0 else 1.0
        
        # Calculate ideal dimensions for the garment based on body
        # For t-shirts: width should cover shoulders plus some extra, height should cover torso
        ideal_width = shoulder_width * sizing_params['width_scale']
        ideal_height = torso_height * sizing_params['height_scale']
        
        # Maintain aspect ratio while fitting to body
        if (ideal_width / ideal_height) > garment_aspect:
            # Width-constrained: use ideal width, calculate height to maintain aspect
            target_width = int(ideal_width)
            target_height = int(target_width / garment_aspect)
        else:
            # Height-constrained: use ideal height, calculate width to maintain aspect
            target_height = int(ideal_height)
            target_width = int(target_height * garment_aspect)
        
        # Resize garment to target dimensions
        garment_resized = self.resize_garment(garment, target_width, target_height)
        
        # Apply rotation if shoulders aren't level (beyond a small threshold)
        if abs(shoulder_angle) > 5:
            center = (garment_resized.shape[1] // 2, garment_resized.shape[0] // 2)
            rotation_matrix = cv2.getRotationMatrix2D(center, shoulder_angle, 1.0)
            garment_resized = cv2.warpAffine(
                garment_resized, rotation_matrix, 
                (garment_resized.shape[1], garment_resized.shape[0]),
                flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_TRANSPARENT
            )
        
        # Apply perspective transform to better fit the body shape
        try:
            # Only apply perspective transform if we have all four corners (shoulders and hips)
            if all([right_shoulder, left_shoulder, right_hip, left_hip]):
                # Define source points (corners of the garment)
                garment_h, garment_w = garment_resized.shape[:2]
                src_pts = np.array([
                    [0, 0],  # Top-left
                    [garment_w, 0],  # Top-right
                    [garment_w, garment_h],  # Bottom-right
                    [0, garment_h]  # Bottom-left
                ], dtype=np.float32)
                
                # Define destination points based on body keypoints
                # Scale factors to find garment edges from body keypoints
                top_width_factor = 1.1  # How much wider than shoulders at top
                bottom_width_factor = 0.9  # How much wider than hips at bottom
                
                # Calculate destination points
                top_left_x = left_shoulder[0] - (shoulder_width * (top_width_factor - 1) / 2)
                top_right_x = right_shoulder[0] + (shoulder_width * (top_width_factor - 1) / 2)
                bottom_left_x = left_hip[0] - (shoulder_width * (bottom_width_factor - 1) / 2)
                bottom_right_x = right_hip[0] + (shoulder_width * (bottom_width_factor - 1) / 2)
                
                # Get y-coordinates (adjust top to be at collar position)
                top_y = (left_shoulder[1] + right_shoulder[1]) / 2 - garment_h * 0.2
                bottom_y = top_y + garment_h * 0.95  # Slightly higher than full height for better look
                
                dst_pts = np.array([
                    [top_left_x, top_y],  # Top-left
                    [top_right_x, top_y],  # Top-right
                    [bottom_right_x, bottom_y],  # Bottom-right
                    [bottom_left_x, bottom_y]  # Bottom-left
                ], dtype=np.float32)
                
                # Get perspective transform matrix
                M = cv2.getPerspectiveTransform(src_pts, dst_pts)
                
                # Apply perspective transform
                # Make output size large enough to contain the warped garment
                output_size = (frame_width, frame_height)
                warped = cv2.warpPerspective(garment_resized, M, output_size, 
                                            flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_TRANSPARENT)
                
                # Crop to the actual garment size to avoid large transparent areas
                # Find non-zero alpha channel pixels
                alpha = warped[:, :, 3]
                coords = cv2.findNonZero(alpha)
                
                if coords is not None and len(coords) > 0:
                    x, y, w, h = cv2.boundingRect(coords)
                    warped = warped[y:y+h, x:x+w]
                    return warped
        
        except Exception as e:
            print(f"Perspective transform failed: {e}")
            # Continue with the regular garment if perspective transform fails
            pass
        
        print(f"Resized garment to fit body: {target_width}x{target_height} px")
        return garment_resized
    
    def resize_garment(self, garment, target_width=None, target_height=None):
        """Resize garment maintaining aspect ratio"""
        if garment is None:
            return None
            
        garment_height, garment_width = garment.shape[:2]
        aspect = garment_width / float(garment_height)
        
        if target_width is not None:
            new_width = target_width
            new_height = int(new_width / aspect)
        elif target_height is not None:
            new_height = target_height
            new_width = int(new_height * aspect)
        else:
            return garment  # No resize if no dimensions provided
        
        # High-quality resize
        resized = cv2.resize(garment, (new_width, new_height), interpolation=cv2.INTER_LANCZOS4)
        return resized

class AdvancedVirtualTryOn:
    """Main class for the virtual try-on system"""
    
    def __init__(self, garment_path, camera_index=0, resolution="640x480", streamlit_mode=False):
        # Parse resolution
        width, height = map(int, resolution.split('x'))
        
        # Set Streamlit mode
        self.streamlit_mode = streamlit_mode
        
        # Initialize components
        self.pose_estimator = HumanPoseEstimator()
        self.garment_processor = GarmentProcessor()
        
        # Load garment
        self.garment = self.garment_processor.load_garment(garment_path)
        
        # Initialize camera if not in Streamlit mode
        if not streamlit_mode:
            self.camera = cv2.VideoCapture(camera_index)
            self.camera.set(cv2.CAP_PROP_FRAME_WIDTH, width)
            self.camera.set(cv2.CAP_PROP_FRAME_HEIGHT, height)
            
            if not self.camera.isOpened():
                raise RuntimeError("Could not open camera")
        
        # Performance tracking
        self.prev_frame_time = 0
        self.new_frame_time = 0
        self.fps = 0
        
        # Garment positioning and sizing parameters - adjusted for better default fit
        self.vertical_offset = 0.05
        self.width_scale = 1.2   # Reduced from 1.5 for a more realistic fit
        self.height_scale = 1.1  # Scale factor for garment height relative to torso
        self.collar_position = 0.20  # Increased to position collar higher on neck
        
        # UI modes
        self.debug_mode = False
        self.show_controls = True
        self.fullscreen_mode = False
        
        # For smoother processing and better performance
        self.skip_frames = 0  # Process every frame by default
        self.frame_counter = 0
        self.last_warped_garment = None
        self.last_keypoints = None
        
        print("Advanced Virtual Try-On initialized.")
        if not streamlit_mode:
            print("Starting camera feed...")
    
    def update_fps(self):
        """Calculate and update FPS"""
        self.new_frame_time = time.time()
        self.fps = 1 / (self.new_frame_time - self.prev_frame_time) if (self.new_frame_time - self.prev_frame_time) > 0 else 0
        self.prev_frame_time = self.new_frame_time
        return int(self.fps)
    
    def overlay_garment(self, frame, keypoints):
        """Overlay the garment on the frame"""
        frame_height, frame_width = frame.shape[:2]
        
        # Check if we have valid keypoints
        if not keypoints or not keypoints[1]:  # Neck keypoint
            # Use last valid keypoints if available for smoothness
            if self.last_keypoints and self.last_warped_garment is not None:
                keypoints = self.last_keypoints
            else:
                return frame
        else:
            # Store last valid keypoints for smooth transitions
            self.last_keypoints = keypoints
        
        try:
            # Get key body keypoints
            neck = keypoints[1]
            right_shoulder = keypoints[2]
            left_shoulder = keypoints[5]
            
            if not all([neck, right_shoulder, left_shoulder]):
                if self.last_warped_garment is not None:
                    # Use last valid garment if available
                    warped_garment = self.last_warped_garment
                else:
                    return frame
            else:
                # Calculate shoulder midpoint for better centering
                shoulder_center_x = (right_shoulder[0] + left_shoulder[0]) / 2
                shoulder_center_y = (right_shoulder[1] + left_shoulder[1]) / 2
                
                # Check if we should skip processing this frame (for performance)
                self.frame_counter += 1
                if self.skip_frames > 0 and self.frame_counter % (self.skip_frames + 1) != 0 and self.last_warped_garment is not None:
                    warped_garment = self.last_warped_garment
                else:
                    # Pass current sizing parameters to warp_garment
                    sizing_params = {
                        'width_scale': self.width_scale,
                        'height_scale': self.height_scale
                    }
                    
                    # Warp garment to fit the body
                    warped_garment = self.garment_processor.warp_garment(
                        self.garment, keypoints, (frame_height, frame_width), sizing_params
                    )
                    
                    # Save for potential reuse
                    self.last_warped_garment = warped_garment
            
            if warped_garment is None:
                return frame
            
            # Calculate position
            garment_height, garment_width = warped_garment.shape[:2]
            
            if all([neck, right_shoulder, left_shoulder]):
                # Center horizontally on shoulders rather than neck for better alignment
                center_x = int(shoulder_center_x)
                
                # Calculate vertical position based on neck and shoulders
                # Position garment higher for more natural look
                shoulder_y = (right_shoulder[1] + left_shoulder[1]) / 2
                center_y = int(neck[1] + (shoulder_y - neck[1]) * 0.5 - (garment_height * self.collar_position))
            else:
                # Fallback to last known position
                center_x = frame_width // 2
                center_y = frame_height // 3
            
            # Calculate top-left corner
            x1 = center_x - garment_width // 2
            y1 = center_y
            
            # Ensure coordinates are within frame
            x1 = max(0, x1)
            y1 = max(0, y1)
            x2 = min(frame_width, x1 + garment_width)
            y2 = min(frame_height, y1 + garment_height)
            
            # Calculate source region in garment
            g_x1 = 0
            g_y1 = 0
            g_x2 = x2 - x1
            g_y2 = y2 - y1
            
            if g_x2 <= 0 or g_y2 <= 0 or g_x1 >= garment_width or g_y1 >= garment_height:
                return frame
            
            # Adjust if needed
            if g_x2 > garment_width:
                g_x2 = garment_width
                x2 = x1 + g_x2
            
            if g_y2 > garment_height:
                g_y2 = garment_height
                y2 = y1 + g_y2
            
            # Extract regions
            roi = frame[y1:y2, x1:x2].copy()
            garment_roi = warped_garment[g_y1:g_y2, g_x1:g_x2].copy()
            
            if roi.shape[:2] != garment_roi.shape[:2]:
                return frame
            
            # Improved alpha blending with edge feathering
            alpha = garment_roi[:, :, 3] / 255.0
            
            # Apply Gaussian blur to alpha channel for softer edges
            alpha_blur = cv2.GaussianBlur(alpha, (5, 5), 0)
            alpha_blur = np.repeat(alpha_blur[:, :, np.newaxis], 3, axis=2)
            
            # Blend images with the smoothed alpha
            blended = roi * (1 - alpha_blur) + garment_roi[:, :, :3] * alpha_blur
            
            # Apply color correction to match lighting
            # This helps the garment look more natural in the scene
            mean_roi = np.mean(roi, axis=(0, 1))
            mean_garment = np.mean(garment_roi[:, :, :3], axis=(0, 1))
            
            # Apply subtle lighting adjustment (limit the effect for realism)
            lighting_factor = 0.3
            lighting_adjustment = (mean_roi - mean_garment) * lighting_factor
            adjusted_garment = np.clip(garment_roi[:, :, :3] + lighting_adjustment, 0, 255)
            
            # Final blending with lighting adjustment
            final_blend = roi * (1 - alpha_blur) + adjusted_garment * alpha_blur
            frame[y1:y2, x1:x2] = final_blend
            
        except Exception as e:
            print(f"Error overlaying garment: {e}")
        
        return frame
    
    def process_frame(self, frame):
        """Process a single frame, returning the processed frame"""
        # Flip for mirror effect
        frame = cv2.flip(frame, 1)
        
        # Estimate pose
        keypoints = self.pose_estimator.estimate_pose(frame)
        
        # Overlay garment
        frame = self.overlay_garment(frame, keypoints)
        
        # Draw skeleton in debug mode
        if self.debug_mode:
            frame = self.pose_estimator.draw_skeleton(frame, keypoints)
        
        # Calculate and display FPS
        fps = self.update_fps()
        cv2.putText(frame, f"FPS: {fps}", (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 
                   1, (0, 255, 0), 2, cv2.LINE_AA)
        
        # Display current garment fit parameters
        if self.show_controls and not self.streamlit_mode:
            # Display fitting instructions
            instructions = [
                "Controls:",
                f"Width Scale: {self.width_scale:.1f} (+/- to adjust)",
                f"Height Scale: {self.height_scale:.1f} (up/down arrows)",
                f"Collar Position: {self.collar_position:.2f} (</> to adjust)",
                f"Performance: {'High' if self.skip_frames>0 else 'Normal'} (f key)",
                f"Display: {'Fullscreen' if self.fullscreen_mode else 'Window'} (s key)",
                "'d' - Toggle debug | 'q' - Quit | 'c' - Hide controls"
            ]
            
            y_pos = 70
            for line in instructions:
                cv2.putText(frame, line, (10, y_pos), cv2.FONT_HERSHEY_SIMPLEX, 
                           0.5, (0, 255, 0), 1, cv2.LINE_AA)
                y_pos += 25
        elif not self.streamlit_mode:
            cv2.putText(frame, "Press 'c' for controls", (10, 70), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 0), 1, cv2.LINE_AA)
        
        return frame

    def run(self):
        """Main application loop"""
        if self.streamlit_mode:
            print("Streamlit mode is active. The main loop will be controlled by the Streamlit app.")
            return
            
        print("Advanced Virtual Try-On started. Press 'q' to quit, 'd' to toggle debug mode.")
        print("Use '+'/'-' to adjust garment width, 'up'/'down' arrows to adjust height.")
        print("Use 'c' to toggle control instructions, 'f' to toggle frame skipping for better performance.")
        print("Press 's' to toggle fullscreen mode.")
        
        # Create a resizable window
        cv2.namedWindow('Advanced Virtual Try-On', cv2.WINDOW_NORMAL)
        
        while self.camera.isOpened():
            success, frame = self.camera.read()
            if not success:
                print("Failed to capture frame")
                break
            
            # Process the frame
            frame = self.process_frame(frame)
            
            # Display the result
            cv2.imshow('Advanced Virtual Try-On', frame)
            
            # Handle key presses
            key = cv2.waitKey(1) & 0xFF
            
            # Quit
            if key == ord('q'):
                break
                
            # Toggle debug mode
            elif key == ord('d'):
                self.debug_mode = not self.debug_mode
                print(f"Debug mode: {'ON' if self.debug_mode else 'OFF'}")
                
            # Toggle control display
            elif key == ord('c'):
                self.show_controls = not self.show_controls
                
            # Toggle fullscreen mode
            elif key == ord('s'):
                self.fullscreen_mode = not self.fullscreen_mode
                if self.fullscreen_mode:
                    cv2.setWindowProperty('Advanced Virtual Try-On', cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_FULLSCREEN)
                    print("Fullscreen mode enabled")
                else:
                    cv2.setWindowProperty('Advanced Virtual Try-On', cv2.WND_PROP_FULLSCREEN, cv2.WINDOW_NORMAL)
                    print("Window mode enabled")
                
            # Adjust width scale
            elif key == ord('+') or key == ord('='):  # = is on the same key as + without shift
                self.width_scale = min(3.0, self.width_scale + 0.1)
                print(f"Width scale: {self.width_scale:.1f}")
                
            elif key == ord('-'):
                self.width_scale = max(0.8, self.width_scale - 0.1)
                print(f"Width scale: {self.width_scale:.1f}")
                
            # Adjust height scale
            elif key == 82:  # Up arrow
                self.height_scale = min(2.0, self.height_scale + 0.1)
                print(f"Height scale: {self.height_scale:.1f}")
                
            elif key == 84:  # Down arrow
                self.height_scale = max(0.6, self.height_scale - 0.1)
                print(f"Height scale: {self.height_scale:.1f}")
                
            # Adjust collar position
            elif key == ord(',') or key == ord('<'):
                self.collar_position = max(0.05, self.collar_position - 0.01)
                print(f"Collar position: {self.collar_position:.2f}")
                
            elif key == ord('.') or key == ord('>'):
                self.collar_position = min(0.3, self.collar_position + 0.01)
                print(f"Collar position: {self.collar_position:.2f}")
                
            # Toggle performance mode
            elif key == ord('f'):
                # Toggle between 0, 1, and 2 frame skips
                self.skip_frames = (self.skip_frames + 1) % 3
                print(f"Performance mode: {'High (skip {self.skip_frames} frames)' if self.skip_frames>0 else 'Normal'}")
        
        # Clean up
        self.clean_up()
    
    def clean_up(self):
        """Clean up resources"""
        if hasattr(self, 'camera') and not self.streamlit_mode and self.camera.isOpened():
            self.camera.release()
        cv2.destroyAllWindows()
        print("Application closed.")

class TkinterUI:
    """Tkinter UI for the virtual try-on application"""
    
    def __init__(self, webcam_index=0, resolution="640x480"):
        self.webcam_index = webcam_index
        self.width, self.height = map(int, resolution.split('x'))
        self.app = None
        self.running = False
        self.garment_path = None
        self.root = None
        self.webcam_label = None
        self.update_interval = 10  # Update every 10ms
    
    def start(self):
        """Start the Tkinter UI"""
        self.root = tk.Tk()
        self.root.title("Virtual Try-On")
        self.root.geometry(f"{self.width + 300}x{self.height + 100}")
        self.root.resizable(True, True)
        self.root.protocol("WM_DELETE_WINDOW", self.on_close)
        
        # Main frame
        main_frame = Frame(self.root)
        main_frame.pack(fill=tk.BOTH, expand=True, padx=10, pady=10)
        
        # Left side - webcam and controls
        left_frame = Frame(main_frame)
        left_frame.pack(side=tk.LEFT, fill=tk.BOTH, expand=True)
        
        # Webcam frame
        webcam_frame = Frame(left_frame, width=self.width, height=self.height)
        webcam_frame.pack(pady=10)
        webcam_frame.pack_propagate(0)  # Don't shrink
        
        self.webcam_label = Label(webcam_frame)
        self.webcam_label.pack(fill=tk.BOTH, expand=True)
        
        # Control buttons
        control_frame = Frame(left_frame)
        control_frame.pack(fill=tk.X, pady=10)
        
        upload_btn = Button(control_frame, text="Upload Garment", command=self.upload_garment)
        upload_btn.pack(side=tk.LEFT, padx=5)
        
        start_btn = Button(control_frame, text="Start Try-On", command=self.start_tryon)
        start_btn.pack(side=tk.LEFT, padx=5)
        
        stop_btn = Button(control_frame, text="Stop", command=self.stop_tryon)
        stop_btn.pack(side=tk.LEFT, padx=5)
        
        # Right side - adjustments and garment preview
        right_frame = Frame(main_frame, width=280)
        right_frame.pack(side=tk.RIGHT, fill=tk.Y, padx=10)
        right_frame.pack_propagate(0)  # Don't shrink
        
        # Garment preview
        preview_label = Label(right_frame, text="Garment Preview")
        preview_label.pack(pady=(0, 5))
        
        self.garment_preview = Label(right_frame, text="No garment selected")
        self.garment_preview.pack(pady=5)
        
        # Adjustments
        adjustments_label = Label(right_frame, text="Adjustments")
        adjustments_label.pack(pady=(15, 5))
        
        # Width scale slider
        width_frame = Frame(right_frame)
        width_frame.pack(fill=tk.X, pady=5)
        
        Label(width_frame, text="Width:").pack(side=tk.LEFT)
        self.width_scale_var = tk.DoubleVar(value=1.2)
        self.width_scale = tk.Scale(width_frame, from_=0.8, to=3.0, resolution=0.1, orient=tk.HORIZONTAL, 
                                    variable=self.width_scale_var, command=self.update_params)
        self.width_scale.pack(side=tk.RIGHT, fill=tk.X, expand=True)
        
        # Height scale slider
        height_frame = Frame(right_frame)
        height_frame.pack(fill=tk.X, pady=5)
        
        Label(height_frame, text="Height:").pack(side=tk.LEFT)
        self.height_scale_var = tk.DoubleVar(value=1.1)
        self.height_scale = tk.Scale(height_frame, from_=0.6, to=2.0, resolution=0.1, orient=tk.HORIZONTAL, 
                                     variable=self.height_scale_var, command=self.update_params)
        self.height_scale.pack(side=tk.RIGHT, fill=tk.X, expand=True)
        
        # Collar position slider
        collar_frame = Frame(right_frame)
        collar_frame.pack(fill=tk.X, pady=5)
        
        Label(collar_frame, text="Collar:").pack(side=tk.LEFT)
        self.collar_var = tk.DoubleVar(value=0.2)
        self.collar_scale = tk.Scale(collar_frame, from_=0.05, to=0.3, resolution=0.01, orient=tk.HORIZONTAL, 
                                     variable=self.collar_var, command=self.update_params)
        self.collar_scale.pack(side=tk.RIGHT, fill=tk.X, expand=True)
        
        # Debug mode checkbox
        debug_frame = Frame(right_frame)
        debug_frame.pack(fill=tk.X, pady=5)
        
        self.debug_var = tk.BooleanVar(value=False)
        debug_check = tk.Checkbutton(debug_frame, text="Show Skeleton", variable=self.debug_var, 
                                    command=self.update_params)
        debug_check.pack(side=tk.LEFT)
        
        # Status label
        self.status_label = Label(right_frame, text="Ready")
        self.status_label.pack(pady=10)
        
        # Start main loop
        self.root.mainloop()
    
    def upload_garment(self):
        """Open file dialog to select a garment image"""
        filetypes = [
            ("Image files", "*.png *.jpg *.jpeg"),
            ("PNG files", "*.png"),
            ("JPEG files", "*.jpg *.jpeg"),
            ("All files", "*.*")
        ]
        
        self.garment_path = filedialog.askopenfilename(
            title="Select Garment Image",
            filetypes=filetypes
        )
        
        if self.garment_path:
            self.status_label.config(text=f"Garment selected: {os.path.basename(self.garment_path)}")
            self.load_preview()
    
    def load_preview(self):
        """Load and display the garment preview"""
        if not self.garment_path:
            return
        
        try:
            # Load image with PIL for preview
            pil_img = Image.open(self.garment_path)
            
            # Resize for preview (keep aspect ratio)
            preview_width = 250
            aspect_ratio = pil_img.width / pil_img.height
            preview_height = int(preview_width / aspect_ratio)
            
            # Limit height
            if preview_height > 300:
                preview_height = 300
                preview_width = int(preview_height * aspect_ratio)
            
            pil_img = pil_img.resize((preview_width, preview_height), Image.LANCZOS)
            
            # Convert to Tkinter format
            tk_img = ImageTk.PhotoImage(pil_img)
            
            # Update preview
            self.garment_preview.config(image=tk_img)
            self.garment_preview.image = tk_img  # Keep a reference
            
        except Exception as e:
            self.status_label.config(text=f"Error loading preview: {e}")
    
    def start_tryon(self):
        """Start the virtual try-on"""
        if not self.garment_path:
            self.status_label.config(text="Please select a garment first")
            return
        
        if self.running:
            self.status_label.config(text="Already running")
            return
        
        try:
            # Initialize the virtual try-on application
            self.app = AdvancedVirtualTryOn(
                self.garment_path, 
                self.webcam_index, 
                f"{self.width}x{self.height}"
            )
            
            # Set initial parameters
            self.update_params()
            
            # Start the webcam
            self.running = True
            self.status_label.config(text="Try-on started")
            
            # Start updating frames
            self.update_frame()
            
        except Exception as e:
            self.status_label.config(text=f"Error starting try-on: {e}")
    
    def update_params(self, *args):
        """Update the application parameters from UI controls"""
        if self.app:
            self.app.width_scale = self.width_scale_var.get()
            self.app.height_scale = self.height_scale_var.get()
            self.app.collar_position = self.collar_var.get()
            self.app.debug_mode = self.debug_var.get()
    
    def update_frame(self):
        """Update the webcam frame"""
        if not self.running or not self.app:
            return
        
        try:
            # Get frame from camera
            ret, frame = self.app.camera.read()
            
            if ret:
                # Process the frame
                processed = self.app.process_frame(frame)
                
                # Convert to PIL format
                pil_img = Image.fromarray(cv2.cvtColor(processed, cv2.COLOR_BGR2RGB))
                
                # Convert to Tkinter format
                tk_img = ImageTk.PhotoImage(image=pil_img)
                
                # Update image
                self.webcam_label.config(image=tk_img)
                self.webcam_label.image = tk_img  # Keep a reference
                
            # Schedule next update
            self.root.after(self.update_interval, self.update_frame)
            
        except Exception as e:
            self.status_label.config(text=f"Error updating frame: {e}")
            self.stop_tryon()
    
    def stop_tryon(self):
        """Stop the virtual try-on"""
        self.running = False
        
        if self.app:
            self.app.clean_up()
            self.app = None
            
        self.status_label.config(text="Try-on stopped")
    
    def on_close(self):
        """Handle window close event"""
        self.stop_tryon()
        if self.root:
            self.root.destroy()

def run_tkinter_app(webcam_index=0, resolution="640x480"):
    """Run the application with Tkinter UI"""
    ui = TkinterUI(webcam_index, resolution)
    ui.start()

def run_streamlit_app():
    """Run the application in Streamlit mode"""
    st.set_page_config(page_title="Virtual Try-On", page_icon="👕", layout="wide")
    
    st.title("Advanced Virtual Try-On")
    st.subheader("Try on clothing virtually using your webcam")
    
    # Sidebar for uploading garment and adjustments
    with st.sidebar:
        st.header("Settings")
        uploaded_garment = st.file_uploader("Upload a garment image (with transparent background)", type=["png", "jpg", "jpeg"])
        
        st.subheader("Fitting Adjustments")
        width_scale = st.slider("Width Scale", 0.8, 3.0, 1.2, 0.1)
        height_scale = st.slider("Height Scale", 0.6, 2.0, 1.1, 0.1)
        collar_position = st.slider("Collar Position", 0.05, 0.3, 0.2, 0.01)
        
        debug_mode = st.checkbox("Show Skeleton", value=False)
        
        st.info("For best results, use a garment image with a transparent background.")
    
    # Main content
    col1, col2 = st.columns(2)
    
    # If no garment uploaded, show sample garments
    if uploaded_garment is None:
        with col1:
            st.warning("Please upload a garment image to begin")
            st.write("No garment image uploaded. Here's how it works:")
            st.write("1. Upload a garment image with transparent background")
            st.write("2. Position yourself in front of the camera")
            st.write("3. Adjust the fit using the controls in the sidebar")
            
            # Example garment placeholder
            st.image("https://i.imgur.com/JFP0rja.png", caption="Example garment (upload your own)")
            
        with col2:
            # Camera placeholder
            st.write("Camera feed will appear here")
            placeholder = st.empty()
            placeholder.image(np.zeros((480, 640, 3), dtype=np.uint8), channels="BGR", caption="Webcam feed will appear here")
        
        return
    
    # Save uploaded garment to temp file
    temp_garment_file = None
    if uploaded_garment is not None:
        temp_garment_file = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
        temp_garment_file.write(uploaded_garment.getvalue())
        temp_garment_file.close()
    
    # Initialize the application with the uploaded garment
    try:
        app = AdvancedVirtualTryOn(temp_garment_file.name, 0, "640x480", streamlit_mode=True)
        
        # Set parameters from sliders
        app.width_scale = width_scale
        app.height_scale = height_scale
        app.collar_position = collar_position
        app.debug_mode = debug_mode
        
        # Show the garment
        with col1:
            st.subheader("Garment:")
            garment_display = cv2.cvtColor(app.garment, cv2.COLOR_BGRA2RGBA)
            st.image(garment_display, caption="Your uploaded garment")
        
        # Start webcam
        with col2:
            st.subheader("Virtual Try-On:")
            webcam_placeholder = st.empty()
            
            cap = cv2.VideoCapture(0)
            cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
            cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
            
            # Check if webcam opened successfully
            if not cap.isOpened():
                st.error("Could not open webcam. Please check your camera connection.")
                return
            
            stop_button = st.button("Stop")
            
            while not stop_button:
                success, frame = cap.read()
                if not success:
                    st.error("Failed to capture frame from webcam")
                    break
                
                # Process the frame
                processed_frame = app.process_frame(frame)
                
                # Convert BGR to RGB for Streamlit
                processed_frame_rgb = cv2.cvtColor(processed_frame, cv2.COLOR_BGR2RGB)
                
                # Display the processed frame
                webcam_placeholder.image(processed_frame_rgb, channels="RGB", caption="Live Try-On")
                
                # Check if stop button was pressed
                if stop_button:
                    break
                
                # Small sleep to reduce CPU usage
                time.sleep(0.01)
                
                # Recheck the stop button status
                stop_button = st.button("Stop")
            
            # Clean up
            cap.release()
            
    except Exception as e:
        st.error(f"Error: {e}")
        if 'app' in locals():
            app.clean_up()
    
    finally:
        # Remove temporary file
        if temp_garment_file:
            try:
                os.unlink(temp_garment_file.name)
            except:
                pass

def main():
    args = parse_args()
    try:
        # Check which mode to run in
        if args.streamlit:
            run_streamlit_app()
        elif args.tkinter or (not args.garment and not args.streamlit):
            # Use tkinter by default if no garment is specified
            run_tkinter_app(args.webcam, args.resolution)
        else:
            # Traditional command-line mode if garment is specified
            width, height = map(int, args.resolution.split('x'))
            app = AdvancedVirtualTryOn(args.garment, args.webcam, args.resolution)
            app.run()
    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    main()