example_usage.py

"""
Example usage of WearIT Garment Mask Pipeline
"""
import os
from pathlib import Path
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np

# Option 1: Using Hugging Face transformers (if available)
try:
    from transformers import pipeline
    USE_HF_PIPELINE = True
except ImportError:
    USE_HF_PIPELINE = False
    print("Transformers not available, using direct import")

# Option 2: Direct import (always works)
from pipeline import GarmentMaskPipeline


def visualize_results(results, save_dir="./visualization"):
    """
    Visualize and save the mask generation results.

    Args:
        results: List of result dictionaries from the pipeline
        save_dir: Directory to save visualizations
    """
    save_path = Path(save_dir)
    save_path.mkdir(parents=True, exist_ok=True)

    for result in results:
        image_id = result["image_id"]
        image = result["image_standardized"]
        masks = result["masks"]

        # Create a figure for each image
        n_masks = len(masks)
        fig, axes = plt.subplots(1, n_masks + 1, figsize=(5 * (n_masks + 1), 5))

        if n_masks == 1:
            axes = [axes] if not isinstance(axes, np.ndarray) else axes

        # Show original image
        axes[0].imshow(image)
        axes[0].set_title(f"Original\n{image_id}")
        axes[0].axis("off")

        # Show masks
        for idx, (garment_type, mask_data) in enumerate(masks.items(), 1):
            mask = mask_data["person_mask"]

            # Overlay mask on image
            img_array = np.array(image)
            mask_array = np.array(mask)

            # Create colored overlay
            overlay = img_array.copy()
            overlay[mask_array > 127] = [255, 0, 0]  # Red for mask

            # Blend
            blended = (0.6 * img_array + 0.4 * overlay).astype(np.uint8)

            axes[idx].imshow(blended)
            axes[idx].set_title(f"{garment_type.capitalize()} Mask")
            axes[idx].axis("off")

            # Save individual mask
            mask.save(save_path / f"{image_id}_{garment_type}_mask.png")

        plt.tight_layout()
        plt.savefig(save_path / f"{image_id}_visualization.png", dpi=150, bbox_inches="tight")
        plt.close()

        # Save standardized image
        image.save(save_path / f"{image_id}_standardized.png")

    print(f"✓ Visualizations saved to {save_dir}")


def example_1_basic_usage():
    """Example 1: Basic single image processing"""
    print("\n" + "="*60)
    print("Example 1: Basic Usage - Single Image")
    print("="*60)

    if USE_HF_PIPELINE:
        # Using Hugging Face pipeline
        pipe = pipeline(
            "image-segmentation",
            model=".",  # Local directory
            trust_remote_code=True,
            device="cuda:0"
        )
    else:
        # Direct instantiation
        pipe = GarmentMaskPipeline(
            device="cuda:0",
            densepose_ckpt="chkpt/DensePose",
            schp_atr_ckpt="chkpt/SCHP/exp-schp-201908301523-atr.pth",
            schp_lip_ckpt="chkpt/SCHP/exp-schp-201908261155-lip.pth"
        )

    # Process single image
    results = pipe(
        "path/to/your/image.jpg",
        garment_types="upper"
    )

    # Access results
    result = results[0]
    upper_mask = result["masks"]["upper"]["person_mask"]
    upper_mask.save("output_upper_mask.png")

    print(f"✓ Generated mask for image: {result['image_id']}")
    print(f"  - Standardized image size: {result['image_standardized'].size}")
    print(f"  - Mask size: {upper_mask.size}")

    return results


def example_2_multiple_garments():
    """Example 2: Generate multiple garment masks per image"""
    print("\n" + "="*60)
    print("Example 2: Multiple Garment Types")
    print("="*60)

    pipe = GarmentMaskPipeline(device="cuda:0")

    # Generate masks for upper, lower, and full body
    results = pipe(
        "path/to/your/image.jpg",
        garment_types=["upper", "lower", "dress"]
    )

    result = results[0]
    print(f"✓ Generated {len(result['masks'])} masks:")
    for garment_type in result['masks'].keys():
        print(f"  - {garment_type}")

    return results


def example_3_batch_processing():
    """Example 3: Batch processing multiple images"""
    print("\n" + "="*60)
    print("Example 3: Batch Processing")
    print("="*60)

    pipe = GarmentMaskPipeline(
        device="cuda:0",
        save_images=True
    )

    # Process multiple images at once
    image_paths = [
        "path/to/image1.jpg",
        "path/to/image2.jpg",
        "path/to/image3.jpg"
    ]

    results = pipe(
        image_paths,
        garment_types=["upper", "lower"],
        image_ids=["person_001", "person_002", "person_003"],
        output_dir="./batch_output"
    )

    print(f"✓ Processed {len(results)} images")
    for result in results:
        print(f"  - {result['image_id']}: {list(result['masks'].keys())}")

    return results


def example_4_custom_configuration():
    """Example 4: Custom configuration"""
    print("\n" + "="*60)
    print("Example 4: Custom Configuration")
    print("="*60)

    # Create pipeline with custom settings
    pipe = GarmentMaskPipeline(
        device="cuda:0",
        output_height=2048,  # Higher resolution output
        process_size=768,     # Larger processing size
        use_convex_hull=False,  # Disable convex hull
        allowed_strategies=["ellipse", "box"],  # Only use ellipse and box
        save_images=True
    )

    results = pipe(
        "path/to/your/image.jpg",
        garment_types="upper",
        output_dir="./custom_output",
        save_mask_s=True,  # Save tight mask before expansion
        save_strong_protect=True  # Save protected areas
    )

    print("✓ Custom configuration applied:")
    print(f"  - Output height: 2048px")
    print(f"  - Allowed strategies: ellipse, box")
    print(f"  - Saved intermediate results to ./custom_output")

    return results


def example_5_pil_images():
    """Example 5: Using PIL Images instead of paths"""
    print("\n" + "="*60)
    print("Example 5: Using PIL Images")
    print("="*60)

    pipe = GarmentMaskPipeline(device="cuda:0")

    # Load images as PIL
    image1 = Image.open("path/to/image1.jpg")
    image2 = Image.open("path/to/image2.jpg")

    # Process PIL images
    results = pipe(
        [image1, image2],
        garment_types="upper"
    )

    print(f"✓ Processed {len(results)} PIL images")

    return results


def example_6_integration_with_inpainting():
    """Example 6: Integration with image inpainting pipeline"""
    print("\n" + "="*60)
    print("Example 6: Integration with Inpainting")
    print("="*60)

    # Generate mask
    pipe = GarmentMaskPipeline(device="cuda:0")
    results = pipe("path/to/your/image.jpg", garment_types="upper")

    result = results[0]
    image = result["image_standardized"]
    mask = result["masks"]["upper"]["person_mask"]

    # Now use with an inpainting model (example with stable diffusion)
    try:
        from diffusers import StableDiffusionInpaintPipeline

        inpaint_pipe = StableDiffusionInpaintPipeline.from_pretrained(
            "stabilityai/stable-diffusion-2-inpainting"
        ).to("cuda")

        # Inpaint the masked region
        inpainted = inpaint_pipe(
            prompt="a stylish red dress",
            image=image,
            mask_image=mask,
            num_inference_steps=50
        ).images[0]

        inpainted.save("inpainted_result.png")
        print("✓ Inpainting completed")

    except ImportError:
        print("⚠ Diffusers not installed, skipping inpainting demo")
        print("  Install with: pip install diffusers")

    return results


def main():
    """Run all examples"""
    print("\n" + "="*60)
    print("WearIT Garment Mask Pipeline - Examples")
    print("="*60)

    # Check if example images exist
    if not os.path.exists("path/to/your/image.jpg"):
        print("\n⚠ WARNING: Please update image paths in the examples!")
        print("  Replace 'path/to/your/image.jpg' with actual image paths.\n")
        return

    # Run examples (uncomment the ones you want to try)

    # results = example_1_basic_usage()
    # visualize_results(results, "./example1_output")

    # results = example_2_multiple_garments()
    # visualize_results(results, "./example2_output")

    # results = example_3_batch_processing()

    # results = example_4_custom_configuration()

    # results = example_5_pil_images()

    # results = example_6_integration_with_inpainting()

    print("\n" + "="*60)
    print("Examples completed!")
    print("="*60 + "\n")


if __name__ == "__main__":
    # Quick test without running all examples
    print("WearIT Garment Mask Pipeline - Example Usage Script")
    print("\nTo run examples, uncomment the desired example in main() function")
    print("and update the image paths.\n")

    # Uncomment to run all examples:
    # main()