File size: 6,847 Bytes
18a82fb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 | 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
) |