Spaces:

readylab
/

VizRef

Sleeping

VizRef / scripts /create_splits.py

Add model and inference code

18a82fb about 1 year ago

6.85 kB

import json
import random
from collections import defaultdict
from sklearn.model_selection import train_test_split
from pathlib import Path


def create_product_aware_splits(metadata_path, output_dir,
                                val_ratio=0.15, test_ratio=0.15,
                                min_images_per_decade=50, source_filter=None):
    """
    Create splits ensuring:
    1. All images of same product stay together
    2. Balanced decades in each split
    3. Stratified by classification if possible
    """

    # Load metadata
    with open(metadata_path, 'r', encoding='utf-8') as f:
        data = json.load(f)
    
    if source_filter == 'mobile':
        # Filter for Mobile Phone Museum AND only first images (img_0)
        filtered_data = []
        for entry in data:
            if entry.get('source') == 'https://www.mobilephonemuseum.com/':
                # Only include first image of each product
                if 'img_0' in entry.get('id', '') or entry.get('image_index') == 0:
                    filtered_data.append(entry)
        
        data = filtered_data  # Replace data with filtered data
        print(f"Mobile Phone Museum first images (img_0): {len(data)}")
        
        if len(data) == 0:
            print("No Mobile Phone Museum entries found!")
            return
    elif source_filter == 'phone_and_calculator':
        # Filter for Mobile Phone Museum and Datamath Calculator sources
        filtered_data = []
        for entry in data:
            source = entry.get('source', '')
            # Include Mobile Phone Museum OR Datamath Calculator Museum
            if source in ['https://www.mobilephonemuseum.com/', 'http://www.datamath.org/']:
                # Only include first image of each product
                if 'img_0' in entry.get('id', '') or entry.get('image_index') == 0:
                    filtered_data.append(entry)
        
        data = filtered_data
        print(f"Mobile Phone Museum + Calculator Museum first images (img_0): {len(data)}")
            
    # Group by product_id to keep product images together
    products = defaultdict(list)
    for entry in data:
        products[entry['product_id']].append(entry)

    # Analyze products by decade and classification
    product_info = {}
    decade_products = defaultdict(list)

    for product_id, images in products.items():
        # All images of a product have same metadata
        first_image = images[0]
        product_info[product_id] = {
            'decade': first_image['decade'],
            'classification': first_image['classification'],
            'country': first_image['country'],
            'images': images
        }
        decade_products[first_image['decade']].append(product_id)

    # Check if we have enough data
    print("\nProducts per decade:")
    for decade, product_ids in sorted(decade_products.items()):
        image_count = sum(len(products[pid]) for pid in product_ids)
        print(f"  {decade}: {len(product_ids)} products, {image_count} images")

    # Split products (not images) to avoid leakage
    train_products = []
    val_products = []
    test_products = []

    # Stratified split by decade
    for decade, product_ids in decade_products.items():
        if len(product_ids) < 3:
            # Too few products, put all in train
            train_products.extend(product_ids)
            print(f"Warning: {decade} has only {len(product_ids)} products, all going to train")
        else:
            # First split off test set
            train_val_ids, test_ids = train_test_split(
                product_ids,
                test_size=test_ratio,
                random_state=42
            )

            # Then split train and val
            train_ids, val_ids = train_test_split(
                train_val_ids,
                test_size=val_ratio / (1 - test_ratio),
                random_state=42
            )

            train_products.extend(train_ids)
            val_products.extend(val_ids)
            test_products.extend(test_ids)

    # Convert back to image entries
    train_data = []
    val_data = []
    test_data = []

    for pid in train_products:
        train_data.extend(product_info[pid]['images'])
    for pid in val_products:
        val_data.extend(product_info[pid]['images'])
    for pid in test_products:
        test_data.extend(product_info[pid]['images'])

    # Shuffle images within each split
    random.shuffle(train_data)
    random.shuffle(val_data)
    random.shuffle(test_data)

    # Save splits
    splits = {
        'train': train_data,
        'val': val_data,
        'test': test_data
    }

    Path(output_dir).mkdir(parents=True, exist_ok=True)

    for split_name, split_data in splits.items():
        # Save full metadata
        with open(Path(output_dir) / f'{split_name}.json', 'w') as f:
            json.dump(split_data, f, indent=2)

        # Save URL list for easy downloading
        urls = [d['url'] for d in split_data]
        with open(Path(output_dir) / f'{split_name}_urls.txt', 'w') as f:
            f.write('\n'.join(urls))

    # Print detailed statistics
    print(f"\n=== Split Statistics ===")
    print(f"Train: {len(train_products)} products, {len(train_data)} images")
    print(f"Val: {len(val_products)} products, {len(val_data)} images")
    print(f"Test: {len(test_products)} products, {len(test_data)} images")

    # Per-decade breakdown
    for split_name, split_data in splits.items():
        print(f"\n{split_name.capitalize()} split decades:")
        decade_counts = defaultdict(int)
        for entry in split_data:
            decade_counts[entry['decade']] += 1
        for decade in sorted(decade_counts.keys()):
            print(f"  {decade}: {decade_counts[decade]} images")

    # Save split summary
    summary = {
        'train_products': len(train_products),
        'val_products': len(val_products),
        'test_products': len(test_products),
        'train_images': len(train_data),
        'val_images': len(val_data),
        'test_images': len(test_data),
        'splits_by_decade': {}
    }

    for split_name, split_data in splits.items():
        decade_counts = defaultdict(int)
        for entry in split_data:
            decade_counts[entry['decade']] += 1
        summary['splits_by_decade'][split_name] = dict(decade_counts)

    with open(Path(output_dir) / 'split_summary.json', 'w') as f:
        json.dump(summary, f, indent=2)


if __name__ == "__main__":
    project_root = Path(__file__).parent.parent
    
    create_product_aware_splits(
        str(project_root / 'data/metadata/processed_metadata.json'),
        str(project_root / 'data/splits'),
        val_ratio=0.15,
        test_ratio=0.15,
        source_filter='phone_and_calculator'  # Include Mobile Phone Museum + Datamath Calculator Museum
    )

	import json
	import random
	from collections import defaultdict
	from sklearn.model_selection import train_test_split
	from pathlib import Path


	def create_product_aware_splits(metadata_path, output_dir,
	val_ratio=0.15, test_ratio=0.15,
	min_images_per_decade=50, source_filter=None):
	"""
	Create splits ensuring:
	1. All images of same product stay together
	2. Balanced decades in each split
	3. Stratified by classification if possible
	"""

	# Load metadata
	with open(metadata_path, 'r', encoding='utf-8') as f:
	data = json.load(f)

	if source_filter == 'mobile':
	# Filter for Mobile Phone Museum AND only first images (img_0)
	filtered_data = []
	for entry in data:
	if entry.get('source') == 'https://www.mobilephonemuseum.com/':
	# Only include first image of each product
	if 'img_0' in entry.get('id', '') or entry.get('image_index') == 0:
	filtered_data.append(entry)

	data = filtered_data # Replace data with filtered data
	print(f"Mobile Phone Museum first images (img_0): {len(data)}")

	if len(data) == 0:
	print("No Mobile Phone Museum entries found!")
	return
	elif source_filter == 'phone_and_calculator':
	# Filter for Mobile Phone Museum and Datamath Calculator sources
	filtered_data = []
	for entry in data:
	source = entry.get('source', '')
	# Include Mobile Phone Museum OR Datamath Calculator Museum
	if source in ['https://www.mobilephonemuseum.com/', 'http://www.datamath.org/']:
	# Only include first image of each product
	if 'img_0' in entry.get('id', '') or entry.get('image_index') == 0:
	filtered_data.append(entry)

	data = filtered_data
	print(f"Mobile Phone Museum + Calculator Museum first images (img_0): {len(data)}")

	# Group by product_id to keep product images together
	products = defaultdict(list)
	for entry in data:
	products[entry['product_id']].append(entry)

	# Analyze products by decade and classification
	product_info = {}
	decade_products = defaultdict(list)

	for product_id, images in products.items():
	# All images of a product have same metadata
	first_image = images[0]
	product_info[product_id] = {
	'decade': first_image['decade'],
	'classification': first_image['classification'],
	'country': first_image['country'],
	'images': images
	}
	decade_products[first_image['decade']].append(product_id)

	# Check if we have enough data
	print("\nProducts per decade:")
	for decade, product_ids in sorted(decade_products.items()):
	image_count = sum(len(products[pid]) for pid in product_ids)
	print(f" {decade}: {len(product_ids)} products, {image_count} images")

	# Split products (not images) to avoid leakage
	train_products = []
	val_products = []
	test_products = []

	# Stratified split by decade
	for decade, product_ids in decade_products.items():
	if len(product_ids) < 3:
	# Too few products, put all in train
	train_products.extend(product_ids)
	print(f"Warning: {decade} has only {len(product_ids)} products, all going to train")
	else:
	# First split off test set
	train_val_ids, test_ids = train_test_split(
	product_ids,
	test_size=test_ratio,
	random_state=42
	)

	# Then split train and val
	train_ids, val_ids = train_test_split(
	train_val_ids,
	test_size=val_ratio / (1 - test_ratio),
	random_state=42
	)

	train_products.extend(train_ids)
	val_products.extend(val_ids)
	test_products.extend(test_ids)

	# Convert back to image entries
	train_data = []
	val_data = []
	test_data = []

	for pid in train_products:
	train_data.extend(product_info[pid]['images'])
	for pid in val_products:
	val_data.extend(product_info[pid]['images'])
	for pid in test_products:
	test_data.extend(product_info[pid]['images'])

	# Shuffle images within each split
	random.shuffle(train_data)
	random.shuffle(val_data)
	random.shuffle(test_data)

	# Save splits
	splits = {
	'train': train_data,
	'val': val_data,
	'test': test_data
	}

	Path(output_dir).mkdir(parents=True, exist_ok=True)

	for split_name, split_data in splits.items():
	# Save full metadata
	with open(Path(output_dir) / f'{split_name}.json', 'w') as f:
	json.dump(split_data, f, indent=2)

	# Save URL list for easy downloading
	urls = [d['url'] for d in split_data]
	with open(Path(output_dir) / f'{split_name}_urls.txt', 'w') as f:
	f.write('\n'.join(urls))

	# Print detailed statistics
	print(f"\n=== Split Statistics ===")
	print(f"Train: {len(train_products)} products, {len(train_data)} images")
	print(f"Val: {len(val_products)} products, {len(val_data)} images")
	print(f"Test: {len(test_products)} products, {len(test_data)} images")

	# Per-decade breakdown
	for split_name, split_data in splits.items():
	print(f"\n{split_name.capitalize()} split decades:")
	decade_counts = defaultdict(int)
	for entry in split_data:
	decade_counts[entry['decade']] += 1
	for decade in sorted(decade_counts.keys()):
	print(f" {decade}: {decade_counts[decade]} images")

	# Save split summary
	summary = {
	'train_products': len(train_products),
	'val_products': len(val_products),
	'test_products': len(test_products),
	'train_images': len(train_data),
	'val_images': len(val_data),
	'test_images': len(test_data),
	'splits_by_decade': {}
	}

	for split_name, split_data in splits.items():
	decade_counts = defaultdict(int)
	for entry in split_data:
	decade_counts[entry['decade']] += 1
	summary['splits_by_decade'][split_name] = dict(decade_counts)

	with open(Path(output_dir) / 'split_summary.json', 'w') as f:
	json.dump(summary, f, indent=2)


	if __name__ == "__main__":
	project_root = Path(__file__).parent.parent

	create_product_aware_splits(
	str(project_root / 'data/metadata/processed_metadata.json'),
	str(project_root / 'data/splits'),
	val_ratio=0.15,
	test_ratio=0.15,
	source_filter='phone_and_calculator' # Include Mobile Phone Museum + Datamath Calculator Museum
	)