jn-v2-ep / app.py
tejastake's picture
Upload 5 files
b866409 verified
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 Get_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 = Get_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_dense_embd(final_image)
return embeddings