Azam-Rabiee/wealthy-estimator-api

1#!/usr/bin/env python32"""3Generate Embeddings for Wealthy Profiles4 5This script reads the inputs.json file, loads images for each profile, 6generates embeddings using the same model as the main application, 7and saves the results to wealthy_profiles.json.8"""9 10import json11import numpy as np12from PIL import Image13import torch14from torchvision import transforms15from transformers import AutoImageProcessor, AutoModel16import os17from pathlib import Path18 19 20def extract_embedding(image_path: str, model, transform, device) -> np.ndarray:21    """Extract embedding from image using the pre-trained model"""22    try:23        # Load and preprocess image24        image = Image.open(image_path)25        26        # Convert to RGB if needed27        if image.mode != 'RGB':28            image = image.convert('RGB')29        30        # Apply transformations31        image_tensor = transform(image).unsqueeze(0).to(device)32        33        with torch.no_grad():34            # Get model output35            outputs = model(image_tensor)36            37            # For ResNet models, the output is the final feature tensor38            # We need to project it to 49 dimensions to match the profiles39            if hasattr(outputs, 'last_hidden_state'):40                # If it has last_hidden_state (unlikely for ResNet), use it41                features = outputs.last_hidden_state.mean(dim=1).cpu().numpy()42            else:43                # For ResNet, use the direct output44                features = outputs.cpu().numpy()45            46            # Flatten the features47            features = features.flatten()48            49            # Project to 49 dimensions by taking first 49 or sampling50            if len(features) >= 49:51                embedding = features[:49]52            else:53                # If we have fewer than 49 dimensions, pad with zeros54                embedding = np.zeros(49)55                embedding[:len(features)] = features56            57            # Normalize embedding58            embedding = embedding / np.linalg.norm(embedding)59            60            return embedding61            62    except Exception as e:63        print(f"Error processing {image_path}: {e}")64        return None65 66 67def main():68    # Configuration - same as in app/config.py69    MODEL_NAME = "microsoft/resnet-50"70    MAX_IMAGE_SIZE = 22471    DEVICE = "cpu"  # Change to "cuda" if GPU available72    73    print(f"Loading model: {MODEL_NAME}")74    75    # Load pre-trained model and processor76    processor = AutoImageProcessor.from_pretrained(MODEL_NAME)77    model = AutoModel.from_pretrained(MODEL_NAME).to(DEVICE)78    model.eval()79    80    # Image preprocessing81    transform = transforms.Compose([82        transforms.Resize((MAX_IMAGE_SIZE, MAX_IMAGE_SIZE)),83        transforms.ToTensor(),84        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])85    ])86    87    print("Model loaded successfully!")88    89    # Read inputs.json90    print("Reading inputs.json...")91    with open('inputs.json', 'r') as f:92        input_data = json.load(f)93    94    print(f"Found {len(input_data['profiles'])} profiles in inputs.json")95    96    # Process each profile and generate embeddings97    processed_profiles = []98    99    for profile in input_data['profiles']:100        name = profile['name']101        net_worth = profile['net_worth']102        occupation = profile['occupation']103        image_path = profile['image']104        105        print(f"\nProcessing {name}...")106        107        if image_path and image_path.strip():108            if os.path.exists(image_path):109                # Generate embedding110                embedding = extract_embedding(image_path, model, transform, DEVICE)111                112                if embedding is not None:113                    processed_profile = {114                        "name": name,115                        "net_worth": net_worth,116                        "occupation": occupation,117                        "embedding": embedding.tolist(),118                        "image": image_path119                    }120                    processed_profiles.append(processed_profile)121                    print(f"  ✓ Generated embedding with {len(embedding)} dimensions")122                else:123                    print(f"  ✗ Failed to generate embedding")124            else:125                print(f"  ✗ Image file not found: {image_path}")126        else:127            print(f"  ✗ No image path provided")128    129    print(f"\nSuccessfully processed {len(processed_profiles)} profiles with embeddings")130    131    # Save the processed profiles to wealthy_profiles.json132    output_data = {133        "profiles": processed_profiles134    }135    136    with open('wealthy_profiles.json', 'w') as f:137        json.dump(output_data, f, indent=2)138    139    print(f"Saved {len(processed_profiles)} profiles to wealthy_profiles.json")140    141    # Show sample of the generated data142    if processed_profiles:143        sample_profile = processed_profiles[0]144        print(f"\nSample profile:")145        print(f"  Name: {sample_profile['name']}")146        print(f"  Net Worth: ${sample_profile['net_worth']:,}")147        print(f"  Occupation: {sample_profile['occupation']}")148        print(f"  Embedding dimensions: {len(sample_profile['embedding'])}")149        print(f"  First 5 embedding values: {sample_profile['embedding'][:5]}")150    151    print("\nDone!")152 153 154if __name__ == "__main__":155    main()