chromedb_service.py

import requests
from PIL import Image
import numpy as np
import tensorflow as tf
import chromadb
from chromadb.config import Settings

# Fonction pour télécharger une image depuis une URL
def download_image(url):
    response = requests.get(url)
    if response.status_code != 200:
        raise Exception(f"Erreur lors du téléchargement de l'image : {response.status_code}")
    return Image.open(io.BytesIO(response.content))

# Fonction pour encoder une image en vecteurs à partir d'une URL
def encode_image_from_url(image_url):
    image = download_image(image_url)
    image = image.resize((224, 224))  # Redimensionner à 224x224
    image_array = np.array(image) / 255.0  # Normaliser les valeurs des pixels
    image_tensor = tf.convert_to_tensor(image_array, dtype=tf.float32)
    image_tensor = tf.expand_dims(image_tensor, axis=0)  # Ajouter une dimension pour le batch

    # Charger le modèle MobileNet
    model = tf.keras.applications.MobileNetV2(weights='imagenet', include_top=False, pooling='avg')
    embeddings = model(image_tensor)
    return embeddings.numpy()[0]  # Retourner les vecteurs sous forme de tableau

# Ajouter une image dans ChromaDB
def add_image_to_chroma(collection_name, id, image_url, metadata):
    vector = encode_image_from_url(image_url)
    chroma_client = chromadb.HttpClient(host='https://stable-diffusion-engine.oneiro-lego.com')
    collection = chroma_client.get_or_create_collection(
        name=collection_name, dimension=len(vector)
    )
    collection.add(
        ids=[id],
        embeddings=[vector],
        metadatas=[metadata]
    )
    print(f"Image {image_url} ajoutée avec succès !")

# Ajouter un document dans ChromaDB
def add_document(collection_name, id, text, metadata):
    chroma_client = chromadb.HttpClient(host='https://stable-diffusion-engine.oneiro-lego.com')
    collection = chroma_client.get_or_create_collection(name=collection_name)
    collection.upsert(
        documents=[text],
        ids=[id],
        metadatas=[metadata]
    )
    print(f"Document {id} ajouté avec succès !")

# Supprimer un document dans ChromaDB
def delete_document(collection_name, id):
    chroma_client = chromadb.HttpClient(host='https://stable-diffusion-engine.oneiro-lego.com')
    collection = chroma_client.get_or_create_collection(name=collection_name)
    collection.delete(ids=[id])
    print(f"Document {id} supprimé avec succès !")

# Supprimer une collection dans ChromaDB
def delete_collection(collection_name):
    chroma_client = chromadb.HttpClient(host='https://stable-diffusion-engine.oneiro-lego.com')
    chroma_client.delete_collection(name=collection_name)
    print(f"Collection {collection_name} supprimée avec succès !")

# Recherche dans une collection
def search(collection_name, query, metadata, n_results):
    chroma_client = chromadb.HttpClient(host='https://stable-diffusion-engine.oneiro-lego.com')
    collection = chroma_client.get_or_create_collection(name=collection_name)
    results = collection.query(
        query_texts=query,
        where=metadata,
        n_results=n_results
    )
    return parse_chromadb_response(results)

# Analyse des réponses de ChromaDB
def parse_chromadb_response(response):
    results = [
        {
            "id": response["ids"][0][i],
            "distance": response["distances"][0][i],
            "document": response["documents"][0][i],
            "metadata": response["metadatas"][0][i] if response["metadatas"] and len(response["metadatas"][0]) > i else None
        }
        for i in range(len(response["ids"][0]))
    ]
    return results