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()