Update app.py

app.py

@@ -1,38 +1,60 @@
-from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi import FastAPI, UploadFile, File, Form
 from transformers import CLIPProcessor, CLIPModel
-from PIL import Image, UnidentifiedImageError
+from PIL import Image
 import torch
 import io
+import uuid
+import chromadb
+from chromadb.config import Settings
 
+# Initialize FastAPI
 app = FastAPI()
 
-# Load the CLIP model and processor
-model = CLIPModel.from_pretrained("laion/CLIP-ViT-B-32-laion2B-s34B-b79K")
-processor = CLIPProcessor.from_pretrained("laion/CLIP-ViT-B-32-laion2B-s34B-b79K")
+# Load CLIP model and processor
+model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
 
-@app.get("/")
-def home():
-    return {"message": "CLIP FastAPI app is running!"}
+# Initialize ChromaDB
+chroma_client = chromadb.Client(Settings(chroma_db_impl="duckdb+parquet", persist_directory="./chroma_storage"))
+collection = chroma_client.get_or_create_collection(name="images")
 
-@app.post("/search/")
-async def search_image(file: UploadFile = File(...), query: str = "a photo"):
-    try:
-        # Read and decode image
-        contents = await file.read()
-        image = Image.open(io.BytesIO(contents)).convert("RGB")
-    except UnidentifiedImageError:
-        raise HTTPException(status_code=400, detail="Invalid image file format.")
-
-    # Preprocess image and text
-    inputs = processor(text=[query], images=image, return_tensors="pt", padding=True)
-
-    # Forward pass through the model
+# Function to extract image embeddings
+def get_image_embedding(image: Image.Image):
+    inputs = processor(images=image, return_tensors="pt")
+    with torch.no_grad():
+        embeddings = model.get_image_features(**inputs)
+        embeddings = embeddings / embeddings.norm(p=2, dim=-1, keepdim=True)
+    return embeddings[0].tolist()
+
+# Function to extract text embeddings
+def get_text_embedding(text: str):
+    inputs = processor(text=[text], return_tensors="pt", padding=True)
     with torch.no_grad():
-        outputs = model(**inputs)
-        logits_per_image = outputs.logits_per_image
-        probs = logits_per_image.softmax(dim=1)
+        embeddings = model.get_text_features(**inputs)
+        embeddings = embeddings / embeddings.norm(p=2, dim=-1, keepdim=True)
+    return embeddings[0].tolist()
 
+@app.get("/")
+def root():
+    return {"message": "CLIP + ChromaDB image-text similarity search"}
+
+@app.post("/add-image/")
+async def add_image(file: UploadFile = File(...), label: str = Form(...)):
+    contents = await file.read()
+    image = Image.open(io.BytesIO(contents)).convert("RGB")
+    embedding = get_image_embedding(image)
+    uid = str(uuid.uuid4())
+    collection.add(documents=[label], embeddings=[embedding], ids=[uid], metadatas=[{"label": label}])
+    return {"message": f"Image '{label}' added with ID {uid}"}
+
+@app.post("/search/")
+async def search_text(query: str = Form(...), top_k: int = 3):
+    embedding = get_text_embedding(query)
+    results = collection.query(query_embeddings=[embedding], n_results=top_k)
     return {
         "query": query,
-        "match_confidence": float(probs[0][0])
+        "results": [
+            {"label": doc, "score": score}
+            for doc, score in zip(results["documents"][0], results["distances"][0])
+        ]
     }