Upload 5 files

Dockerfile

@@ -0,0 +1,22 @@
+FROM python:3.9
+
+WORKDIR /app
+
+COPY . /app
+
+RUN pip3 install fastapi uvicorn transformers==4.42.3 pillow protobuf==4.25.3 fastapi-health boto3
+
+RUN useradd -m -u 1000 user
+
+USER user
+
+ENV HOME=/home/user \
+ PATH=/home/user/.local/bin:$PATH
+
+WORKDIR $HOME/app
+
+COPY --chown=user . $HOME/app
+
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]
+
+#uvicorn app:app --host 0.0.0.0 --port 7860 

app.py

@@ -0,0 +1,47 @@
+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

process_img.py

@@ -0,0 +1,49 @@
+import numpy as np
+from PIL import Image, ImageOps
+import logging
+
+class Image_Processor:
+    def __init__(self):
+        pass
+    def is_image_white_by_percentage(self,image_path, white_threshold):
+        image = image_path.convert('RGB')
+        image_np = np.array(image)
+        white_pixel = np.array([255, 255, 255])
+        white_pixels_count = np.sum(np.all(image_np == white_pixel, axis=-1))
+        total_pixels = image_np.shape[0] * image_np.shape[1]
+        white_pixel_percentage = (white_pixels_count / total_pixels) * 100
+        return white_pixel_percentage > white_threshold
+
+    def padding_white(self,image, output_size=(512, 512)):
+        # Ensure image is in RGB mode before padding
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        new_image = ImageOps.pad(image, output_size, method=Image.Resampling.LANCZOS, color=(255, 255, 255))
+        return new_image
+
+    def resize_image_with_aspect_ratio(self,img):
+        target_size=512
+        width, height = img.size
+        original_aspect_ratio = width / height
+        if width > height:
+            new_width = target_size
+            new_height = int(target_size / original_aspect_ratio)
+        else:
+            new_height = target_size
+            new_width = int(target_size * original_aspect_ratio)
+        resized_img = img.resize((new_width, new_height))
+        return resized_img
+    
+    def get_processed_img(self,image):
+        white_thresh = self.is_image_white_by_percentage(image,50)
+        if white_thresh == True:
+            resized_image = self.resize_image_with_aspect_ratio(image)
+            final_image = self.padding_white(resized_image)
+            logging.info('Resized and Padded Image')
+        else:
+            #final_image = self.resize_image_with_aspect_ratio(image)
+            final_image = image.resize((512,512))
+            logging.info('Resized Image')
+        
+        final_image = final_image.convert('L') if final_image.mode != 'L' else final_image
+        return final_image

requirements.txt

@@ -0,0 +1,7 @@
+fastapi
+uvicorn
+pillow
+torch
+transformers
+fastapi-health
+boto3

vector_emb.py

@@ -0,0 +1,32 @@
+from sentence_transformers import SentenceTransformer
+import torch.nn.functional as F
+import torch
+
+class Load_EmbeddingModels:
+    def __init__(self, model_name ='jinaai/jina-clip-v2'):
+        self.device = "cuda" if torch.cuda.is_available() else "cpu"
+        self.img_model_ID = model_name 
+        self.img_model = self.get_cp_img_model_info(self.img_model_ID)
+        
+    def get_cp_img_model_info(self, model_name):
+        print('Loading SentenceTransformer model')    
+        model =SentenceTransformer(model_name, trust_remote_code=True)
+        model = model.to(self.device)
+        return model
+
+    def get_single_image_embedding_cp_im(self, my_image):    
+        embedding = self.img_model.encode(
+            my_image,
+            normalize_embeddings=True
+        )
+        values = embedding.tolist()
+        return values
+
+class Get_EmbeddingModels:
+    def __init__(self, model_name='jinaai/jina-clip-v2'):
+        self.embed_model = Load_EmbeddingModels(model_name)
+    
+    def get_dense_embd(self, img):
+        embd = self.embed_model.get_single_image_embedding_cp_im(img)
+        return embd
+