Add garment segmentation application files

README.md

@@ -1,12 +1,31 @@
 ---
-title: Easel AI Engineering
-emoji: 👁
-colorFrom: blue
-colorTo: pink
+title: Garment-based Segmentation
+emoji: 👕
+colorFrom: green
+colorTo: blue
 sdk: gradio
-sdk_version: 5.23.0
+sdk_version: 3.19.0
 app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# Garment-based Segmentation with SegFormer and Fashion-CLIP
+
+This application uses AI models to segment specific clothing items in images by matching a garment to a person.
+
+## Features
+
+- **Targeted Garment Segmentation**: Segment only the specific garment type that matches a reference image
+- **Interactive Interface**: User-friendly interface with clear sections for person and garment images
+- **Visualization Options**: View original images, segmentation maps, and overlays
+
+## How to Use
+
+1. Upload an image of a person wearing clothes
+2. Upload a reference image of a garment (e.g., a t-shirt, pants, dress)
+3. Click "Process Images" to generate the targeted segmentation
+4. View the results in the gallery
+
+## Technical Details
+
+This application combines SegFormer for image segmentation with Fashion-CLIP for garment classification to create a system that can identify and segment specific garment types in images.

app.py

@@ -0,0 +1,138 @@
+import gradio as gr
+import os
+
+# Import from our modules
+from constants import class_names
+from segmentation import segment_image
+from utils import process_url, process_person_and_garment
+
+# Define fixed size for consistent image display
+FIXED_IMAGE_SIZE = (400, 400)
+
+def create_interface():
+    """Create the Gradio interface with improved image consistency"""
+    with gr.Blocks(title="Garment-based Segmentation") as demo:
+        gr.Markdown("""
+        # Garment-based Segmentation with SegFormer and Fashion-CLIP
+        
+        This application uses AI models to segment specific clothing items in images by matching a garment to a person.
+        """)
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                # Person image section
+                gr.Markdown("### Person Image")
+                person_image = gr.Image(
+                    type="pil", 
+                    label="Upload a person wearing clothes", 
+                    height=300,
+                    sources=["upload", "webcam", "clipboard"],
+                    elem_id="person-image-upload"
+                )
+                
+                # Garment image section
+                gr.Markdown("### Garment Image")
+                garment_image = gr.Image(
+                    type="pil", 
+                    label="Upload a garment to detect", 
+                    height=300,
+                    sources=["upload", "webcam", "clipboard"],
+                    elem_id="garment-image-upload"
+                )
+                
+                with gr.Row():
+                    show_original_dual = gr.Checkbox(label="Show Original", value=True)
+                    show_segmentation_dual = gr.Checkbox(label="Show Segmentation", value=True)
+                    show_overlay_dual = gr.Checkbox(label="Show Overlay", value=True)
+                
+                process_button = gr.Button(
+                    "Process Images", 
+                    variant="primary",
+                    size="lg",
+                    elem_id="process-button"
+                )
+            
+            with gr.Column(scale=2):
+                dual_output_images = gr.Gallery(
+                    label="Results", 
+                    columns=3, 
+                    height=450,
+                    object_fit="contain",
+                    elem_id="dual_gallery"
+                )
+                result_text = gr.Textbox(label="Result", interactive=False, lines=4)
+        
+        # Set up event handler for dual image processing
+        process_button.click(
+            fn=lambda p_img, g_img, orig, seg, over: process_person_and_garment(p_img, g_img, orig, seg, over, FIXED_IMAGE_SIZE),
+            inputs=[person_image, garment_image, show_original_dual, show_segmentation_dual, show_overlay_dual],
+            outputs=[dual_output_images, result_text]
+        )
+        
+        # Add custom CSS for consistent image sizes and improved UI
+        gr.HTML("""
+        <style>
+        .gradio-container img {
+            max-height: 400px !important;
+            object-fit: contain !important;
+        }
+        #dual_gallery {
+            min-height: 450px;
+        }
+        /* Larger upload buttons */
+        #person-image-upload .upload-button, 
+        #garment-image-upload .upload-button {
+            font-size: 1.2em !important;
+            padding: 12px 20px !important;
+            border-radius: 8px !important;
+            margin: 10px auto !important;
+            display: block !important;
+            width: 80% !important;
+            text-align: center !important;
+            background-color: #4CAF50 !important;
+            color: white !important;
+            border: none !important;
+            cursor: pointer !important;
+            transition: background-color 0.3s ease !important;
+        }
+        #person-image-upload .upload-button:hover, 
+        #garment-image-upload .upload-button:hover {
+            background-color: #45a049 !important;
+        }
+        /* Larger process button */
+        #process-button {
+            font-size: 1.3em !important;
+            padding: 15px 25px !important;
+            margin: 15px auto !important;
+            display: block !important;
+            width: 90% !important;
+        }
+        /* Better section headers */
+        h3 {
+            font-size: 1.5em !important;
+            margin-top: 20px !important;
+            margin-bottom: 15px !important;
+            color: #2c3e50 !important;
+            border-bottom: 2px solid #3498db !important;
+            padding-bottom: 8px !important;
+        }
+        /* Better main heading */
+        h1 {
+            color: #2c3e50 !important;
+            text-align: center !important;
+            margin-bottom: 30px !important;
+            font-size: 2.5em !important;
+        }
+        /* Better checkbox layout */
+        .gradio-checkbox {
+            margin: 10px 5px !important;
+        }
+        </style>
+        """)
+    
+    return demo
+
+# Main application entry point
+if __name__ == "__main__":
+    demo = create_interface()
+    demo.launch()

classification.py

@@ -0,0 +1,71 @@
+import torch
+from models import clip_model, clip_processor, segformer_model, segformer_processor
+from constants import fashion_categories, fashion_clip_to_segformer, class_names, category_to_segment_mapping, garment_to_segments
+from segmentation import identify_garment_segformer
+
+def identify_garment_clip(image):
+    """Identify the garment type using Fashion-CLIP model"""
+    # Prepare text prompts
+    texts = [f"a photo of a {category}" for category in fashion_categories]
+    
+    # Process inputs
+    inputs = clip_processor(text=texts, images=image, return_tensors="pt", padding=True)
+    
+    # Get predictions
+    with torch.no_grad():
+        outputs = clip_model(**inputs)
+        logits_per_image = outputs.logits_per_image
+        probs = logits_per_image.softmax(dim=1)
+    
+    # Get the top prediction
+    top_idx = torch.argmax(probs[0]).item()
+    top_category = fashion_categories[top_idx]
+    confidence = probs[0][top_idx].item() * 100
+    
+    # Map to SegFormer class if possible
+    if top_category in fashion_clip_to_segformer:
+        segformer_idx = fashion_clip_to_segformer[top_category]
+        segformer_class = class_names[segformer_idx]
+        return top_category, segformer_idx, confidence
+    else:
+        # Fallback to using SegFormer directly
+        return top_category, None, confidence
+
+def get_segments_for_garment(garment_image):
+    """Get the segments that should be included for a given garment image"""
+    # First try to identify the garment using Fashion-CLIP
+    clip_category, segformer_idx, confidence = identify_garment_clip(garment_image)
+    
+    # If CLIP couldn't map to a SegFormer class, fall back to SegFormer
+    if segformer_idx is None:
+        garment_name, segformer_idx = identify_garment_segformer(garment_image)
+        method = "SegFormer (fallback)"
+        confidence_text = ""
+    else:
+        garment_name = class_names[segformer_idx]
+        method = "Fashion-CLIP"
+        confidence_text = f" with {confidence:.2f}% confidence"
+    
+    if segformer_idx is None:
+        return None, None, "No clear garment detected in the garment image"
+    
+    # Get all segments that should be included based on the detected garment
+    if method == "Fashion-CLIP" and clip_category in category_to_segment_mapping:
+        # Use the detailed mapping for CLIP categories
+        selected_class = category_to_segment_mapping[clip_category]
+    elif segformer_idx in garment_to_segments:
+        # Fall back to the SegFormer class-based mapping
+        segment_indices = garment_to_segments[segformer_idx]
+        selected_class = [class_names[idx] for idx in segment_indices]
+    else:
+        # Fallback to just the detected garment if no mapping exists
+        selected_class = [class_names[segformer_idx]]
+    
+    # Prepare a more descriptive result text
+    included_segments = ", ".join(selected_class)
+    if method == "Fashion-CLIP":
+        result_text = f"Detected garment: {clip_category} (mapped to {garment_name})\nUsing {method}{confidence_text}\nSegmented parts: {included_segments}"
+    else:
+        result_text = f"Detected garment: {garment_name}\nUsing {method}\nSegmented parts: {included_segments}"
+    
+    return selected_class, segformer_idx, result_text

constants.py

@@ -0,0 +1,168 @@
+import matplotlib.pyplot as plt
+
+# Define the class names for the segmentation model
+class_names = [
+    "Background", "Hat", "Hair", "Sunglasses", "Upper-clothes", "Skirt", "Pants", 
+    "Dress", "Belt", "Left-shoe", "Right-shoe", "Face", "Left-leg", "Right-leg", 
+    "Left-arm", "Right-arm", "Bag", "Scarf"
+]
+
+# Define a color map for visualization
+color_map = plt.cm.get_cmap('tab20', len(class_names))
+
+# Define a mapping of garment types to related segments that should be included
+garment_to_segments = {
+    0: [0],                   # Background --> segment background only
+    1: [1, 2, 11],            # Hat --> segment hat, hair, and face
+    2: [2],                   # Hair --> segment hair only
+    3: [3, 11],               # Sunglasses --> segment sunglasses and face
+    4: [4, 14, 15],           # Upper-clothes --> segment upper clothes, left arm, right arm
+    5: [5, 6, 12, 13],        # Skirt --> segment skirt, pants, left leg, right leg
+    6: [6, 12, 13],           # Pants --> segment pants, left leg, right leg
+    7: [4, 5, 6, 7, 12, 13, 14, 15],  # Dress --> segment whole body except face and hair
+    8: [8],                   # Belt --> segment belt only
+    9: [9],                   # Left-shoe --> segment left shoe only
+    10: [10],                 # Right-shoe --> segment right shoe only
+    11: [11],                 # Face --> segment face only
+    12: [12],                 # Left-leg --> segment left leg only
+    13: [13],                 # Right-leg --> segment right leg only
+    14: [14],                 # Left-arm --> segment left arm only
+    15: [15],                 # Right-arm --> segment right arm only
+    16: [16],                 # Bag --> segment bag only
+    17: [17, 2, 11]           # Scarf --> segment scarf, hair and face
+}
+
+# Define categories for Fashion-CLIP
+fashion_categories = [
+    # Upper body
+    "t-shirt", "shirt", "blouse", "tank top", "polo shirt", "sweatshirt", "hoodie",
+    
+    # Outerwear
+    "jacket", "coat", "blazer", "cardigan", "vest", "windbreaker",
+    
+    # Dresses
+    "dress", "shirt dress", "sundress", "evening gown", "maxi dress", "mini dress",
+    
+    # Lower body
+    "jeans", "pants", "trousers", "shorts", "skirt", "leggings", "joggers", "sweatpants",
+    
+    # Formal wear
+    "suit", "tuxedo", "formal shirt", "formal dress",
+    
+    # Undergarments
+    "bra", "underwear", "boxers", "briefs", "lingerie",
+    
+    # Sleepwear
+    "pajamas", "nightgown", "bathrobe",
+    
+    # Swimwear
+    "swimsuit", "bikini", "swim trunks",
+    
+    # Footwear
+    "shoes", "boots", "sneakers", "sandals", "high heels", "loafers", "flats",
+    
+    # Accessories
+    "hat", "cap", "scarf", "gloves", "belt", "tie", "socks",
+    
+    # Bags
+    "handbag", "backpack", "purse", "tote bag"
+]
+
+# Mapping from Fashion-CLIP categories to SegFormer classes
+fashion_clip_to_segformer = {
+    # Upper body items -> Upper-clothes (4)
+    "t-shirt": 4, "shirt": 4, "blouse": 4, "tank top": 4, "polo shirt": 4, "sweatshirt": 4, "hoodie": 4,
+    "cardigan": 4, "vest": 4, "formal shirt": 4,
+    
+    # Outerwear -> Upper-clothes (4)
+    "jacket": 4, "coat": 4, "blazer": 4, "windbreaker": 4,
+    
+    # Dresses -> Dress (7)
+    "dress": 7, "shirt dress": 7, "sundress": 7, "evening gown": 7, "maxi dress": 7, "mini dress": 7,
+    "formal dress": 7,
+    
+    # Lower body -> Pants (6) or Skirt (5)
+    "jeans": 6, "pants": 6, "trousers": 6, "shorts": 6, "leggings": 6, "joggers": 6, "sweatpants": 6,
+    "skirt": 5,
+    
+    # Formal wear -> Upper-clothes (4) or Dress (7)
+    "suit": 4, "tuxedo": 4,
+    
+    # Footwear -> Left-shoe/Right-shoe (9/10)
+    "shoes": 9, "boots": 9, "sneakers": 9, "sandals": 9, "high heels": 9, "loafers": 9, "flats": 9,
+    
+    # Accessories
+    "hat": 1, "cap": 1, "scarf": 17, "belt": 8,
+    
+    # Bags
+    "handbag": 16, "backpack": 16, "purse": 16, "tote bag": 16
+}
+
+# Detailed mapping from categories to segment names
+category_to_segment_mapping = {
+    # Upper body items map to Upper-clothes and arms
+    "t-shirt": ["Upper-clothes", "Left-arm", "Right-arm"],
+    "shirt": ["Upper-clothes", "Left-arm", "Right-arm"],
+    "blouse": ["Upper-clothes", "Left-arm", "Right-arm"],
+    "tank top": ["Upper-clothes", "Left-arm", "Right-arm"],
+    "polo shirt": ["Upper-clothes", "Left-arm", "Right-arm"],
+    "sweatshirt": ["Upper-clothes", "Left-arm", "Right-arm"],
+    "hoodie": ["Upper-clothes", "Left-arm", "Right-arm"],
+    
+    # Outerwear maps to Upper-clothes and arms
+    "jacket": ["Upper-clothes", "Left-arm", "Right-arm"],
+    "coat": ["Upper-clothes", "Left-arm", "Right-arm"],
+    "blazer": ["Upper-clothes", "Left-arm", "Right-arm"],
+    "cardigan": ["Upper-clothes", "Left-arm", "Right-arm"],
+    "vest": ["Upper-clothes"],
+    "windbreaker": ["Upper-clothes", "Left-arm", "Right-arm"],
+    
+    # Dresses map to Dress
+    "dress": ["Dress", "Upper-clothes", "Skirt", "Pants", "Left-arm", "Right-arm", "Left-leg", "Right-leg"],
+    "shirt dress": ["Dress", "Upper-clothes", "Skirt", "Pants", "Left-arm", "Right-arm", "Left-leg", "Right-leg"],
+    "sundress": ["Dress", "Upper-clothes", "Skirt", "Pants", "Left-arm", "Right-arm", "Left-leg", "Right-leg"],
+    "evening gown": ["Dress", "Upper-clothes", "Skirt", "Pants", "Left-arm", "Right-arm", "Left-leg", "Right-leg"],
+    "maxi dress": ["Dress", "Upper-clothes", "Skirt", "Pants", "Left-arm", "Right-arm", "Left-leg", "Right-leg"],
+    "mini dress": ["Dress", "Upper-clothes", "Skirt", "Pants", "Left-arm", "Right-arm", "Left-leg", "Right-leg"],
+    "formal dress": ["Dress", "Upper-clothes", "Skirt", "Pants", "Left-arm", "Right-arm", "Left-leg", "Right-leg"],
+    
+    # Lower body items map to Pants or Skirt and legs
+    "jeans": ["Pants", "Left-leg", "Right-leg"],
+    "pants": ["Pants", "Left-leg", "Right-leg"],
+    "trousers": ["Pants", "Left-leg", "Right-leg"],
+    "shorts": ["Pants", "Left-leg", "Right-leg"],
+    "skirt": ["Skirt", "Pants", "Left-leg", "Right-leg"],
+    "leggings": ["Pants", "Left-leg", "Right-leg"],
+    "joggers": ["Pants", "Left-leg", "Right-leg"],
+    "sweatpants": ["Pants", "Left-leg", "Right-leg"],
+    
+    # Formal wear maps depending on type
+    "suit": ["Upper-clothes", "Left-arm", "Right-arm", "Pants", "Left-leg", "Right-leg"],
+    "tuxedo": ["Upper-clothes", "Left-arm", "Right-arm", "Pants", "Left-leg", "Right-leg"],
+    "formal shirt": ["Upper-clothes", "Left-arm", "Right-arm"],
+    "formal dress": ["Dress", "Upper-clothes", "Skirt", "Pants", "Left-arm", "Right-arm", "Left-leg", "Right-leg"],
+    
+    # Footwear maps to shoes
+    "shoes": ["Left-shoe", "Right-shoe"],
+    "boots": ["Left-shoe", "Right-shoe"],
+    "sneakers": ["Left-shoe", "Right-shoe"],
+    "sandals": ["Left-shoe", "Right-shoe"],
+    "high heels": ["Left-shoe", "Right-shoe"],
+    "loafers": ["Left-shoe", "Right-shoe"],
+    "flats": ["Left-shoe", "Right-shoe"],
+    
+    # Accessories map to their respective parts
+    "hat": ["Hat"],
+    "cap": ["Hat"],
+    "scarf": ["Scarf", "Face", "Hair"],
+    "gloves": ["Left-arm", "Right-arm"],
+    "belt": ["Belt"],
+    "tie": ["Upper-clothes"],
+    "socks": ["Left-leg", "Right-leg"],
+    
+    # Bags map to Bag
+    "handbag": ["Bag"],
+    "backpack": ["Bag"],
+    "purse": ["Bag"],
+    "tote bag": ["Bag"]
+}

models.py

@@ -0,0 +1,21 @@
+import torch
+import torch.nn as nn
+from transformers import SegformerImageProcessor, AutoModelForSemanticSegmentation, CLIPProcessor, CLIPModel
+
+# Load the SegFormer model and processor for segmentation
+def load_segformer_model():
+    """Load and return the SegFormer model and processor"""
+    processor = SegformerImageProcessor.from_pretrained("mattmdjaga/segformer_b2_clothes")
+    model = AutoModelForSemanticSegmentation.from_pretrained("mattmdjaga/segformer_b2_clothes")
+    return processor, model
+
+# Load Fashion-CLIP model for garment classification
+def load_clip_model():
+    """Load and return the Fashion-CLIP model and processor"""
+    model = CLIPModel.from_pretrained("patrickjohncyh/fashion-clip")
+    processor = CLIPProcessor.from_pretrained("patrickjohncyh/fashion-clip")
+    return model, processor
+
+# Initialize models
+segformer_processor, segformer_model = load_segformer_model()
+clip_model, clip_processor = load_clip_model()

requirements.txt

@@ -0,0 +1,8 @@
+torch>=1.10.0
+torchvision>=0.11.1
+transformers>=4.15.0
+Pillow>=8.4.0
+requests>=2.26.0
+matplotlib>=3.5.0
+gradio>=3.19.0
+numpy>=1.21.0

segmentation.py

@@ -0,0 +1,150 @@
+import torch
+import torch.nn as nn
+import numpy as np
+import matplotlib.pyplot as plt
+from PIL import Image
+import io
+from collections import Counter
+import gradio as gr
+
+from models import segformer_model, segformer_processor
+from constants import class_names, color_map
+
+def segment_image(image, selected_classes=None, show_original=True, show_segmentation=True, show_overlay=True, fixed_size=(400, 400)):
+    """Segment the image based on selected classes with consistent output sizes"""
+    # Process the image
+    inputs = segformer_processor(images=image, return_tensors="pt")
+    
+    # Get model predictions
+    outputs = segformer_model(**inputs)
+    logits = outputs.logits.cpu()
+    
+    # Upsample the logits to match the original image size
+    upsampled_logits = nn.functional.interpolate(
+        logits,
+        size=image.size[::-1],  # (height, width)
+        mode="bilinear",
+        align_corners=False,
+    )
+    
+    # Get the predicted segmentation map
+    pred_seg = upsampled_logits.argmax(dim=1)[0].numpy()
+    
+    # Filter classes if specified
+    if selected_classes and len(selected_classes) > 0:
+        # Create a mask for selected classes
+        mask = np.zeros_like(pred_seg, dtype=bool)
+        for class_name in selected_classes:
+            if class_name in class_names:
+                class_idx = class_names.index(class_name)
+                mask = np.logical_or(mask, pred_seg == class_idx)
+        
+        # Apply the mask to keep only selected classes, set others to background (0)
+        filtered_seg = np.zeros_like(pred_seg)
+        filtered_seg[mask] = pred_seg[mask]
+        pred_seg = filtered_seg
+    
+    # Create a colored segmentation map
+    colored_seg = np.zeros((pred_seg.shape[0], pred_seg.shape[1], 3))
+    for class_idx in range(len(class_names)):
+        mask = pred_seg == class_idx
+        if mask.any():
+            colored_seg[mask] = color_map(class_idx)[:3]
+    
+    # Create an overlay of the segmentation on the original image
+    image_array = np.array(image)
+    overlay = image_array.copy()
+    alpha = 0.5  # Transparency factor
+    mask = pred_seg > 0  # Exclude background
+    overlay[mask] = overlay[mask] * (1 - alpha) + colored_seg[mask] * 255 * alpha
+    
+    # Prepare output images based on user selection
+    outputs = []
+    
+    if show_original:
+        # Resize original image to ensure consistent size
+        resized_original = image.resize(fixed_size)
+        outputs.append(resized_original)
+    
+    if show_segmentation:
+        seg_image = Image.fromarray((colored_seg * 255).astype('uint8'))
+        # Ensure segmentation has consistent size
+        seg_image = seg_image.resize(fixed_size)
+        outputs.append(seg_image)
+    
+    if show_overlay:
+        overlay_image = Image.fromarray(overlay.astype('uint8'))
+        # Ensure overlay has consistent size
+        overlay_image = overlay_image.resize(fixed_size)
+        outputs.append(overlay_image)
+    
+    # Create a legend for the segmentation classes
+    fig, ax = plt.subplots(figsize=(10, 2))
+    fig.patch.set_alpha(0.0)
+    ax.axis('off')
+    
+    # Create legend patches
+    legend_elements = []
+    for i, class_name in enumerate(class_names):
+        if i == 0 and selected_classes:  # Skip background in legend if filtering
+            continue
+        if not selected_classes or class_name in selected_classes:
+            color = color_map(i)[:3]
+            legend_elements.append(plt.Rectangle((0, 0), 1, 1, color=color))
+    
+    # Only add legend if there are elements to show
+    if legend_elements:
+        legend_class_names = [name for name in class_names if name != "Background" and (not selected_classes or name in selected_classes)]
+        ax.legend(legend_elements, legend_class_names, loc='center', ncol=6)
+    
+    # Save the legend to a bytes buffer
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png', bbox_inches='tight', transparent=True)
+    buf.seek(0)
+    legend_img = Image.open(buf)
+    
+    plt.close(fig)
+    
+    outputs.append(legend_img)
+    
+    return outputs
+
+def identify_garment_segformer(image):
+    """Identify the dominant garment type using SegFormer"""
+    # Process the image
+    inputs = segformer_processor(images=image, return_tensors="pt")
+    
+    # Get model predictions
+    outputs = segformer_model(**inputs)
+    logits = outputs.logits.cpu()
+    
+    # Upsample the logits to match the original image size
+    upsampled_logits = nn.functional.interpolate(
+        logits,
+        size=image.size[::-1],  # (height, width)
+        mode="bilinear",
+        align_corners=False,
+    )
+    
+    # Get the predicted segmentation map
+    pred_seg = upsampled_logits.argmax(dim=1)[0].numpy()
+    
+    # Count the pixels for each class (excluding background)
+    class_counts = Counter(pred_seg.flatten())
+    if 0 in class_counts:  # Remove background
+        del class_counts[0]
+    
+    # Find the most common clothing item
+    clothing_classes = [4, 5, 6, 7]  # Upper-clothes, Skirt, Pants, Dress
+    
+    # Filter to only include clothing items
+    clothing_counts = {k: v for k, v in class_counts.items() if k in clothing_classes}
+    
+    if not clothing_counts:
+        return "No garment detected", None
+    
+    # Get the most common clothing item
+    dominant_class = max(clothing_counts.items(), key=lambda x: x[1])[0]
+    dominant_class_name = class_names[dominant_class]
+    
+    return dominant_class_name, dominant_class

utils.py

@@ -0,0 +1,34 @@
+import requests
+from PIL import Image
+import gradio as gr
+
+from segmentation import segment_image
+from classification import get_segments_for_garment
+
+def process_url(url, selected_classes, show_original, show_segmentation, show_overlay, fixed_size=(400, 400)):
+    """Process an image from a URL"""
+    try:
+        image = Image.open(requests.get(url, stream=True).raw)
+        return segment_image(image, selected_classes, show_original, show_segmentation, show_overlay, fixed_size)
+    except Exception as e:
+        return [gr.update(value=None)] * 4, f"Error: {str(e)}"
+
+def process_person_and_garment(person_image, garment_image, show_original, show_segmentation, show_overlay, fixed_size=(400, 400)):
+    """Process person and garment images for targeted segmentation"""
+    if person_image is None or garment_image is None:
+        return [gr.update(value=None)] * 4, "Please provide both person and garment images"
+    
+    try:
+        # Get segments that should be included based on the garment
+        selected_class, segformer_idx, result_text = get_segments_for_garment(garment_image)
+        
+        if selected_class is None:
+            return [gr.update(value=None)] * 4, result_text
+        
+        # Process the person image with the selected garment classes
+        result_images = segment_image(person_image, selected_class, show_original, show_segmentation, show_overlay, fixed_size)
+        
+        return result_images, result_text
+    
+    except Exception as e:
+        return [gr.update(value=None)] * 4, f"Error: {str(e)}"