AtlasMapFromDataset

from nomic import AtlasDataset
import pandas as pd
from PIL import Image
import requests
from io import BytesIO
import concurrent.futures
import uuid
import gc
import os


Image.MAX_IMAGE_PIXELS = None  
def load_image_urls(csv_file):
    df = pd.read_csv(csv_file)
    image_urls = df['original_url'].tolist()  
    return image_urls


def download_image(url):
    try:
        response = requests.get(url, timeout=10)
        response.raise_for_status()  
        img = Image.open(BytesIO(response.content))
        return img, url  
    except requests.exceptions.RequestException as e:
        print(f"Error downloading {url}: {e}")
        return None, url

def download_images(image_urls, max_workers=5):  
    images = []
    data = []
    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:
        futures = {executor.submit(download_image, url): url for url in image_urls}
        for future in concurrent.futures.as_completed(futures):
            img, url = future.result()
            if img:
                images.append(img)
                data.append({'id': str(uuid.uuid4()), 'url': url})  
    return images, data


def create_atlas_map_in_batches(dataset_name, image_urls, batch_size=100):  
    # Load the existing AtlasDataset
    dataset = AtlasDataset(dataset_name, unique_id_field="id")


   
    for i in range(0, len(image_urls), batch_size):
        batch_urls = image_urls[i:i + batch_size]
        images, data = download_images(batch_urls)
        if images:
          
            dataset.add_data(data=data, blobs=images)
            print(f"Processed batch {i // batch_size + 1}")
            del images
            gc.collect() 

            dataset.save(f"{dataset_name}_intermediate_{i}.hdf5")
    

            dataset.save(f"{dataset_name}_intermediate_{i}.hdf5")
            print(f"Saved intermediate state at batch {i // batch_size + 1}")


   
    atlas_map = dataset.create_index(
        topic_model=True,
        embedding_model='nomic-embed-vision-v1.5'
        
    )
    return atlas_map


if __name__ == "__main__":
    
    dataset_name = input("Enter the name of the dataset: ").strip()
    csv_file_path = input("Enter the path to the CSV file: ").strip()

    image_urls = load_image_urls(csv_file_path)


    if image_urls:
 
        intermediate_files = [f for f in os.listdir('.') if f.startswith(f"{dataset_name}_intermediate_")]
        if intermediate_files:
            latest_file = max(intermediate_files, key=os.path.getctime)
            print(f"Resuming from {latest_file}")
            dataset = AtlasDataset.load(latest_file)
        
            atlas_map = create_atlas_map_in_batches(dataset_name, image_urls, batch_size=500)


        
        print(f"Atlas map has been created for dataset '{dataset_name}'")
    else:
        print("Error: No image URLs were loaded.")