Create chromedb_service.py

chromedb_service.py

@@ -0,0 +1,88 @@
+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