Create app.py

app.py

@@ -0,0 +1,41 @@
+from fastapi import FastAPI, UploadFile, File, HTTPException
+from fastapi.responses import JSONResponse
+from transformers import pipeline
+from PIL import Image
+import torch
+import io
+
+# Set up the device for the model
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+# Load the feature extraction pipeline
+feature_extraction_pipeline = pipeline(
+    task="image-feature-extraction", 
+    model="google/vit-base-patch16-384", 
+    device=0 if torch.cuda.is_available() else -1
+)
+
+# Initialize FastAPI
+app = FastAPI()
+
+# Endpoint to extract image features
+@app.post("/extract-features/")
+async def extract_features(file: UploadFile = File(...)):
+    try:
+        # Validate file format
+        if file.content_type not in ["image/jpeg", "image/png", "image/jpg"]:
+            raise HTTPException(status_code=400, detail="Unsupported file format. Upload a JPEG or PNG image.")
+        
+        # Read image
+        contents = await file.read()
+        image = Image.open(io.BytesIO(contents)).convert("RGB")
+
+        # Get feature embeddings
+        features = feature_extraction_pipeline(image)
+        
+        # Return the embedding vector
+        return JSONResponse(content={"features": features})
+    
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+