Spaces:
Running
Running
File size: 3,571 Bytes
4e54f0a 1a2d238 4e54f0a 1a2d238 4e54f0a 1a2d238 4e54f0a 1a2d238 4e54f0a 1a2d238 4e54f0a e326e89 4e54f0a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 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, 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],
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
|