Spaces:

ravi19
/

FastAPI

No application file

FastAPI / api /routes.py

Deploy FastAPI to HF Space

b36cb8b over 1 year ago

3.66 kB

"""
Routes for the Image Similarity Search API
Contains all endpoints for the application using your original route implementation
"""

import uuid
import base64
import io
from typing import List, Optional
from fastapi import APIRouter, FastAPI, File, UploadFile, Form, Query, Path # type: ignore
from pydantic import BaseModel
from PIL import Image

from services.embedding_service import ImageEmbeddingModel
from services.vector_db_service import VectorDatabaseClient


class Base64ImageRequest(BaseModel):
    """Request model for base64 encoded images"""
    image_data: str


def register_routes(
    app: FastAPI,
    embedding_model: ImageEmbeddingModel,
    vector_db: VectorDatabaseClient,
):
    """Register all routes with the FastAPI app"""

    @app.api_route("/", methods=["GET", "HEAD"])
    async def read_root():
        return {"status": "API running"}
    
    @app.post("/add-image/")
    async def add_image(
        file: UploadFile = File(...),
        item_name: str = Form(...),
        design_name: str = Form(...),
        item_price: float = Form(...)
    ):
        """Upload an image with product details and store its embedding"""
        # Process the image to get embedding
        # image_data = await file.read()
        embedding = await embedding_model.get_embedding_from_upload(file)
        
        # Generate a unique ID
        image_id = str(uuid.uuid4())
        
        # Store additional metadata in payload
        payload = {
            "filename": file.filename,
            "item_name": item_name,
            "design_name": design_name,
            "item_price": item_price
        }
        
        # Store in vector database
        vector_db.add_image(image_id, embedding, payload)
        
        return {"message": "Image added successfully", "id": image_id}
    
    @app.post("/add-images-from-folder/")
    async def add_images_from_folder(folder_path: str):
        """Process and add all images from a specified folder"""
        embeddings = embedding_model.get_embeddings_from_folder(folder_path)
        return {"embeddings": embeddings}
    
    @app.post("/search-by-image/")
    async def search_by_image(file: UploadFile = File(...)):
        """Search for similar images by uploading a file"""
        # Process the image to get embedding
        # image_data = await file.read()
        embedding = await embedding_model.get_embedding_from_upload(file)
        
        # Search using the embedding
        results = vector_db.search_by_vector(embedding, limit=1)
        
        # return [
        #     {
        #         "id": r.id,
        #         "score": r.score,
        #         "payload": r.payload
        #     }
        #     for r in results
        # ]
        return results
    
    @app.post("/search-by-image-scan/")
    async def search_by_image_scan(request: Base64ImageRequest):
        """Search for similar images using a base64 encoded image"""
        # Decode base64 image
        image_data = request.image_data
        image_bytes = base64.b64decode(image_data.split(',')[1] if ',' in image_data else image_data)
        
        # Convert to PIL Image
        image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
        
        # Process image to get embedding
        embedding = embedding_model.get_embedding_from_pil(image)
        
        # Search using the embedding
        results = vector_db.search_by_vector(embedding, limit=1)
        
        return results
    
    @app.get("/collections")
    def list_collections():
        """List all available collections in the vector database"""
        return vector_db.list_collections()

	"""
	Routes for the Image Similarity Search API
	Contains all endpoints for the application using your original route implementation
	"""

	import uuid
	import base64
	import io
	from typing import List, Optional
	from fastapi import APIRouter, FastAPI, File, UploadFile, Form, Query, Path # type: ignore
	from pydantic import BaseModel
	from PIL import Image

	from services.embedding_service import ImageEmbeddingModel
	from services.vector_db_service import VectorDatabaseClient


	class Base64ImageRequest(BaseModel):
	"""Request model for base64 encoded images"""
	image_data: str


	def register_routes(
	app: FastAPI,
	embedding_model: ImageEmbeddingModel,
	vector_db: VectorDatabaseClient,
	):
	"""Register all routes with the FastAPI app"""

	@app.api_route("/", methods=["GET", "HEAD"])
	async def read_root():
	return {"status": "API running"}

	@app.post("/add-image/")
	async def add_image(
	file: UploadFile = File(...),
	item_name: str = Form(...),
	design_name: str = Form(...),
	item_price: float = Form(...)
	):
	"""Upload an image with product details and store its embedding"""
	# Process the image to get embedding
	# image_data = await file.read()
	embedding = await embedding_model.get_embedding_from_upload(file)

	# Generate a unique ID
	image_id = str(uuid.uuid4())

	# Store additional metadata in payload
	payload = {
	"filename": file.filename,
	"item_name": item_name,
	"design_name": design_name,
	"item_price": item_price
	}

	# Store in vector database
	vector_db.add_image(image_id, embedding, payload)

	return {"message": "Image added successfully", "id": image_id}

	@app.post("/add-images-from-folder/")
	async def add_images_from_folder(folder_path: str):
	"""Process and add all images from a specified folder"""
	embeddings = embedding_model.get_embeddings_from_folder(folder_path)
	return {"embeddings": embeddings}

	@app.post("/search-by-image/")
	async def search_by_image(file: UploadFile = File(...)):
	"""Search for similar images by uploading a file"""
	# Process the image to get embedding
	# image_data = await file.read()
	embedding = await embedding_model.get_embedding_from_upload(file)

	# Search using the embedding
	results = vector_db.search_by_vector(embedding, limit=1)

	# return [
	# {
	# "id": r.id,
	# "score": r.score,
	# "payload": r.payload
	# }
	# for r in results
	# ]
	return results

	@app.post("/search-by-image-scan/")
	async def search_by_image_scan(request: Base64ImageRequest):
	"""Search for similar images using a base64 encoded image"""
	# Decode base64 image
	image_data = request.image_data
	image_bytes = base64.b64decode(image_data.split(',')[1] if ',' in image_data else image_data)

	# Convert to PIL Image
	image = Image.open(io.BytesIO(image_bytes)).convert("RGB")

	# Process image to get embedding
	embedding = embedding_model.get_embedding_from_pil(image)

	# Search using the embedding
	results = vector_db.search_by_vector(embedding, limit=1)

	return results

	@app.get("/collections")
	def list_collections():
	"""List all available collections in the vector database"""
	return vector_db.list_collections()