# -*- coding: utf-8 -*- """VTON_GarmentMasker.ipynb Automatically generated by Colab. Original file is located at https://colab.research.google.com/drive/1Y22abu3jZQ5qCKP7DTR6kYvXdQbHnJCu Using YOLO Clothing Classification Model """ # !pip install gradio # !pip install ultralytics # !pip install segment-anything # !wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth import torch import numpy as np import cv2 from PIL import Image from torchvision import transforms from ultralytics import YOLO from segment_anything import SamPredictor, sam_model_registry from transformers import YolosForObjectDetection, YolosImageProcessor import gradio as gr import os import urllib.request class GarmentMaskingPipeline: def __init__(self): self.device = "cuda" if torch.cuda.is_available() else "cpu" print(f"Using device: {self.device}") self.yolo_model, self.sam_predictor, self.classification_model = self.load_models() self.clothing_to_body_parts = { 'shirt': ['torso', 'arms'], 't-shirt': ['torso', 'upper_arms'], 'blouse': ['torso', 'arms'], 'dress': ['torso', 'legs'], 'skirt': ['lower_torso', 'legs'], 'pants': ['legs'], 'shorts': ['upper_legs'], 'jacket': ['torso', 'arms'], 'coat': ['torso', 'arms'] } self.body_parts_positions = { 'face': (0.0, 0.2), 'torso': (0.2, 0.5), 'arms': (0.2, 0.5), 'upper_arms': (0.2, 0.35), 'lower_torso': (0.4, 0.6), 'legs': (0.5, 0.9), 'upper_legs': (0.5, 0.7), 'feet': (0.9, 1.0) } def load_models(self): print("Loading models...") # Download models if they don't exist self.download_models() # Load YOLO model yolo_model = YOLO('yolov8n.pt') # Load SAM model sam = sam_model_registry["vit_h"](checkpoint="sam_vit_h_4b8939.pth") sam.to(self.device) predictor = SamPredictor(sam) # Load YOLOS-Fashionpedia model for clothing classification print("Loading YOLOS-Fashionpedia model...") model_name = "valentinafeve/yolos-fashionpedia" processor = YolosImageProcessor.from_pretrained(model_name) classification_model = YolosForObjectDetection.from_pretrained(model_name) classification_model.to(self.device) classification_model.eval() print("Models loaded successfully!") return yolo_model, predictor, classification_model def download_models(self): """Download required model files if they don't exist""" models = { "yolov8n.pt": "https://github.com/ultralytics/assets/releases/download/v0.0.0/yolov8n.pt", "sam_vit_h_4b8939.pth": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth" } for filename, url in models.items(): if not os.path.exists(filename): print(f"Downloading {filename}...") urllib.request.urlretrieve(url, filename) print(f"Downloaded {filename}") else: print(f"{filename} already exists") # The YOLOS-Fashionpedia model will be downloaded automatically by transformers def classify_clothing(self, clothing_image): if not isinstance(clothing_image, Image.Image): clothing_image = Image.fromarray(clothing_image) # Process image with YOLOS processor processor = YolosImageProcessor.from_pretrained("valentinafeve/yolos-fashionpedia") inputs = processor(images=clothing_image, return_tensors="pt").to(self.device) # Run inference with torch.no_grad(): outputs = self.classification_model(**inputs) # Process results target_sizes = torch.tensor([clothing_image.size[::-1]]).to(self.device) results = processor.post_process_object_detection( outputs, target_sizes=target_sizes, threshold=0.1 )[0] # Extract detected labels and confidence scores labels = results["labels"] scores = results["scores"] # Get class names from model config id2label = self.classification_model.config.id2label # Define Fashionpedia to our category mapping fashionpedia_to_clothing = { 'shirt': 'shirt', 'blouse': 'shirt', 'top': 't-shirt', 't-shirt': 't-shirt', 'sweater': 'shirt', 'jacket': 'jacket', 'cardigan': 'jacket', 'coat': 'coat', 'jumper': 'shirt', 'dress': 'dress', 'skirt': 'skirt', 'shorts': 'shorts', 'pants': 'pants', 'jeans': 'pants', 'leggings': 'pants', 'jumpsuit': 'dress' } # Find the garment with highest confidence if len(labels) > 0: detections = [(id2label[label.item()].lower(), score.item()) for label, score in zip(labels, scores)] detections.sort(key=lambda x: x[1], reverse=True) for label, score in detections: # Look for clothing keywords in the label for keyword, category in fashionpedia_to_clothing.items(): if keyword in label: return category # If no mapping found, use the first detection as is return 't-shirt' # Default to t-shirt if nothing detected return 't-shirt' def create_garment_mask(self, person_image, garment_image): clothing_type = self.classify_clothing(garment_image) parts_to_mask = self.clothing_to_body_parts.get(clothing_type, []) results = self.yolo_model(person_image, classes=[0]) mask = np.zeros(person_image.shape[:2], dtype=np.uint8) if results and len(results[0].boxes.data) > 0: person_boxes = results[0].boxes.data person_areas = [(box[2] - box[0]) * (box[3] - box[1]) for box in person_boxes] largest_person_index = np.argmax(person_areas) person_box = person_boxes[largest_person_index][:4].cpu().numpy().astype(int) self.sam_predictor.set_image(person_image) masks, _, _ = self.sam_predictor.predict(box=person_box, multimask_output=False) person_mask = masks[0].astype(np.uint8) h, w = person_mask.shape for part in parts_to_mask: if part in self.body_parts_positions: top_ratio, bottom_ratio = self.body_parts_positions[part] top_px, bottom_px = int(h * top_ratio), int(h * bottom_ratio) part_mask = np.zeros_like(person_mask) part_mask[top_px:bottom_px, :] = 1 part_mask = np.logical_and(part_mask, person_mask).astype(np.uint8) mask = np.logical_or(mask, part_mask).astype(np.uint8) # Remove face from the mask face_top_px, face_bottom_px = int(h * 0.0), int(h * 0.2) face_mask = np.zeros_like(person_mask) face_mask[face_top_px:face_bottom_px, :] = 1 face_mask = np.logical_and(face_mask, person_mask).astype(np.uint8) mask = np.logical_and(mask, np.logical_not(face_mask)).astype(np.uint8) # Remove feet from the mask feet_top_px, feet_bottom_px = int(h * 0.9), int(h * 1.0) feet_mask = np.zeros_like(person_mask) feet_mask[feet_top_px:feet_bottom_px, :] = 1 feet_mask = np.logical_and(feet_mask, person_mask).astype(np.uint8) mask = np.logical_and(mask, np.logical_not(feet_mask)).astype(np.uint8) return mask * 255 def process(self, person_image_pil, garment_image_pil, mask_color_hex="#00FF00", opacity=0.5): """Process the input images and return the masked result""" # Convert PIL to numpy array person_image = np.array(person_image_pil) garment_image = np.array(garment_image_pil) # Convert to RGB if needed if person_image.shape[2] == 4: # RGBA person_image = person_image[:, :, :3] if garment_image.shape[2] == 4: # RGBA garment_image = garment_image[:, :, :3] # Create garment mask garment_mask = self.create_garment_mask(person_image, garment_image) # Convert hex color to RGB r = int(mask_color_hex[1:3], 16) g = int(mask_color_hex[3:5], 16) b = int(mask_color_hex[5:7], 16) color = (r, g, b) # Create a colored mask colored_mask = np.zeros_like(person_image) for i in range(3): colored_mask[:, :, i] = garment_mask * (color[i] / 255.0) # Create binary mask for visualization binary_mask = np.stack([garment_mask, garment_mask, garment_mask], axis=2) # Overlay mask on original image mask_3d = garment_mask[:, :, np.newaxis] / 255.0 overlay = person_image * (1 - opacity * mask_3d) + colored_mask * opacity overlay = overlay.astype(np.uint8) # Get classification result clothing_type = self.classify_clothing(garment_image) parts_to_mask = self.clothing_to_body_parts.get(clothing_type, []) return overlay, binary_mask, f"Detected garment: {clothing_type}\nBody parts to mask: {', '.join(parts_to_mask)}" def process_images(person_img, garment_img, mask_color, opacity): """Gradio processing function""" try: pipeline = GarmentMaskingPipeline() result = pipeline.process(person_img, garment_img, mask_color, opacity) return result except Exception as e: import traceback error_msg = f"Error processing images: {str(e)}\n{traceback.format_exc()}" print(error_msg) return None, None, error_msg def create_gradio_interface(): """Create and launch the Gradio interface""" with gr.Blocks(title="VTON SAM Garment Masking Pipeline") as interface: gr.Markdown(""" # Virtual Try-On Garment Masking Pipeline with SAM and YOLOS-Fashionpedia Upload a person image and a garment image to generate a mask for a virtual try-on application. The system will: 1. Detect the person using YOLO 2. Create a high-quality segmentation using SAM (Segment Anything Model) 3. Classify the garment type using YOLOS-Fashionpedia 4. Generate a mask of the area where the garment should be placed **Note**: This system uses state-of-the-art AI segmentation and fashion detection models for accurate results. """) with gr.Row(): with gr.Column(): person_input = gr.Image(label="Person Image (Image A)", type="pil") garment_input = gr.Image(label="Garment Image (Image B)", type="pil") with gr.Row(): mask_color = gr.ColorPicker(label="Mask Color", value="#00FF00") opacity = gr.Slider(label="Mask Opacity", minimum=0.1, maximum=0.9, value=0.5, step=0.1) submit_btn = gr.Button("Generate Mask") with gr.Column(): masked_output = gr.Image(label="Person with Masked Region") mask_output = gr.Image(label="Standalone Mask") result_text = gr.Textbox(label="Detection Results", lines=3) # Set up the processing flow submit_btn.click( fn=process_images, inputs=[person_input, garment_input, mask_color, opacity], outputs=[masked_output, mask_output, result_text] ) gr.Markdown(""" ## How It Works 1. **Person Detection**: Uses YOLO to detect and locate the person in the image 2. **Segmentation**: Uses SAM (Segment Anything Model) to create a high-quality segmentation mask 3. **Garment Classification**: Uses YOLOS-Fashionpedia to identify the garment type with fashion-specific detection 4. **Mask Generation**: Creates a mask based on the garment type and body part mapping ## Supported Garment Types - Shirts, Blouses, Tops, and T-shirts - Sweaters and Cardigans - Dresses and Jumpsuits - Skirts - Pants, Jeans, and Leggings - Shorts - Jackets and Coats """) return interface if __name__ == "__main__": # Create and launch the Gradio interface interface = create_gradio_interface() interface.launch(debug=True,share=True)