Image Embeddings and UMAP Generation

This notebook may contain text, code and images generated by artificial intelligence. Used model: claude-sonnet-4-20250514, vision model: claude-sonnet-4-20250514, endpoint: None, bia-bob version: 0.34.3.. It is good scientific practice to check the code and results it produces carefully. Read more about code generation using bia-bob

Image Embeddings and UMAP Generation#

This notebook demonstrates how to:

Generate vision embeddings using the openai/clip-vit-base-patch32
Create a 3D UMAP visualization of the embeddings
Save the results for VEST

import os
import numpy as np
import pandas as pd
from PIL import Image
import torch
import torch.nn.functional as F
from transformers import AutoImageProcessor, AutoModel
from umap import UMAP
import random
from pathlib import Path
import requests
from io import BytesIO
import torch

import transformers
transformers.__version__

'5.12.1'

base_dir = Path("./")
images_dir = base_dir / "images"

Load Vision Embedding Model#

from transformers import CLIPProcessor, CLIPModel
import pandas as pd
import torch
from PIL import Image
import requests
import os

# Initialize models
clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

Generate Vision Embeddings for All Images#

import os
import numpy as np
import pandas as pd
from PIL import Image
import torch
import torch.nn.functional as F

# Get list of all image files in the images folder and subfolders
image_files = []
for root, dirs, files in os.walk(images_dir):
    for f in files:
        if f.lower().endswith('.png') or f.lower().endswith('.jpg'):
            image_files.append(os.path.join(root, f))



print(f"{len(image_files)} images found")

9285 images found

image_files = random.sample(image_files, 500)

print(f"Processing {len(image_files)} images for embeddings...")

# Initialize lists to store results
filenames = []
embeddings = []
images = []

# Loop through all image files
for i, image_path in enumerate(image_files):
    # Obtain filename from path
    filename = os.path.relpath(image_path, images_dir)
    
    # Load the image
    current_image = Image.open(image_path)
    images.append(np.asarray(current_image))

    # Apply the processing pipeline
    current_inputs = clip_processor(images=current_image, return_tensors="pt")
    with torch.no_grad():
        current_np_emb = clip_model.get_image_features(**current_inputs)[0].cpu().squeeze()[0].tolist()

    # Store results
    filenames.append(filename)
    embeddings.append(current_np_emb)

    if (i + 1) % 100 == 0:
        print(f"Processed {i + 1}/{len(image_files)} images")



# Create DataFrame
df = pd.DataFrame({
    'filename': filenames,
    'embedding': embeddings
})

print(f"Successfully processed {len(df)} images")
display(df.head())

Processing 500 images for embeddings...
Processed 100/500 images
Processed 200/500 images
Processed 300/500 images
Processed 400/500 images
Processed 500/500 images
Successfully processed 500 images

	filename	embedding
0	train\Image_6438.jpg	[-0.10860812664031982, 0.46936142444610596, -0...
1	train\Image_32.jpg	[-0.29819875955581665, 0.053064510226249695, -...
2	train\Image_4834.jpg	[-0.1498401165008545, -0.14732249081134796, -0...
3	train\Image_3971.jpg	[-0.21588081121444702, 0.004175275564193726, -...
4	train\Image_3736.jpg	[-0.17965900897979736, 0.013487622141838074, -...

Create 3D UMAP Visualization from Embeddings#

# Convert embeddings list to numpy array matrix
embedding_matrix = np.stack(df['embedding'].values)

print(f"Embedding matrix shape: {embedding_matrix.shape}")
print("Creating 3D UMAP coordinates...")

# Create 3D UMAP
umap_reducer = UMAP(n_components=3, random_state=42)
umap_coords_actual = umap_reducer.fit_transform(embedding_matrix)

# Add UMAP coordinates to dataframe
df['x'] = umap_coords_actual[:, 0]
df['y'] = umap_coords_actual[:, 1]
df['z'] = umap_coords_actual[:, 2]

print("UMAP coordinates created successfully")
print(f"X range: {df['x'].min():.2f} to {df['x'].max():.2f}")
print(f"Y range: {df['y'].min():.2f} to {df['y'].max():.2f}")
print(f"Z range: {df['z'].min():.2f} to {df['z'].max():.2f}")
display(df.head())

Embedding matrix shape: (500, 768)
Creating 3D UMAP coordinates...

C:\structure\code\rag-2026\.venv\Lib\site-packages\umap\umap_.py:1952: UserWarning: n_jobs value 1 overridden to 1 by setting random_state. Use no seed for parallelism.
  warn(

UMAP coordinates created successfully
X range: 2.09 to 8.42
Y range: 7.93 to 14.67
Z range: 6.17 to 11.60

	filename	embedding	x	y	z
0	train\Image_6438.jpg	[-0.10860812664031982, 0.46936142444610596, -0...	6.298563	9.551894	10.315451
1	train\Image_32.jpg	[-0.29819875955581665, 0.053064510226249695, -...	3.993875	11.151982	9.934458
2	train\Image_4834.jpg	[-0.1498401165008545, -0.14732249081134796, -0...	2.789860	12.051387	6.814284
3	train\Image_3971.jpg	[-0.21588081121444702, 0.004175275564193726, -...	3.978969	11.955753	10.324073
4	train\Image_3736.jpg	[-0.17965900897979736, 0.013487622141838074, -...	5.000830	13.291059	9.683495

Save Results to CSV File#

# Save the dataframe to CSV
output_path = base_dir / "data.csv"

# Convert embedding arrays to strings for CSV storage
df_to_save = df.copy()
df_to_save['embedding'] = df_to_save['embedding'].apply(lambda x: ','.join(map(str, x)))

df_to_save.to_csv(output_path, index=False)

print(f"Results saved to {output_path}")
print(f"Final dataframe shape: {df.shape}")
print(f"Columns: {list(df.columns)}")
display(df[['filename', 'x', 'y', 'z']].head(10))

Results saved to data.csv
Final dataframe shape: (500, 5)
Columns: ['filename', 'embedding', 'x', 'y', 'z']

	filename	x	y	z
0	train\Image_6438.jpg	6.298563	9.551894	10.315451
1	train\Image_32.jpg	3.993875	11.151982	9.934458
2	train\Image_4834.jpg	2.789860	12.051387	6.814284
3	train\Image_3971.jpg	3.978969	11.955753	10.324073
4	train\Image_3736.jpg	5.000830	13.291059	9.683495
5	train\Image_3534.jpg	4.364135	12.640450	8.587366
6	test\Image_2716.jpg	8.276859	11.268472	9.252500
7	train\Image_6140.jpg	7.159472	11.854764	10.952312
8	train\Image_4421.jpg	5.640723	12.595654	9.444682
9	train\Image_1690.jpg	5.330312	10.914335	9.631583