Update app.py

app.py

@@ -1,3 +1,5 @@
+import cv2
+import numpy as np
 import argparse
 import base64
 import io
@@ -7,7 +9,8 @@ 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
@@ -18,21 +21,30 @@ import torch
 import uvicorn
 from dotenv import load_dotenv
 from fastapi import FastAPI, File, Form, HTTPException, UploadFile
-from PIL import Image, ImageEnhance
+from PIL import Image
 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)
@@ -74,35 +86,85 @@ def enhance_image(pil_image):
     return enhanced_pil_image
 
     
-print("checking pinecone Index")
 if INDEX_NAME not in pinecone.list_indexes():
-    # 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")
+    pinecone.create_index(INDEX_NAME, dimension=512, metric='cosine')
     
 print("Connecting to Pinecone Index")
 index = pinecone.Index(INDEX_NAME)
     
+os.makedirs(TMP_DIR, exist_ok=True)
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
+# licence_model = torch.hub.load(
+#     "ultralytics/yolov5", "custom", path="Licenseplate_model.pt", device="cpu", force_reload=True
+# )
+# licence_model.cpu()
 
-os.makedirs(TMP_DIR, exist_ok=True)
+# detector = cv2.dnn.DetectionModel(
+#     "res10_300x300_ssd_iter_140000_fp16.caffemodel", "deploy.prototxt"
+# )
 
+# 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] = []
+
 
 # define response
 @app.get("/")
@@ -110,6 +172,280 @@ 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
@@ -204,7 +540,8 @@ 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 = image_sim_model.encode(query_image, convert_to_tensor=True).tolist()
+            prompt_embedding = fclip.encode_images([query_image], batch_size=32)[0]
             if INDEX_NAME not in pinecone.list_indexes():
                 return {"similar_images": [], "status": "No index found for images"}
 
@@ -288,9 +625,10 @@ 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 = 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]
             print(f"Indexing images to pinecone Index : {INDEX_NAME}")
             index.upsert(
                 vectors=list(zip(ids, corpus_embeddings, meta_datas)), namespace=namespace
@@ -311,4 +649,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)