Update app.py

app.py

@@ -35,10 +35,13 @@ async def extract_features(file: UploadFile = File(...)):
         image = Image.open(io.BytesIO(contents)).convert("RGB")
 
         # Get feature embeddings
-        features = feature_extraction_pipeline(image)
+        features = feature_extraction_pipeline(image)  # Shape: (1, 577, 768)
+        
+        # Extract CLS token embedding
+        cls_embedding = np.array(features)[0, 0, :]  # Shape: (768,)
         
         # Return the embedding vector
-        return JSONResponse(content={"features": features})
+        return JSONResponse(content={"features": cls_embedding.tolist()})
     
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))