|
|
from fastapi import FastAPI, Depends, HTTPException
|
|
|
from pydantic import BaseModel
|
|
|
from fastapi_health import health
|
|
|
from PIL import Image
|
|
|
import logging
|
|
|
import sys
|
|
|
from io import BytesIO
|
|
|
import base64
|
|
|
|
|
|
from process_img import Image_Processor
|
|
|
from vector_emb import EmbeddingModels
|
|
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
logging.basicConfig(
|
|
|
level=logging.getLevelName("INFO"),
|
|
|
handlers=[logging.StreamHandler(sys.stdout)],
|
|
|
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s")
|
|
|
logging.info('Logging module started')
|
|
|
|
|
|
def get_session():
|
|
|
return True
|
|
|
|
|
|
def is_database_online(session: bool = Depends(get_session)):
|
|
|
return session
|
|
|
|
|
|
app = FastAPI()
|
|
|
app.add_api_route("/healthz", health([is_database_online]))
|
|
|
|
|
|
|
|
|
|
|
|
model = EmbeddingModels()
|
|
|
img_Processor = Image_Processor()
|
|
|
|
|
|
class ImageBase64(BaseModel):
|
|
|
base64_string: str
|
|
|
|
|
|
class TextInput(BaseModel):
|
|
|
text: str
|
|
|
|
|
|
@app.post("/design-dense/")
|
|
|
async def embed_image(data: ImageBase64):
|
|
|
base64_string = data.base64_string
|
|
|
image_data = base64.b64decode(base64_string)
|
|
|
image = Image.open(BytesIO(image_data))
|
|
|
final_image = img_Processor.get_processed_img(image)
|
|
|
embeddings = model.get_single_image_embedding(final_image)
|
|
|
return embeddings
|
|
|
|
|
|
|
|
|
@app.post("/sparse/")
|
|
|
async def embed_text(item: TextInput):
|
|
|
try:
|
|
|
logging.info(f'Received text for embedding: {item.text}')
|
|
|
embeddings = model.get_single_sparse_text_embedding(item.text)
|
|
|
logging.info('Embedding process completed')
|
|
|
return embeddings
|
|
|
except Exception as e:
|
|
|
logging.error(f'Error during embedding process: {e}')
|
|
|
raise HTTPException(status_code=500, detail=str(e))
|
|
|
|
|
|
|
|
|
@app.post("/design-sparse/")
|
|
|
async def embed_text(item: TextInput):
|
|
|
try:
|
|
|
logging.info(f'Received text for embedding: {item.text}')
|
|
|
embeddings = model.get_single_sparse_text_embedding(item.text)
|
|
|
embeddings = model.normalize_sparse_vector_values(embeddings)
|
|
|
logging.info('Embedding process completed')
|
|
|
return embeddings
|
|
|
except Exception as e:
|
|
|
logging.error(f'Error during embedding process: {e}')
|
|
|
raise HTTPException(status_code=500, detail=str(e)) |
