Create app.py

app.py

@@ -0,0 +1,316 @@
+# -*- 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)