cnn.py

import os
from PIL import Image
import torch
from transformers import CLIPProcessor, CLIPModel
from sklearn.preprocessing import normalize
from sklearn.metrics.pairwise import cosine_distances
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns

IMAGES_DIR = "imagenes"
VALID_EXTENSIONS = (".jpg", ".jpeg", ".png", ".avif", ".webp", ".bmp", ".tiff")

# Configure GPU device if available
device = "cuda" if torch.cuda.is_available() else "cpu"
print("Using device:", device)

# Load CLIP model
model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32").to(device)
processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")

# Helper functions
def load_images(path):
    return [f for f in os.listdir(path) if f.lower().endswith(VALID_EXTENSIONS)]

def get_embeddings(images):
    embeddings = []
    names = []
    for img_name in images:
        try:
            img_path = os.path.join(IMAGES_DIR, img_name)
            image = Image.open(img_path).convert("RGB")
            inputs = processor(images=image, return_tensors="pt").to(device)
            with torch.no_grad():
                emb = model.get_image_features(**inputs)
            embeddings.append(emb.cpu().numpy()[0])
            names.append(img_name)
        except Exception as e:
            print(f"Error processing {img_name}: {e}")
    embeddings = normalize(np.array(embeddings))  # important to normalize for cosine
    return embeddings, names

# Load images and generate embeddings
images = load_images(IMAGES_DIR)
if not images:
    print("No valid images found.")
    exit()

print(f"Processing {len(images)} images...")
embeddings, names = get_embeddings(images)

# Calculate cosine distance matrix
dist_matrix = cosine_distances(embeddings)

# Print statistics
print("Minimum distance:", dist_matrix.min())
print("Maximum distance:", dist_matrix.max())
print("Average distance:", dist_matrix.mean())

# Draw heatmap
plt.figure(figsize=(10, 8))
sns.heatmap(dist_matrix, xticklabels=names, yticklabels=names, cmap="viridis", annot=False)
plt.title("Cosine Distance Map Between Images (CLIP embeddings)")
plt.xticks(rotation=90)
plt.yticks(rotation=0)
plt.tight_layout()
plt.show()