experiments/zero_shot_fashion.py

"""
Zero-Shot Fashion Segmentation Experiment

This experiment demonstrates zero-shot learning for fashion segmentation
using SAM 2 with advanced text prompting and attention mechanisms.
"""

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import os
import json
from typing import List, Dict, Tuple
import argparse
from tqdm import tqdm

# Add parent directory to path
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from models.sam2_zeroshot import SAM2ZeroShot, ZeroShotEvaluator
from utils.data_loader import FashionDataLoader
from utils.metrics import SegmentationMetrics
from utils.visualization import visualize_segmentation


class FashionZeroShotExperiment:
    """Zero-shot learning experiment for fashion segmentation."""
    
    def __init__(
        self,
        sam2_checkpoint: str,
        data_dir: str,
        output_dir: str,
        device: str = "cuda",
        use_attention_maps: bool = True,
        temperature: float = 0.1
    ):
        self.device = device
        self.output_dir = output_dir
        
        # Create output directory
        os.makedirs(output_dir, exist_ok=True)
        
        # Initialize model
        self.model = SAM2ZeroShot(
            sam2_checkpoint=sam2_checkpoint,
            device=device,
            use_attention_maps=use_attention_maps,
            temperature=temperature
        )
        
        # Initialize evaluator
        self.evaluator = ZeroShotEvaluator()
        
        # Initialize data loader
        self.data_loader = FashionDataLoader(data_dir)
        
        # Initialize metrics
        self.metrics = SegmentationMetrics()
        
        # Fashion-specific classes
        self.classes = ["shirt", "pants", "dress", "shoes"]
        
        # Prompt strategies to test
        self.prompt_strategies = [
            "basic",      # Simple class names
            "descriptive", # Enhanced descriptions
            "contextual",  # Context-aware prompts
            "detailed"     # Detailed descriptions
        ]
    
    def run_single_experiment(
        self, 
        image: torch.Tensor, 
        ground_truth: Dict[str, torch.Tensor],
        strategy: str = "descriptive"
    ) -> Dict:
        """Run a single zero-shot experiment."""
        # Perform segmentation
        predictions = self.model.segment(image, "fashion", self.classes)
        
        # Evaluate results
        evaluation = self.evaluator.evaluate(predictions, ground_truth)
        
        return {
            'predictions': predictions,
            'evaluation': evaluation,
            'strategy': strategy
        }
    
    def run_comparative_experiment(
        self, 
        num_images: int = 50
    ) -> Dict:
        """Run comparative experiment with different prompt strategies."""
        results = {
            'strategies': {strategy: [] for strategy in self.prompt_strategies},
            'overall_comparison': {},
            'class_analysis': {cls: {strategy: [] for strategy in self.prompt_strategies} 
                             for cls in self.classes}
        }
        
        print(f"Running comparative zero-shot experiment on {num_images} images...")
        
        for i in tqdm(range(num_images)):
            # Load test image and ground truth
            image, ground_truth = self.data_loader.get_test_sample()
            
            # Test each strategy
            for strategy in self.prompt_strategies:
                # Modify model's prompt strategy for this experiment
                if strategy == "basic":
                    # Use simple prompts
                    self.model.advanced_prompts["fashion"] = {
                        "shirt": ["shirt"],
                        "pants": ["pants"],
                        "dress": ["dress"],
                        "shoes": ["shoes"]
                    }
                elif strategy == "descriptive":
                    # Use descriptive prompts
                    self.model.advanced_prompts["fashion"] = {
                        "shirt": ["fashion photography of shirts", "clothing item top"],
                        "pants": ["fashion photography of pants", "lower body clothing"],
                        "dress": ["fashion photography of dresses", "full body garment"],
                        "shoes": ["fashion photography of shoes", "footwear item"]
                    }
                elif strategy == "contextual":
                    # Use contextual prompts
                    self.model.advanced_prompts["fashion"] = {
                        "shirt": ["in a fashion setting, shirt", "worn by a person, shirt"],
                        "pants": ["in a fashion setting, pants", "worn by a person, pants"],
                        "dress": ["in a fashion setting, dress", "worn by a person, dress"],
                        "shoes": ["in a fashion setting, shoes", "worn by a person, shoes"]
                    }
                elif strategy == "detailed":
                    # Use detailed prompts
                    self.model.advanced_prompts["fashion"] = {
                        "shirt": ["high quality fashion photograph of a shirt with clear details", 
                                "professional clothing photography showing shirt"],
                        "pants": ["high quality fashion photograph of pants with clear details",
                                "professional clothing photography showing pants"],
                        "dress": ["high quality fashion photograph of a dress with clear details",
                                "professional clothing photography showing dress"],
                        "shoes": ["high quality fashion photograph of shoes with clear details",
                                "professional clothing photography showing shoes"]
                    }
                
                # Run experiment
                experiment_result = self.run_single_experiment(image, ground_truth, strategy)
                
                # Store results
                results['strategies'][strategy].append(experiment_result['evaluation'])
                
                # Store class-specific results
                for class_name in self.classes:
                    iou_key = f"{class_name}_iou"
                    dice_key = f"{class_name}_dice"
                    
                    if iou_key in experiment_result['evaluation']:
                        results['class_analysis'][class_name][strategy].append({
                            'iou': experiment_result['evaluation'][iou_key],
                            'dice': experiment_result['evaluation'][dice_key]
                        })
                
                # Visualize every 10 images
                if i % 10 == 0:
                    self.visualize_comparison(
                        i, image, ground_truth, 
                        {s: results['strategies'][s][-1] for s in self.prompt_strategies},
                        strategy
                    )
        
        # Compute overall comparison
        for strategy in self.prompt_strategies:
            strategy_results = results['strategies'][strategy]
            if strategy_results:
                results['overall_comparison'][strategy] = {
                    'mean_iou': np.mean([r.get('mean_iou', 0) for r in strategy_results]),
                    'mean_dice': np.mean([r.get('mean_dice', 0) for r in strategy_results]),
                    'std_iou': np.std([r.get('mean_iou', 0) for r in strategy_results]),
                    'std_dice': np.std([r.get('mean_dice', 0) for r in strategy_results])
                }
        
        return results
    
    def run_attention_analysis(self, num_images: int = 20) -> Dict:
        """Run analysis of attention-based prompt generation."""
        results = {
            'with_attention': [],
            'without_attention': [],
            'attention_points': []
        }
        
        print(f"Running attention analysis on {num_images} images...")
        
        for i in tqdm(range(num_images)):
            # Load test image and ground truth
            image, ground_truth = self.data_loader.get_test_sample()
            
            # Test with attention maps
            self.model.use_attention_maps = True
            with_attention = self.run_single_experiment(image, ground_truth, "attention")
            
            # Test without attention maps
            self.model.use_attention_maps = False
            without_attention = self.run_single_experiment(image, ground_truth, "no_attention")
            
            # Store results
            results['with_attention'].append(with_attention['evaluation'])
            results['without_attention'].append(without_attention['evaluation'])
            
            # Analyze attention points
            if with_attention['predictions']:
                # Extract attention points (simplified)
                attention_points = self.extract_attention_points(image, self.classes)
                results['attention_points'].append(attention_points)
        
        return results
    
    def extract_attention_points(self, image: torch.Tensor, classes: List[str]) -> List[Tuple[int, int]]:
        """Extract attention points for visualization."""
        # Simplified attention point extraction
        h, w = image.shape[-2:]
        points = []
        
        for class_name in classes:
            # Generate some sample points (in practice, these would come from attention maps)
            center_x, center_y = w // 2, h // 2
            points.append((center_x, center_y))
            
            # Add some variation
            points.append((center_x + w // 4, center_y))
            points.append((center_x, center_y + h // 4))
        
        return points
    
    def visualize_comparison(
        self,
        image_idx: int,
        image: torch.Tensor,
        ground_truth: Dict[str, torch.Tensor],
        strategy_results: Dict,
        best_strategy: str
    ):
        """Visualize comparison between different strategies."""
        fig, axes = plt.subplots(3, 4, figsize=(20, 15))
        
        # Original image
        axes[0, 0].imshow(image.permute(1, 2, 0).cpu().numpy())
        axes[0, 0].set_title("Original Image")
        axes[0, 0].axis('off')
        
        # Ground truth
        for i, class_name in enumerate(self.classes):
            if class_name in ground_truth:
                axes[0, i+1].imshow(ground_truth[class_name].cpu().numpy(), cmap='gray')
                axes[0, i+1].set_title(f"GT: {class_name}")
            axes[0, i+1].axis('off')
        
        # Best strategy predictions
        best_result = strategy_results[best_strategy]
        for i, class_name in enumerate(self.classes):
            if class_name in best_result:
                axes[1, i].imshow(best_result[class_name].cpu().numpy(), cmap='gray')
                axes[1, i].set_title(f"Best: {class_name}")
            axes[1, i].axis('off')
        
        # Strategy comparison
        strategies = list(strategy_results.keys())
        metrics = ['mean_iou', 'mean_dice']
        
        for i, metric in enumerate(metrics):
            values = [strategy_results[s].get(metric, 0) for s in strategies]
            axes[2, i].bar(strategies, values)
            axes[2, i].set_title(f"{metric.replace('_', ' ').title()}")
            axes[2, i].tick_params(axis='x', rotation=45)
        
        # Add text summary
        summary_text = f"Best Strategy: {best_strategy}\n"
        for strategy, result in strategy_results.items():
            summary_text += f"{strategy}: IoU={result.get('mean_iou', 0):.3f}, Dice={result.get('mean_dice', 0):.3f}\n"
        
        axes[2, 2].text(0.1, 0.5, summary_text, transform=axes[2, 2].transAxes, 
                       verticalalignment='center', fontsize=10)
        axes[2, 2].axis('off')
        axes[2, 3].axis('off')
        
        plt.tight_layout()
        plt.savefig(os.path.join(self.output_dir, f"comparison_{image_idx}.png"))
        plt.close()
    
    def save_results(self, results: Dict, experiment_type: str = "comparative"):
        """Save experiment results."""
        # Save detailed results
        with open(os.path.join(self.output_dir, f'{experiment_type}_results.json'), 'w') as f:
            json.dump(results, f, indent=2)
        
        # Save summary
        if experiment_type == "comparative":
            summary = {
                'experiment_type': experiment_type,
                'num_images': len(results['strategies'][list(results['strategies'].keys())[0]]),
                'overall_comparison': results['overall_comparison'],
                'best_strategy': max(results['overall_comparison'].items(), 
                                   key=lambda x: x[1]['mean_iou'])[0]
            }
        else:
            summary = {
                'experiment_type': experiment_type,
                'attention_analysis': {
                    'with_attention_mean_iou': np.mean([r.get('mean_iou', 0) for r in results['with_attention']]),
                    'without_attention_mean_iou': np.mean([r.get('mean_iou', 0) for r in results['without_attention']]),
                    'attention_improvement': np.mean([r.get('mean_iou', 0) for r in results['with_attention']]) - 
                                          np.mean([r.get('mean_iou', 0) for r in results['without_attention']])
                }
            }
        
        with open(os.path.join(self.output_dir, f'{experiment_type}_summary.json'), 'w') as f:
            json.dump(summary, f, indent=2)
        
        print(f"Results saved to {self.output_dir}")
        if experiment_type == "comparative":
            print(f"Best strategy: {summary['best_strategy']}")
            print(f"Best mean IoU: {summary['overall_comparison'][summary['best_strategy']]['mean_iou']:.3f}")


def main():
    parser = argparse.ArgumentParser(description="Zero-shot fashion segmentation experiment")
    parser.add_argument("--sam2_checkpoint", type=str, required=True, help="Path to SAM 2 checkpoint")
    parser.add_argument("--data_dir", type=str, required=True, help="Path to fashion dataset")
    parser.add_argument("--output_dir", type=str, default="results/zero_shot_fashion", help="Output directory")
    parser.add_argument("--num_images", type=int, default=50, help="Number of test images")
    parser.add_argument("--device", type=str, default="cuda", help="Device to use")
    parser.add_argument("--experiment_type", type=str, default="comparative", 
                       choices=["comparative", "attention"], help="Type of experiment")
    parser.add_argument("--temperature", type=float, default=0.1, help="CLIP temperature")
    
    args = parser.parse_args()
    
    # Run experiment
    experiment = FashionZeroShotExperiment(
        sam2_checkpoint=args.sam2_checkpoint,
        data_dir=args.data_dir,
        output_dir=args.output_dir,
        device=args.device,
        temperature=args.temperature
    )
    
    if args.experiment_type == "comparative":
        results = experiment.run_comparative_experiment(num_images=args.num_images)
    else:
        results = experiment.run_attention_analysis(num_images=args.num_images)
    
    experiment.save_results(results, args.experiment_type)


if __name__ == "__main__":
    main()