Update app.py

app.py

@@ -1,5 +1,3 @@
-import cv2
-import numpy as np
 import argparse
 import base64
 import io
@@ -9,8 +7,7 @@ import sys
 import traceback
 import uuid
 from typing import List, Optional
-from PIL import ImageEnhance
-import traceback
+
 import cv2
 import numpy as np
 import pandas as pd
@@ -21,30 +18,21 @@ import torch
 import uvicorn
 from dotenv import load_dotenv
 from fastapi import FastAPI, File, Form, HTTPException, UploadFile
-from PIL import Image
+from PIL import Image, ImageEnhance
 from pydantic import BaseModel
 from sentence_transformers import SentenceTransformer, util
-from transformers import (
-    AutoFeatureExtractor,
-    AutoModel,
-    DonutProcessor,
-    VisionEncoderDecoderModel,
-)
-from fashion_clip.fashion_clip import FashionCLIP
+
 load_dotenv()
 pinecone.init(api_key=os.getenv("PINECONE_KEY"), environment=os.getenv("PINECONE_ENV"))
-DETECTION_URL = "/object-detection/"
-CLASSIFICATION_URL = "/object-classification/"
-QUALITY_ASSESSMENT_URL = "/quality-assessment/"
-FACE_URL = "/face-anonymization/"
-LICENCE_URL = "/licenceplate-anonymization/"
-DOCUMENT_QA = "/document-qa/"
+
 IMAGE_SIMILARITY_DEMO = "/find-similar-image/"
 IMAGE_SIMILARITY_PINECONE_DEMO = "/find-similar-image-pinecone/"
 INDEX_NAME = "imagesearch-demo"
 INDEX_DIMENSION = 512
 TMP_DIR = "tmp"
 
+image_sim_model = SentenceTransformer("clip-ViT-B-32")
+
 def enhance_image(pil_image):
     # Convert PIL Image to OpenCV format
     open_cv_image = np.array(pil_image)
@@ -86,84 +74,34 @@ def enhance_image(pil_image):
     return enhanced_pil_image
 
     
+print("checking pinecone Index")
 if INDEX_NAME not in pinecone.list_indexes():
-    pinecone.create_index(INDEX_NAME, dimension=512, metric='cosine')
+    # delete the current index and create the new index if it does not exist
+    for delete_index in pinecone.list_indexes():
+        print(f"Deleting exitsing pinecone Index : {delete_index}")
+
+        pinecone.delete_index(delete_index)
+    print(f"Creating new pinecone Index : {INDEX_NAME}")
+    pinecone.create_index(INDEX_NAME, dimension=INDEX_DIMENSION, metric="cosine")
     
 print("Connecting to Pinecone Index")
 index = pinecone.Index(INDEX_NAME)
     
-os.makedirs(TMP_DIR, exist_ok=True)
 
-# licence_model = torch.hub.load(
-#     "ultralytics/yolov5", "custom", path="Licenseplate_model.pt", device="cpu", force_reload=True
-# )
-# licence_model.cpu()
+device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# detector = cv2.dnn.DetectionModel(
-#     "res10_300x300_ssd_iter_140000_fp16.caffemodel", "deploy.prototxt"
-# )
+os.makedirs(TMP_DIR, exist_ok=True)
 
-# processor = DonutProcessor.from_pretrained("naver-clova-ix/donut-base-finetuned-docvqa")
-# doc_qa_model = VisionEncoderDecoderModel.from_pretrained(
-#     "naver-clova-ix/donut-base-finetuned-docvqa"
-# )
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
-# doc_qa_model.to(device)
 
 
 os.makedirs(TMP_DIR, exist_ok=True)
 
-# model = torch.hub.load(
-#     "ultralytics/yolov5", "custom", path="best.pt", device="cpu", force_reload=True
-# )
-# model.cpu()
-
-# classes = [
-#     "gas-distribution-meter",
-#     "gas-distribution-piping",
-#     "gas-distribution-regulator",
-#     "gas-distribution-valve",
-# ]
-
-# class_to_idx = {
-#     "gas-distribution-meter": 0,
-#     "gas-distribution-piping": 1,
-#     "gas-distribution-regulator": 2,
-#     "gas-distribution-valve": 3,
-# }
-
-# idx_to_classes = {v: k for k, v in class_to_idx.items()}
-# modelname = "resnet50d"
-# model_weights = "best_classifer_model.pt"
-# num_classes = len(classes)
-
-# classifier_model = timm.create_model(
-#     "resnet50d", pretrained=True, num_classes=num_classes, drop_path_rate=0.05
-# )
-# classifier_model.load_state_dict(
-#     torch.load(model_weights, map_location=torch.device("cpu"))["model_state_dict"]
-# )
-
-# musiq_metric = pyiqa.create_metric("musiq-koniq", device=torch.device("cpu"))
-# image_sim_model = SentenceTransformer("patrickjohncyh/fashion-clip")
-# from transformers import AutoProcessor, AutoModelForZeroShotImageClassification
-
-# processor = AutoProcessor.from_pretrained("patrickjohncyh/fashion-clip")
-# model = AutoModelForZeroShotImageClassification.from_pretrained("patrickjohncyh/fashion-clip")
-
-# model_ckpt = "nateraw/vit-base-beans"
-# extractor = AutoFeatureExtractor.from_pretrained(model_ckpt)
-# image_sim_model = AutoModel.from_pretrained(model_ckpt)
-fclip = FashionCLIP('fashion-clip')
 
-app = FastAPI(title="CV Demos")
 
-# Define the Response
-class Prediction(BaseModel):
-    filename: str
-    contenttype: str
-    prediction: List[float] = []
+
+app = FastAPI(title="CV Demos")
 
 
 # define response
@@ -172,280 +110,6 @@ def root_route():
     return {"error": f"Use GET {IMAGE_SIMILARITY_PINECONE_DEMO} instead of the root route!"}
 
 
-# @app.post(
-#     DETECTION_URL,
-# )
-# async def predict(file: UploadFile = File(...), quality_check: bool = False):
-#     try:
-#         extension = file.filename.split(".")[-1] in ("jpg", "jpeg", "png")
-#         if not extension:
-#             return "Image must be jpg or png format!"
-#         # read image contain
-#         contents = await file.read()
-#         pil_image = Image.open(io.BytesIO(contents))
-#         if quality_check:
-#             print("RUNNING QUALITY CEHCK BEFORE OBJEFCT DETECTION!!!")
-#             tmp_file = f"{TMP_DIR}/tmp.png"
-#             pil_image.save(tmp_file)
-#             score = musiq_metric(tmp_file)
-#             if score < 50:
-#                 return {
-#                     "Error": "Image quality is not sufficient enough to be considered for object detection"
-#                 }
-
-#         results = model(pil_image, size=640)  # reduce size=320 for faster inference
-#         return results.pandas().xyxy[0].to_json(orient="records")
-#     except:
-#         e = sys.exc_info()[1]
-#         raise HTTPException(status_code=500, detail=str(e))
-
-
-# @app.post(CLASSIFICATION_URL)
-# async def classify(file: UploadFile = File(...)):
-#     try:
-#         extension = file.filename.split(".")[-1] in ("jpg", "jpeg", "png")
-#         if not extension:
-#             return "Image must be jpg or png format!"
-#         # read image contain
-#         contents = await file.read()
-#         pil_image = Image.open(io.BytesIO(contents))
-#         data_mean = (0.485, 0.456, 0.406)
-#         data_std = (0.229, 0.224, 0.225)
-#         image_size = (224, 224)
-#         eval_transforms = timm.data.create_transform(
-#             input_size=image_size, mean=data_mean, std=data_std
-#         )
-#         eval_transforms(pil_image).unsqueeze(dim=0).shape
-#         classifier_model.eval()
-#         print("RUNNING Image Classification!!!")
-#         max_class_idx = np.argmax(
-#             classifier_model(eval_transforms(pil_image).unsqueeze(dim=0)).detach().numpy()
-#         )
-#         predicted_class = idx_to_classes[max_class_idx]
-#         print(f"Predicted Class idx: {max_class_idx} with name : {predicted_class}")
-#         return {"object": predicted_class}
-
-#     except:
-#         e = sys.exc_info()[1]
-#         raise HTTPException(status_code=500, detail=str(e))
-
-
-# @app.post(QUALITY_ASSESSMENT_URL)
-# async def quality_check(file: UploadFile = File(...)):
-#     try:
-#         extension = file.filename.split(".")[-1] in ("jpg", "jpeg", "png")
-#         if not extension:
-#             return "Image must be jpg or png format!"
-#         # read image contain
-#         contents = await file.read()
-#         pil_image = Image.open(io.BytesIO(contents))
-#         tmp_file = f"{TMP_DIR}/tmp.png"
-#         pil_image.save(tmp_file)
-#         score = musiq_metric(tmp_file).detach().numpy().tolist()
-#         return {"score": score}
-
-#     except:
-#         e = sys.exc_info()[1]
-#         raise HTTPException(status_code=500, detail=str(e))
-
-
-# def anonymize_simple(image, factor=3.0):
-#     # automatically determine the size of the blurring kernel based
-#     # on the spatial dimensions of the input image
-#     (h, w) = image.shape[:2]
-#     kW = int(w / factor)
-#     kH = int(h / factor)
-#     # ensure the width of the kernel is odd
-#     if kW % 2 == 0:
-#         kW -= 1
-#     # ensure the height of the kernel is odd
-#     if kH % 2 == 0:
-#         kH -= 1
-#     # apply a Gaussian blur to the input image using our computed
-#     # kernel size
-#     return cv2.GaussianBlur(image, (kW, kH), 0)
-
-
-# def anonymize_pixelate(image, blocks=3):
-#     # divide the input image into NxN blocks
-#     (h, w) = image.shape[:2]
-#     xSteps = np.linspace(0, w, blocks + 1, dtype="int")
-#     ySteps = np.linspace(0, h, blocks + 1, dtype="int")
-#     # loop over the blocks in both the x and y direction
-#     for i in range(1, len(ySteps)):
-#         for j in range(1, len(xSteps)):
-#             # compute the starting and ending (x, y)-coordinates
-#             # for the current block
-#             startX = xSteps[j - 1]
-#             startY = ySteps[i - 1]
-#             endX = xSteps[j]
-#             endY = ySteps[i]
-#             # extract the ROI using NumPy array slicing, compute the
-#             # mean of the ROI, and then draw a rectangle with the
-#             # mean RGB values over the ROI in the original image
-#             roi = image[startY:endY, startX:endX]
-#             (B, G, R) = [int(x) for x in cv2.mean(roi)[:3]]
-#             cv2.rectangle(image, (startX, startY), (endX, endY), (B, G, R), -1)
-#     # return the pixelated blurred image
-#     return image
-
-
-# # define response
-# @app.get("/")
-# def root_route():
-#     return {"error": f"Use GET {FACE_URL}  or {LICENCE_URL} instead of the root route!"}
-
-
-# @app.post(
-#     FACE_URL,
-# )
-# async def face_anonymize(
-#     file: UploadFile = File(...), blur_type="simple", quality_check: bool = False
-# ):
-#     """
-#     https://pyimagesearch.com/2020/04/06/blur-and-anonymize-faces-with-opencv-and-python/
-#     """
-#     try:
-#         extension = file.filename.split(".")[-1] in ("jpg", "jpeg", "png")
-#         if not extension:
-#             return "Image must be jpg or png format!"
-#         # read image contain
-#         contents = await file.read()
-#         pil_image = Image.open(io.BytesIO(contents)).convert("RGB")
-#         detector = cv2.dnn.DetectionModel(
-#             "res10_300x300_ssd_iter_140000_fp16.caffemodel", "deploy.prototxt"
-#         )
-#         open_cv_image = np.array(pil_image)
-#         # Convert RGB to BGR
-#         open_cv_image = open_cv_image[:, :, ::-1].copy()
-#         (h, w) = open_cv_image.shape[:2]
-#         # Getting the detections
-#         detections = detector.detect(open_cv_image)
-#         if len(detections[2]) > 0:
-#             for face in detections[2]:
-#                 (x, y, w, h) = face.astype("int")
-#                 # extract the face ROI
-
-#                 face = open_cv_image[y : y + h, x : x + w]
-#                 if blur_type == "simple":
-#                     face = anonymize_simple(face)
-#                 else:
-#                     face = anonymize_pixelate(face)
-#                 open_cv_image[y : y + h, x : x + w] = face
-
-#         _, encoded_img = cv2.imencode(".PNG", open_cv_image)
-
-#         encoded_img = base64.b64encode(encoded_img)
-#         return {
-#             "filename": file.filename,
-#             "dimensions": str(open_cv_image.shape),
-#             "encoded_img": encoded_img,
-#         }
-#     except:
-#         e = sys.exc_info()[1]
-#         print(traceback.format_exc())
-#         raise HTTPException(status_code=500, detail=str(e))
-
-
-# @app.post(LICENCE_URL)
-# async def licence_anonymize(file: UploadFile = File(...), blur_type="simple"):
-#     """https://www.kaggle.com/code/gowrishankarp/license-plate-detection-yolov5-pytesseract/notebook#Visualize"""
-#     try:
-#         extension = file.filename.split(".")[-1] in ("jpg", "jpeg", "png")
-#         if not extension:
-#             return "Image must be jpg or png format!"
-#         # read image contain
-#         contents = await file.read()
-#         pil_image = Image.open(io.BytesIO(contents))
-#         results = licence_model(pil_image, size=640)  # reduce size=320 for faster inference
-#         pil_image = pil_image.convert("RGB")
-#         open_cv_image = np.array(pil_image)
-#         open_cv_image = open_cv_image[:, :, ::-1].copy()
-#         df = results.pandas().xyxy[0]
-#         for i, row in df.iterrows():
-#             xmin = int(row["xmin"])
-#             ymin = int(row["ymin"])
-#             xmax = int(row["xmax"])
-#             ymax = int(row["ymax"])
-#             licence = open_cv_image[ymin:ymax, xmin:xmax]
-#             if blur_type == "simple":
-#                 licence = anonymize_simple(licence)
-#             else:
-#                 licence = anonymize_pixelate(licence)
-#             open_cv_image[ymin:ymax, xmin:xmax] = licence
-
-#         _, encoded_img = cv2.imencode(".PNG", open_cv_image)
-
-#         encoded_img = base64.b64encode(encoded_img)
-#         return {
-#             "filename": file.filename,
-#             "dimensions": str(open_cv_image.shape),
-#             "encoded_img": encoded_img,
-#         }
-
-#     except:
-#         e = sys.exc_info()[1]
-#         raise HTTPException(status_code=500, detail=str(e))
-
-
-# def process_document(image, question):
-#     # prepare encoder inputs
-#     pixel_values = processor(image, return_tensors="pt").pixel_values
-
-#     # prepare decoder inputs
-#     task_prompt = "<s_docvqa><s_question>{user_input}</s_question><s_answer>"
-#     prompt = task_prompt.replace("{user_input}", question)
-#     decoder_input_ids = processor.tokenizer(
-#         prompt, add_special_tokens=False, return_tensors="pt"
-#     ).input_ids
-
-#     # generate answer
-#     outputs = doc_qa_model.generate(
-#         pixel_values.to(device),
-#         decoder_input_ids=decoder_input_ids.to(device),
-#         max_length=doc_qa_model.decoder.config.max_position_embeddings,
-#         early_stopping=True,
-#         pad_token_id=processor.tokenizer.pad_token_id,
-#         eos_token_id=processor.tokenizer.eos_token_id,
-#         use_cache=True,
-#         num_beams=1,
-#         bad_words_ids=[[processor.tokenizer.unk_token_id]],
-#         return_dict_in_generate=True,
-#     )
-
-#     # postprocess
-#     sequence = processor.batch_decode(outputs.sequences)[0]
-#     sequence = sequence.replace(processor.tokenizer.eos_token, "").replace(
-#         processor.tokenizer.pad_token, ""
-#     )
-#     sequence = re.sub(r"<.*?>", "", sequence, count=1).strip()  # remove first task start token
-
-#     return processor.token2json(sequence)
-
-
-# @app.post(DOCUMENT_QA)
-# async def document_qa(question: str = Form(...), file: UploadFile = File(...)):
-
-#     try:
-#         extension = file.filename.split(".")[-1] in ("jpg", "jpeg", "png")
-#         if not extension:
-#             return "Image must be jpg or png format!"
-#         # read image contain
-#         contents = await file.read()
-#         pil_image = Image.open(io.BytesIO(contents))
-#         # tmp_file = f"{TMP_DIR}/tmp.png"
-#         # pil_image.save(tmp_file)
-#         # answer_git_large = generate_answer_git(git_processor_large, git_model_large, image, question)
-
-#         answer = process_document(pil_image, question)["answer"]
-
-#         return {"answer": answer}
-
-#     except:
-#         e = sys.exc_info()[1]
-#         raise HTTPException(status_code=500, detail=str(e))
-
-
 @app.post(IMAGE_SIMILARITY_DEMO)
 async def image_search_local(
     images_to_search: List[UploadFile], query_image: UploadFile = File(...), top_k: int = 5
@@ -540,8 +204,7 @@ async def image_search_pinecone(
             query_image = Image.open(io.BytesIO(contents))
             print("Indexing query image...")
             query_image = enhance_image(query_image)
-            # prompt_embedding = image_sim_model.encode(query_image, convert_to_tensor=True).tolist()
-            prompt_embedding = fclip.encode_images([query_image], batch_size=32)[0]
+            prompt_embedding = image_sim_model.encode(query_image, convert_to_tensor=True).tolist()
             if INDEX_NAME not in pinecone.list_indexes():
                 return {"similar_images": [], "status": "No index found for images"}
 
@@ -625,10 +288,9 @@ async def image_search_pinecone(
                 ids.append(str(uuid.uuid1()).replace("-",""))
 
             print("Encoding images to vectors...")
-            # corpus_embeddings = image_sim_model.encode(
-            #     search_images, convert_to_tensor=True, show_progress_bar=True
-            # ).tolist()
-            corpus_embeddings = fclip.encode_images(search_images, batch_size=32)[0]
+            corpus_embeddings = image_sim_model.encode(
+                search_images, convert_to_tensor=True, show_progress_bar=True
+            ).tolist()
             print(f"Indexing images to pinecone Index : {INDEX_NAME}")
             index.upsert(
                 vectors=list(zip(ids, corpus_embeddings, meta_datas)), namespace=namespace
@@ -649,4 +311,4 @@ if __name__ == "__main__":
     parser.add_argument("--port", default=8000, type=int, help="port number")
     # parser.add_argument('--model', nargs='+', default=['yolov5s'], help='model(s) to run, i.e. --model yolov5n yolov5s')
     opt = parser.parse_args()
-    uvicorn.run(app, port=opt.port)
+    uvicorn.run(app, port=opt.port)a