explore_splits.py

#!/usr/bin/env python3
"""
Example: Loading splits with comprehensive metadata per sample.

This script demonstrates how to load splits and display
detailed metadata for each sample including concept information, supercategories,
and attribute-specific details.
"""

import json
import csv
import numpy as np
import pandas as pd
from pathlib import Path
from collections import defaultdict
from typing import Dict, List, Set, Tuple, Any

def load_supercategories(file_path: str = "data/things/Categories_final_20200131.tsv") -> Dict[str, Set[str]]:
    """Load supercategory mappings from the THINGS dataset."""
    supercategories = defaultdict(set)
    missing = 0
    
    with open(file_path, "r") as f:
        reader = csv.reader(f, delimiter="\t")
        header = next(reader)
        supercat_set = header[2:]  # Skip word and definition columns
        
        for row in reader:
            word = row[0]
            # Find all supercategories this concept belongs to (all "1" values)
            concept_supercats = []
            for i, value in enumerate(row[2:]):
                if value == "1":
                    concept_supercats.append(supercat_set[i])
            
            if concept_supercats:
                # Add concept to all its supercategories
                for supercat in concept_supercats:
                    supercategories[supercat].add(word)
            else:
                # No supercategory assigned (no "1" values)
                missing += 1
                continue
    
    print(f"✓ Loaded supercategories: {len(supercategories)} categories, {missing} missing assignments")
    return dict(supercategories)

def load_concept_metadata(file_path: str = "data/things/concepts-and-categories.json") -> Dict[str, Dict[str, str]]:
    """Load concept metadata including categories and definitions."""
    with open(file_path, "r") as f:
        data = json.load(f)
    
    metadata = {}
    for item in data:
        metadata[item["concept"]] = {
            "category": item["category"],
            "definition": item["definition"],
            "id": item["id"]
        }
    
    print(f"✓ Loaded metadata for {len(metadata)} concepts")
    return metadata

def load_splits(split_file: str = "splits/k_means_split.json") -> Dict[str, Any]:
    """Load split data from JSON file."""
    with open(split_file, "r") as f:
        return json.load(f)

def load_image_paths(file_path: str = "data/all_filepath.txt") -> List[str]:
    """Load image file paths."""
    with open(file_path, "r") as f:
        return [line.strip() for line in f if line.strip()]

def load_labels(file_path: str = "data/labels_images.txt") -> List[int]:
    """Load image labels."""
    with open(file_path, "r") as f:
        return [int(line.strip()) for line in f if line.strip()]

def extract_concept_name_from_path(image_file: str) -> str:
    """Extract concept name from image file path."""
    # Image files are in format like "aardvark_01b", "abacus_03s", etc.
    # Extract the concept name (part before the underscore)
    if "_" in image_file:
        return image_file.split("_")[0]
    return image_file

def load_feature_to_concepts() -> Dict[str, List[str]]:
    """Load feature to concepts mapping from McRae norms."""
    import json
    with open('data/mcrae-norms-grouped-with-concepts.json', 'r') as f:
        data = json.load(f)
    
    feature_to_concepts = {}
    for norm in data:
        feature = norm['norm']
        concepts = norm.get('concepts-mcrae', [])
        feature_to_concepts[feature] = concepts
    
    return feature_to_concepts

def concept_has_attribute(concept_name: str, attribute: str, feature_to_concepts: Dict[str, List[str]]) -> bool:
    """Check if a concept has a specific attribute."""
    if attribute not in feature_to_concepts:
        return False
    return concept_name in feature_to_concepts[attribute]

def get_sample_metadata(concept_idx: int, image_paths: List[str], labels: List[int],
                       concept_metadata: Dict, supercategories: Dict[str, Set[str]], 
                       concept_to_supercat: Dict[str, str], feature_to_concepts: Dict[str, List[str]] = None,
                       current_attribute: str = None) -> Dict[str, Any]:
    """Get comprehensive metadata for a concept (concept_idx represents concept ID, not sample index)."""
    # Find the first sample that belongs to this concept
    sample_idx = None
    for i, label in enumerate(labels):
        if label == concept_idx:
            sample_idx = i
            break
    
    if sample_idx is None:
        return {"error": f"No samples found for concept ID {concept_idx}"}
    
    if sample_idx >= len(image_paths) or sample_idx >= len(labels):
        return {"error": f"Sample index {sample_idx} out of range"}
    
    image_file = image_paths[sample_idx]
    label = labels[sample_idx]
    
    # Extract concept name from image path
    concept_name = extract_concept_name_from_path(image_file)
    
    # Get concept metadata
    concept_info = concept_metadata.get(concept_name, {})
    
    # Get supercategory
    supercategory = concept_to_supercat.get(concept_name, "Unknown")
    
    # Check if concept has the current attribute
    has_attribute = None
    if feature_to_concepts and current_attribute:
        has_attribute = concept_has_attribute(concept_name, current_attribute, feature_to_concepts)
    
    return {
        "concept_idx": concept_idx,
        "sample_idx": sample_idx,
        "concept_name": concept_name,
        "label": label,
        "image_file": image_file,
        "category": concept_info.get("category", "Unknown"),
        "definition": concept_info.get("definition", "No definition available"),
        "supercategory": supercategory,
        "concept_id": concept_info.get("id", "Unknown"),
        "has_attribute": has_attribute
    }

def print_sample_metadata(metadata: Dict[str, Any], attribute_name: str, split_type: str):
    """Print formatted metadata for a sample."""
    if "error" in metadata:
        print(f"  ERROR: {metadata['error']}")
        return
        
    print(f"  📋 Concept {metadata['concept_idx']} (Sample {metadata['sample_idx']}, {split_type}):")
    print(f"     🏷️  Concept: {metadata['concept_name']}")
    print(f"     🔢 Concept label: {metadata['label']}")
    print(f"     📂 Category: {metadata['category']}")
    print(f"     🏗️  Supercategory: {metadata['supercategory']}")
    print(f"     📝 Definition: {metadata['definition'][:80]}...")
    print(f"     🖼️  Image: {metadata['image_file']}")
    print(f"     🎯 Attribute: {attribute_name}")
    if metadata.get('has_attribute') is not None:
        has_attr = "✅ YES" if metadata['has_attribute'] else "❌ NO"
        print(f"     🔍 Has {attribute_name}: {has_attr}")
    print()

def analyze_kmeans_attribute_splits(split_data: Dict[str, Any], attribute_name: str, 
                                 image_paths: List[str], labels: List[int],
                                 concept_metadata: Dict, supercategories: Dict[str, Set[str]], 
                                 concept_to_supercat: Dict[str, str], feature_to_concepts: Dict[str, List[str]],
                                 max_samples: int = 5):
    """Analyze splits for a specific attribute with detailed metadata."""
    print(f"\n{'='*80}")
    print(f"🎯 {attribute_name}")
    print(f"{'='*80}")
    
    # Find the attribute in the splits
    attribute_idx = None
    for i, attribute in enumerate(split_data.get("attributes", [])):
        if attribute == attribute_name:
            attribute_idx = i
            break
    
    if attribute_idx is None:
        print(f"❌ Attribute '{attribute_name}' not found in this split")
        return
    
    # Get the split for this attribute
    splits = split_data["splits"][attribute_idx]
    train_indices = splits[0]
    test_indices = splits[1]
    
    # Count positive and negative samples
    train_positive = 0
    train_negative = 0
    test_positive = 0
    test_negative = 0
    
    for concept_idx in train_indices:
        metadata = get_sample_metadata(concept_idx, image_paths, labels, concept_metadata, 
                                     supercategories, concept_to_supercat, feature_to_concepts, attribute_name)
        if metadata.get('has_attribute') is True:
            train_positive += 1
        elif metadata.get('has_attribute') is False:
            train_negative += 1
    
    for concept_idx in test_indices:
        metadata = get_sample_metadata(concept_idx, image_paths, labels, concept_metadata, 
                                     supercategories, concept_to_supercat, feature_to_concepts, attribute_name)
        if metadata.get('has_attribute') is True:
            test_positive += 1
        elif metadata.get('has_attribute') is False:
            test_negative += 1
    
    print(f"📊 Split Statistics:")
    train_ratio = len(train_indices) / (len(train_indices) + len(test_indices)) * 100
    print(f"   🚂 Train samples: {len(train_indices)} (Train ratio: {train_ratio:.2f}%)")
    print(f"      ✅ Positive: {train_positive} | ❌ Negative: {train_negative}")
    print(f"   🧪 Test samples: {len(test_indices)}")
    print(f"      ✅ Positive: {test_positive} | ❌ Negative: {test_negative}")
    print(f"   📈 Total samples: {len(train_indices) + len(test_indices)}")
    print(f"      ✅ Total Positive: {train_positive + test_positive} | ❌ Total Negative: {train_negative + test_negative}")
    
    # Show train samples with metadata (both positive and negative)
    print(f"\n🚂 TRAIN SAMPLES (showing first {min(max_samples, len(train_indices))}):")
    print("-" * 60)
    positive_samples = []
    negative_samples = []
    
    for concept_idx in train_indices:
        metadata = get_sample_metadata(concept_idx, image_paths, labels, concept_metadata, 
                                     supercategories, concept_to_supercat, feature_to_concepts, attribute_name)
        if metadata.get('has_attribute') is True:
            positive_samples.append(metadata)
        elif metadata.get('has_attribute') is False:
            negative_samples.append(metadata)
    
    # Show positive samples first, then negative
    samples_shown = 0
    for metadata in positive_samples[:max_samples//2]:
        print_sample_metadata(metadata, attribute_name, "TRAIN")
        samples_shown += 1
    
    for metadata in negative_samples[:max_samples//2]:
        if samples_shown >= max_samples:
            break
        print_sample_metadata(metadata, attribute_name, "TRAIN")
        samples_shown += 1
    
    # Show test samples with metadata (both positive and negative)
    print(f"\n🧪 TEST SAMPLES (showing first {min(max_samples, len(test_indices))}):")
    print("-" * 60)
    positive_samples = []
    negative_samples = []
    
    for concept_idx in test_indices:
        metadata = get_sample_metadata(concept_idx, image_paths, labels, concept_metadata, 
                                     supercategories, concept_to_supercat, feature_to_concepts, attribute_name)
        if metadata.get('has_attribute') is True:
            positive_samples.append(metadata)
        elif metadata.get('has_attribute') is False:
            negative_samples.append(metadata)
    
    # Show positive samples first, then negative
    samples_shown = 0
    for metadata in positive_samples[:max_samples//2]:
        print_sample_metadata(metadata, attribute_name, "TEST")
        samples_shown += 1
    
    for metadata in negative_samples[:max_samples//2]:
        if samples_shown >= max_samples:
            break
        print_sample_metadata(metadata, attribute_name, "TEST")
        samples_shown += 1
    
    # Analyze supercategory distribution
    print(f"\n📊 SUPERCATEGORY DISTRIBUTION:")
    print("-" * 60)
    
    # Count supercategories in train and test
    train_supercats = defaultdict(int)
    test_supercats = defaultdict(int)
    
    for idx in train_indices:
        if idx < len(image_paths):
            concept_name = extract_concept_name_from_path(image_paths[idx])
            supercat = concept_to_supercat.get(concept_name, "Unknown")
            train_supercats[supercat] += 1
    
    for idx in test_indices:
        if idx < len(image_paths):
            concept_name = extract_concept_name_from_path(image_paths[idx])
            supercat = concept_to_supercat.get(concept_name, "Unknown")
            test_supercats[supercat] += 1
    
    # Display distribution
    all_supercats = set(train_supercats.keys()) | set(test_supercats.keys())
    print(f"{'Supercategory':<25} {'Train':<8} {'Test':<8} {'Total':<8} {'Train%':<8}")
    print("-" * 65)
    
    for supercat in sorted(all_supercats):
        train_count = train_supercats.get(supercat, 0)
        test_count = test_supercats.get(supercat, 0)
        total = train_count + test_count
        train_pct = (train_count / total * 100) if total > 0 else 0
        print(f"{supercat:<25} {train_count:<8} {test_count:<8} {total:<8} {train_pct:<7.1f}%")

def main():
    """Main function to demonstrate split loading with metadata."""
    print("🎯 SPLIT WITH METADATA DEMONSTRATION")
    print("=" * 80)
    
    # Load metadata
    print("\n📚 Loading metadata...")
    concept_metadata = load_concept_metadata()
    supercategories = load_supercategories()
    feature_to_concepts = load_feature_to_concepts()
    
    # Create concept to supercategory mapping (assign largest supercategory)
    concept_to_supercat = {}
    # Get all unique concepts across all supercategories
    all_concepts = set().union(*supercategories.values())
    
    for concept in all_concepts:
        largest_supercat = None
        largest_size = 0
        # Find all supercategories this concept belongs to and pick the largest
        for supercat, concepts in supercategories.items():
            if concept in concepts and len(concepts) > largest_size:
                largest_supercat = supercat
                largest_size = len(concepts)
        concept_to_supercat[concept] = largest_supercat
    
    # Load image paths and labels
    print("\n🖼️  Loading image data...")
    image_paths = load_image_paths()
    labels = load_labels()
    
    print(f"✓ Loaded {len(image_paths)} image paths")
    print(f"✓ Loaded {len(labels)} labels")
    
    # Load available attributes
    with open("splits/intersect_attributes.json", "r") as f:
        available_attributes = json.load(f)
    
    print(f"✓ Available attributes: {len(available_attributes)}")
    
    # Show available attributes count
    print(f"\n📋 AVAILABLE ATTRIBUTES: {len(available_attributes)}")
    
    # Find all split files in the splits folder
    splits_folder = Path("splits")
    split_files = list(splits_folder.glob("*.json"))
    split_files = [f for f in split_files if f.name != "intersect_attributes.json"]  # Exclude the attributes list
    
    print(f"\n🔄 Found {len(split_files)} split files to analyze:")
    for split_file in split_files:
        print(f"   - {split_file.name}")
    
    # Analyze specific attributes
    attributes_to_analyze = ["has_wheels"]
    
    # Iterate through each split file
    for split_file in split_files:
        print(f"\n{'='*80}")
        print(f"📁 ANALYZING SPLIT FILE: {split_file.name}")
        print(f"{'='*80}")
        
        # Load splits data
        splits_data = load_splits(str(split_file))
        print(f"✓ Loaded splits for {len(splits_data.get('attributes', []))} attributes")
        
        print(f"\n🔍 Analyzing {split_file.stem} splits for {len(attributes_to_analyze)} attributes...")
        
        for attribute_name in attributes_to_analyze:
            if attribute_name in available_attributes:
                analyze_kmeans_attribute_splits(
                    splits_data, attribute_name, image_paths, labels,
                    concept_metadata, supercategories, concept_to_supercat, feature_to_concepts,
                    max_samples=3
                )
            else:
                print(f"⚠️  Attribute '{attribute_name}' not available in {split_file.stem} splits")
        break
    
    # Summary statistics
    print(f"\n📈 SUMMARY STATISTICS")
    print("=" * 80)
    
    print(f"📊 Total samples in dataset: {len(image_paths)}")
    print(f"📊 Total concepts with metadata: {len(concept_metadata)}")
    print(f"📊 Total supercategories: {len(supercategories)}")
    print(f"📊 Total split files analyzed: {len(split_files)}")
    
    # Show top supercategories
    supercat_counts = {name: len(concepts) for name, concepts in supercategories.items()}
    sorted_supercats = sorted(supercat_counts.items(), key=lambda x: x[1], reverse=True)
    
    print(f"\n🏆 Top 5 supercategories by concept count:")
    for i, (supercat, count) in enumerate(sorted_supercats[:5], 1):
        print(f"   {i}. {supercat}: {count} concepts")
    
    print(f"\n✅ Split analysis completed successfully!")
    print("=" * 80)

if __name__ == "__main__":
    main()