Spaces:
Sleeping
Sleeping
File size: 8,267 Bytes
e81469a | 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 189 190 191 192 193 194 195 196 197 198 199 200 201 202 | """Load galaxy data from HuggingFace datasets and update profiles."""
import os
import logging
from typing import Dict, Any
import requests
from PIL import Image
import io
try:
from datasets import load_dataset
DATASETS_AVAILABLE = True
except ImportError:
DATASETS_AVAILABLE = False
logging.warning("datasets library not available, using fallback data")
logger = logging.getLogger(__name__)
def load_galaxy_data() -> Dict[str, Dict[str, Any]]:
"""Load galaxy data from HuggingFace datasets and return structured profiles."""
if not DATASETS_AVAILABLE:
logger.warning("Using fallback galaxy data")
return get_fallback_galaxy_data()
try:
# Load both datasets
euclid_dataset = load_dataset("mwalmsley/gz_euclid", "tiny", split="test")
descriptions_dataset = load_dataset("Smith42/dating_pool_but_galaxies", split="train")
# Convert to dictionaries for easier access
euclid_data = {row['id_str']: row for row in euclid_dataset}
descriptions_data = {row['id_str']: row for row in descriptions_dataset}
galaxy_profiles = {}
# Match on id_str and create profiles with numbered galaxy IDs
matched_galaxies = []
for galaxy_id_str in euclid_data.keys():
if galaxy_id_str in descriptions_data:
matched_galaxies.append(galaxy_id_str)
# Sort for consistent ordering
matched_galaxies.sort()
# Create numbered galaxy profiles (galaxy_01, galaxy_02, etc.)
for i, galaxy_id_str in enumerate(matched_galaxies, 1):
numbered_id = f"galaxy_{i:02d}"
euclid_row = euclid_data[galaxy_id_str]
desc_row = descriptions_data[galaxy_id_str]
# Extract name and description from caption
caption = desc_row.get('caption', '')
bio = ""
description = ""
name = f"Galaxy {i}"
# Parse the caption to extract bio and name if available
if caption:
lines = caption.split('\\n')
if lines and lines[0].startswith('Bio:'):
bio_part = lines[0].replace('Bio:', '').strip()
bio = bio_part[:100] + "..." if len(bio_part) > 100 else bio_part
description = bio_part
else:
description = caption.strip()
bio = description[:100] + "..." if len(description) > 100 else description
galaxy_profiles[numbered_id] = {
"name": name,
"bio": bio,
"description": description,
"tags": ["Cosmic", "Mysterious"],
"color": generate_color_from_id(galaxy_id_str),
"id_str": galaxy_id_str,
"image_data": euclid_row.get('image'), # PIL Image object
"euclid_features": {k: v for k, v in euclid_row.items() if k not in ['image', 'id_str']}
}
logger.info(f"Loaded {len(galaxy_profiles)} galaxy profiles from HuggingFace datasets")
return galaxy_profiles
except Exception as e:
logger.error(f"Error loading from HuggingFace datasets: {e}")
return get_fallback_galaxy_data()
def generate_color_from_id(id_str: str) -> str:
"""Generate a consistent color from galaxy ID string."""
colors = [
"#A688C9", "#D68B8B", "#E8A0D0", "#7BC9A0", "#E87D5A",
"#6B8DD6", "#D4A76A", "#C9A688", "#B8A0C9", "#A0D4E8",
"#E8D4A0", "#A0E8B8", "#E8A0A0", "#A0A0E8", "#E8E8A0",
"#D0A0E8", "#A0E8D0", "#E8C4A0", "#C4A0E8", "#A0C4E8",
"#E8A0C4", "#B4E8A0", "#A0E8E8"
]
# Use hash to get consistent color index
hash_val = sum(ord(c) for c in id_str)
return colors[hash_val % len(colors)]
def get_fallback_galaxy_data() -> Dict[str, Dict[str, Any]]:
"""Fallback galaxy data using original profiles with numbered IDs."""
fallback_profiles = {
"galaxy_01": {
"name": "Velvet Vortex",
"bio": "Smooth operator with a soft spiral glow.",
"description": "I keep my arms tight and my luminosity low-key. Looking for someone who appreciates subtlety over spectacle.",
"tags": ["Smooth Spiral", "Low-Key", "Gentle Glow"],
"color": "#A688C9",
"image_url": None,
"image_data": None
},
"galaxy_02": {
"name": "Crimson Drift",
"bio": "Redshifted and proud.",
"description": "I have been expanding away from everyone for billions of years and I am NOT slowing down. Commitment-phobic? Maybe. Mysterious? Definitely.",
"tags": ["High Redshift", "Distant", "Loner Vibes"],
"color": "#D68B8B",
"image_url": None,
"image_data": None
}
# Add more as needed...
}
# Generate numbered galaxy profiles if we have fewer than 23
for i in range(1, 24):
galaxy_id = f"galaxy_{i:02d}"
if galaxy_id not in fallback_profiles:
fallback_profiles[galaxy_id] = {
"name": f"Galaxy {i}",
"bio": "A mysterious galaxy in the cosmic dating scene.",
"description": "This galaxy is looking for its perfect cosmic match.",
"tags": ["Mysterious", "Cosmic"],
"color": "#A688C9",
"image_url": None,
"image_data": None
}
return fallback_profiles
def download_galaxy_image(image_data, galaxy_id: str, images_dir: str = "images") -> bool:
"""Download and save a galaxy image from HuggingFace dataset."""
if not image_data:
return False
try:
# Ensure images directory exists
os.makedirs(images_dir, exist_ok=True)
# If image_data is a PIL Image (from datasets)
if hasattr(image_data, 'save') and hasattr(image_data, 'mode'):
image_path = os.path.join(images_dir, f"{galaxy_id}.jpg")
# Convert to RGB if needed (in case it's RGBA or other format)
if image_data.mode != 'RGB':
image_data = image_data.convert('RGB')
image_data.save(image_path, 'JPEG', quality=85)
logger.info(f"Saved image for {galaxy_id}")
return True
# If image_data has a URL (fallback)
elif isinstance(image_data, dict) and 'url' in image_data:
response = requests.get(image_data['url'])
if response.status_code == 200:
image = Image.open(io.BytesIO(response.content))
if image.mode != 'RGB':
image = image.convert('RGB')
image_path = os.path.join(images_dir, f"{galaxy_id}.jpg")
image.save(image_path, 'JPEG', quality=85)
logger.info(f"Downloaded and saved image for {galaxy_id}")
return True
return False
except Exception as e:
logger.error(f"Error downloading image for {galaxy_id}: {e}")
return False
def update_galaxy_images():
"""Download all galaxy images from the dataset."""
galaxy_data = load_galaxy_data()
for galaxy_id, profile in galaxy_data.items():
if profile.get('image_data'):
download_galaxy_image(profile['image_data'], galaxy_id)
elif profile.get('image_url'):
# Try to download from URL
try:
response = requests.get(profile['image_url'])
if response.status_code == 200:
image = Image.open(io.BytesIO(response.content))
download_galaxy_image(image, galaxy_id)
except Exception as e:
logger.error(f"Error downloading from URL for {galaxy_id}: {e}")
if __name__ == "__main__":
# Test the data loading
data = load_galaxy_data()
print(f"Loaded {len(data)} galaxies")
for galaxy_id, profile in list(data.items())[:3]:
print(f"{galaxy_id}: {profile['name']} - {profile['bio'][:50]}...")
# Update images
update_galaxy_images() |