moPharma1/computer-vision-backend
0
1import argparse2import base643import io4import os5import re6import sys7import traceback8import uuid9from typing import List, Optional10from transformers import AutoProcessor, AutoModelForZeroShotImageClassification11from transformers import CLIPProcessor, CLIPModel12 13import cv214import numpy as np15import pandas as pd16import pinecone17import pyiqa18import timm19import torch20import uvicorn21from dotenv import load_dotenv22from fastapi import FastAPI, File, Form, HTTPException, UploadFile23from PIL import Image, ImageEnhance24from pydantic import BaseModel25from sentence_transformers import SentenceTransformer, util26 27load_dotenv()28pinecone.init(api_key=os.getenv("PINECONE_KEY"), environment=os.getenv("PINECONE_ENV"))29 30IMAGE_SIMILARITY_DEMO = "/find-similar-image/"31IMAGE_SIMILARITY_PINECONE_DEMO = "/find-similar-image-pinecone/"32INDEX_NAME = "imagesearch-demo"33INDEX_DIMENSION = 51234TMP_DIR = "tmp"35model_type = "huggingface"36# model_type = "sentence_transformers"37processor = AutoProcessor.from_pretrained("patrickjohncyh/fashion-clip")38model = CLIPModel.from_pretrained("patrickjohncyh/fashion-clip")39 40image_sim_model = SentenceTransformer("clip-ViT-B-32")41 42def enhance_image(pil_image):43 # Convert PIL Image to OpenCV format44 open_cv_image = np.array(pil_image)45 # Convert RGB to BGR46 open_cv_image = open_cv_image[:, :, ::-1].copy()47 48 # Convert to grayscale49 gray = cv2.cvtColor(open_cv_image, cv2.COLOR_BGR2GRAY)50 51 # Histogram equalization52 equ = cv2.equalizeHist(gray)53 54 # Adaptive Histogram Equalization55 clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))56 adaptive_hist_eq = clahe.apply(gray)57 58 # Gaussian Blurring59 gaussian_blurred = cv2.GaussianBlur(adaptive_hist_eq, (5,5), 0)60 61 # Noise reduction62 denoised = cv2.medianBlur(gaussian_blurred, 3)63 64 # Brightness & Contrast adjustment65 lab = cv2.cvtColor(open_cv_image, cv2.COLOR_BGR2Lab)66 l, a, b = cv2.split(lab)67 cl = clahe.apply(l)68 limg = cv2.merge((cl, a, b))69 enhanced_image = cv2.cvtColor(limg, cv2.COLOR_Lab2BGR)70 71 # Convert back to PIL Image72 enhanced_pil_image = Image.fromarray(cv2.cvtColor(enhanced_image, cv2.COLOR_BGR2RGB))73 74 # IMAGE AUGMENTATION75 # For demonstration purposes, let's do a simple brightness adjustment.76 # In practice, choose the augmentations that suit your task.77 enhancer = ImageEnhance.Brightness(enhanced_pil_image)78 enhanced_pil_image = enhancer.enhance(1.2) # Brighten the image by 20%79 80 return enhanced_pil_image81 82 83print("checking pinecone Index")84if INDEX_NAME not in pinecone.list_indexes():85 # delete the current index and create the new index if it does not exist86 for delete_index in pinecone.list_indexes():87 print(f"Deleting exitsing pinecone Index : {delete_index}")88 89 pinecone.delete_index(delete_index)90 print(f"Creating new pinecone Index : {INDEX_NAME}")91 pinecone.create_index(INDEX_NAME, dimension=INDEX_DIMENSION, metric="cosine")92 93print("Connecting to Pinecone Index")94index = pinecone.Index(INDEX_NAME)95 96 97device = "cuda" if torch.cuda.is_available() else "cpu"98 99os.makedirs(TMP_DIR, exist_ok=True)100 101 102device = "cuda" if torch.cuda.is_available() else "cpu"103 104 105os.makedirs(TMP_DIR, exist_ok=True)106 107 108 109 110app = FastAPI(title="CV Demos")111 112 113# define response114@app.get("/")115def root_route():116 return {"error": f"Use GET {IMAGE_SIMILARITY_PINECONE_DEMO} instead of the root route!"}117 118 119@app.post(IMAGE_SIMILARITY_DEMO)120async def image_search_local(121 images_to_search: List[UploadFile], query_image: UploadFile = File(...), top_k: int = 5, 122):123 print(124 f"Recived images of length: {len(images_to_search)} needs to retrieve top k : {top_k} similar images as result"125 )126 try:127 extension = query_image.filename.split(".")[-1] in ("jpg", "jpeg", "png")128 search_images = []129 search_filenames = []130 print("Processing request...")131 for image in images_to_search:132 if image.filename.split(".")[-1] not in ("jpg", "jpeg", "png"):133 return "Image must be jpg or png format!"134 # read image contain135 search_filenames.append(image.filename)136 contents = await image.read()137 search_images.append(Image.open(io.BytesIO(contents)))138 print("Indexing images to search...")139 140 corpus_embeddings = image_sim_model.encode(141 search_images, convert_to_tensor=True, show_progress_bar=True142 )143 if not extension:144 return "Image must be jpg or png format!"145 # read image contain146 contents = await query_image.read()147 query_image = Image.open(io.BytesIO(contents))148 print("Indexing query image...")149 150 prompt_embedding = image_sim_model.encode(query_image, convert_to_tensor=True)151 print("Searching query image...")152 153 hits = util.semantic_search(prompt_embedding, corpus_embeddings, top_k=top_k)154 # hits = pd.DataFrame(hits[0], columns=['corpus_id', 'score'])155 # tmp_file = f"{TMP_DIR}/tmp.png"156 # pil_image.save(tmp_file)157 # answer_git_large = generate_answer_git(git_processor_large, git_model_large, image, question)158 print("Creating the result..")159 similar_images = []160 print("hits ", hits)161 for hit in hits[0]:162 # print("Finding the image ")163 # print("Type of images list ", type(search_images), "similar image id ", hit['corpus_id'])164 open_cv_image = np.array(search_images[hit["corpus_id"]].convert("RGB"))[:, :, ::-1]165 # print("cv2.imencode the image ")166 _, encoded_img = cv2.imencode(".PNG", open_cv_image)167 # print("base64 the image ")168 encoded_img = base64.b64encode(encoded_img)169 # print("Appending the image ")170 similar_images.append(171 {172 "filename": search_filenames[hit["corpus_id"]],173 "dimensions": str(open_cv_image.shape),174 "score": hit["score"],175 "encoded_img": encoded_img,176 }177 )178 print("Sending result..")179 180 return {"similar_images": similar_images}181 182 except:183 e = sys.exc_info()[1]184 raise HTTPException(status_code=500, detail=str(e))185 186 187def get_clip_vectors(image):188 img_inputs = processor(images=image, return_tensors="pt")189 img_emb = model.get_image_features(**img_inputs)190 191 return img_emb192 193@app.post(IMAGE_SIMILARITY_PINECONE_DEMO)194async def image_search_pinecone(195 images_to_search: Optional[List[UploadFile]] = File(None),196 query_image: Optional[UploadFile] = File(None),197 top_k: int = 5,198 namespace="av_local",199 action="query",200 model_type= "huggingface"201):202 203 try:204 # Function to delete all files from the database205 print(f"Received request with images_to_search: {images_to_search} and query_image: {query_image} with action: {action}")206 if action == "delete":207 index = pinecone.Index(INDEX_NAME)208 delete_response = index.delete(delete_all=True, namespace=namespace)209 return {f"Deleted the namespace: {namespace}": delete_response}210 211 elif action == "query" and query_image is not None:212 extension = query_image.filename.split(".")[-1] in ("jpg", "jpeg", "png", "JPG", "PNG", "JPEG")213 if not extension:214 return "Image must be jpg or png format!"215 # read image contain216 contents = await query_image.read()217 query_image = Image.open(io.BytesIO(contents))218 print("Indexing query image...")219 query_image = enhance_image(query_image)220 if model_type =="huggingface":221 prompt_embedding = get_clip_vectors(query_image).tolist()222 else:223 prompt_embedding = image_sim_model.encode(query_image, convert_to_tensor=True).tolist()224 if INDEX_NAME not in pinecone.list_indexes():225 return {"similar_images": [], "status": "No index found for images"}226 227 else:228 index = pinecone.Index(INDEX_NAME)229 query_response = index.query(230 namespace=namespace,231 top_k=top_k,232 include_values=True,233 include_metadata=True,234 vector=prompt_embedding,235 )236 result_images = [d["metadata"]["file_path"] for d in query_response["matches"]]237 print("Creating the result..")238 similar_images = []239 print("Retrieved matches ", query_response["matches"])240 for file_path in result_images:241 try:242 # print("Finding the image ")243 # print("Type of images list ", type(search_images), "similar image id ", hit['corpus_id'])244 open_cv_image = cv2.imread(file_path)245 # print("cv2.imencode the image ")246 _, encoded_img = cv2.imencode(".PNG", open_cv_image)247 # print("base64 the image ")248 encoded_img = base64.b64encode(encoded_img)249 # print("Appending the image ")250 similar_images.append(251 {252 "filename": file_path,253 "dimensions": str(open_cv_image.shape),254 "score": 0,255 "encoded_img": encoded_img,256 }257 )258 except:259 similar_images.append(260 {261 "filename": file_path,262 "dimensions": None,263 "score": 0,264 "encoded_img": None,265 }266 )267 print("Sending result..")268 269 return {"similar_images": similar_images}270 271 elif action == "index" and (images_to_search is not None) and (len(images_to_search) > 0):272 print(273 f"Recived images of length: {len(images_to_search)} needs to retrieve top k : {top_k} similar images as result"274 )275 print(f"Action indexing is executing for : {len(images_to_search)} images")276 # if the index does not already exist, we create it277 # check if the abstractive-question-answering index exists278 print("checking pinecone Index")279 if INDEX_NAME not in pinecone.list_indexes():280 # delete the current index and create the new index if it does not exist281 for delete_index in pinecone.list_indexes():282 print(f"Deleting exitsing pinecone Index : {delete_index}")283 284 pinecone.delete_index(delete_index)285 print(f"Creating new pinecone Index : {INDEX_NAME}")286 pinecone.create_index(INDEX_NAME, dimension=INDEX_DIMENSION, metric="cosine")287 # instantiate connection to your Pinecone index288 print(f"Connecting to pinecone Index : {INDEX_NAME}")289 index = pinecone.Index(INDEX_NAME)290 search_images = []291 meta_datas = []292 ids = []293 print("Processing request...")294 for image in images_to_search:295 if image.filename.split(".")[-1] not in ("jpg", "jpeg", "png", "JPG", "PNG", "JPEG"):296 return "Image must be jpg or png format!"297 # read image contain298 contents = await image.read()299 pil_image = Image.open(io.BytesIO(contents))300 # pil_image = enhance_image(pil_image)301 tmp_file = f"{TMP_DIR}/{image.filename}"302 pil_image.save(tmp_file)303 meta_datas.append({"file_path": tmp_file})304 search_images.append(pil_image)305 ids.append(str(uuid.uuid1()).replace("-",""))306 307 print("Encoding images to vectors...")308 if model_type =="huggingface":309 corpus_embeddings = get_clip_vectors(search_images).tolist()310 else:311 corpus_embeddings = image_sim_model.encode(312 search_images, convert_to_tensor=True, show_progress_bar=True313 ).tolist()314 print(f"Indexing images to pinecone Index : {INDEX_NAME}")315 index.upsert(316 vectors=list(zip(ids, corpus_embeddings, meta_datas)), namespace=namespace317 )318 319 320 return {"similar_images": [], "status": "Indexing succesfull for uploaded files"}321 else:322 return {"similar_images": []}323 except Exception as e:324 e = sys.exc_info()[1]325 print(f"exception happened {e} {str(traceback.print_exc())}")326 raise HTTPException(status_code=500, detail=str(e))327 328 329if __name__ == "__main__":330 parser = argparse.ArgumentParser(description="Fast API exposing YOLOv5 model")331 parser.add_argument("--port", default=8000, type=int, help="port number")332 # parser.add_argument('--model', nargs='+', default=['yolov5s'], help='model(s) to run, i.e. --model yolov5n yolov5s')333 opt = parser.parse_args()334 uvicorn.run(app, port=opt.port)