app.py

# -*- 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)