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 )