image_embedder.py

import torch
import timm
from PIL import Image
import os
import json
from tqdm import tqdm
from neo4j import GraphDatabase
from dotenv import load_dotenv
import numpy as np
import unicodedata

load_dotenv()
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")

# Normalizza i nomi rimuovendo accenti e convertendo in minuscolo
def normalize_name(name):
    if not isinstance(name, str):
        return ""

    nfkd_form = unicodedata.normalize('NFKD', name.lower())
    return "".join([c for c in nfkd_form if not unicodedata.combining(c)])

def setup():
    # num_classes=0 per ottenere l'embedding senza il layer di classificazione
    model = timm.create_model('efficientnet_b0', pretrained=True, num_classes=0).to(device)
    model.eval() # modalità di valutazione (disabilitiamo dropout e batchnorm che servono solo in training)

    data_config = timm.data.resolve_model_data_config(model)
    # timm crea automaticamente la pipeline di preprocessing corretta per il modello
    transform = timm.data.create_transform(**data_config, is_training=False)
    print("Image model e transform settati.")

    uri = os.getenv("NEO4J_URI")
    user = os.getenv("NEO4J_USER")
    password = os.getenv("NEO4J_PASSWORD")
    driver = GraphDatabase.driver(uri, auth=(user, password))
    print("Connessione a Neo4j stabilita.")

    return model, transform, driver

def get_players_from_neo4j(driver):
    print("Recupero dei giocatori da Neo4j...")

    with driver.session(database=os.getenv("NEO4J_DATABASE")) as session:
        results = session.run("""
                MATCH (p:Player)
                RETURN p.playerId AS playerId, p.givenName AS givenName, p.familyName AS familyName
            """)
        player_map = {
            normalize_name(f"{r['givenName']} {r['familyName']}"): r['playerId']
            for r in results
        }
        print(f"Recuperati {len(player_map)} giocatori.")

        return player_map

def extract_image_embedding(image_path, model, transform):
    """
    Calcola l'embedding di un'immagine data il suo percorso.
    """
    img = Image.open(image_path).convert('RGB')
    # applichiamo la trasformazione e aggiungiamo una dimensione batch (dato che PyTorch si aspetta un batch di immagini)
    img_tensor = transform(img).unsqueeze(0).to(device)

    # Non calcoliamo i gradienti, dato che siamo in fase di inferenza
    with torch.no_grad():
        # Calcoliamo l'embedding (tolist serve per convertire il tensore in una lista Python, facile da serializzare in JSON)
        embedding = model(img_tensor).squeeze().cpu().numpy().tolist()
    return embedding

if __name__ == "__main__":
    model, transform, driver = setup()
    player_name_to_id_map = get_players_from_neo4j(driver)

    image_dir = "player_images"

    image_embeddings = {}

    for group_folder in tqdm(os.listdir(image_dir), desc="Processando Gruppi"):
        group_path = os.path.join(image_dir, group_folder)

        if os.path.isdir(group_path) and group_folder.startswith("Group"):

            for country_folder in os.listdir(group_path):
                country_path = os.path.join(group_path, country_folder)

                if os.path.isdir(country_path):

                    for image_folder in os.listdir(country_path):

                        image_folder_path = os.path.join(country_path, image_folder)

                        # Estraiamo il nome del giocatore dalla cartella
                        player_clean_name = image_folder.replace("Images_", "").replace("_", " ").strip()
                        normalized_player_name = normalize_name(player_clean_name)

                        # Cerchiamo l'ID del giocatore nella mappa creata in precedenza
                        # Questo è il punto di collegamento tra i file e il database
                        player_id = player_name_to_id_map.get(normalized_player_name)

                        if player_id:

                            for image_filename in os.listdir(image_folder_path):
                                if image_filename.lower().endswith(('.png', '.jpg', '.jpeg')):

                                    # Calcoliamo l'embedding dell'immagine
                                    embedding = extract_image_embedding(os.path.join(image_folder_path, image_filename), model, transform)

                                    if embedding is not None:
                                        if player_id not in image_embeddings:
                                            image_embeddings[player_id] = []

                                        # Salviamo l'embedding insieme al percorso della cartella
                                        image_embeddings[player_id].append({
                                            "image_folder": image_folder_path,
                                            "embedding": embedding
                                        })

    final_image_embeddings = {}
    print("Calcolo degli embeddings medi...")

    # Per ogni giocatore, calcoliamo l'embedding medio delle sue immagini
    for player_id, embeddings in tqdm(image_embeddings.items(), desc="Calcolando embedding medi"):
        if embeddings:
            # Estrai solo gli embeddings dai dizionari
            embedding_vectors = [item["embedding"] for item in embeddings]
            # Prendi il primo image_folder (dovrebbe essere lo stesso per tutte le immagini del giocatore)
            image_folder = embeddings[0]["image_folder"]
            avg_embedding = np.mean(np.array(embedding_vectors), axis=0).tolist()
            final_image_embeddings[player_id] = {
                "embedding": avg_embedding,
                "image_folder": image_folder
            }

    with open("image_embeddings_avg.json", "w", encoding="utf-8") as f:
        json.dump(final_image_embeddings, f, ensure_ascii=False, indent=2)

    print("Embeddings delle immagini salvati in image_embeddings_avg.json")