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()