Create clip_base.py

clip_base.py

@@ -0,0 +1,145 @@
+import os
+import torch
+import requests
+import zipfile
+from PIL import Image
+from io import BytesIO
+import tensorflow as tf
+import pandas as pd
+from transformers import CLIPProcessor, TFCLIPModel
+
+
+class TfliteConverter:
+
+    def __init__(self, model, image_data, title_list, unique_title_list):
+        self.model = model
+        self.image_data = image_data
+        self.title_list = title_list
+        self.unique_title_list = unique_title_list
+
+    def build_model(self):
+        concrete_func = self.serving_fn.get_concrete_function()
+
+        converter = tf.lite.TFLiteConverter.from_concrete_functions([concrete_func])
+
+        converter.target_spec.supported_ops = [
+            tf.lite.OpsSet.TFLITE_BUILTINS,
+            tf.lite.OpsSet.SELECT_TF_OPS
+        ]
+
+        tflite_model = converter.convert()
+
+        model_path = os.path.join(os.getcwd(), 'clip_tflite_model.tflite')
+        json_path = os.path.join(os.getcwd(), 'categories.json')
+        zip_path = os.path.join(os.getcwd(), 'model_package.zip')
+
+        with open(model_path, 'wb') as f:
+            f.write(tflite_model)
+
+        categories_df = pd.DataFrame({
+            'id': range(len(self.unique_title_list)),
+            'title': self.unique_title_list
+        })
+        categories_df.to_json(json_path, orient='records', indent=2)
+
+        # Create ZIP file containing both files
+        with zipfile.ZipFile(zip_path, 'w', zipfile.ZIP_DEFLATED) as zipf:
+            zipf.write(model_path, os.path.basename(model_path))
+            zipf.write(json_path, os.path.basename(json_path))
+
+        # # Clean up temporary files
+        # os.remove(model_path)
+        # os.remove(json_path)
+
+        return zip_path
+
+    @tf.function(input_signature=[
+        tf.TensorSpec(shape=[1, 224, 224, 3], dtype=tf.float32, name='input_images')
+    ])
+    def serving_fn(self, input_images):
+        x = tf.transpose(input_images, (0, 3, 1, 2))
+
+        with torch.no_grad():
+            img_embeddings = self.model.get_image_features(pixel_values=x)
+
+        labels_embeddings = [item['embeddings'] for item in self.image_data]
+
+        labels_embeddings = tf.stack(labels_embeddings, axis=0)
+
+        similarities = tf.reduce_sum(tf.multiply(img_embeddings, labels_embeddings), axis=-1)
+        norm_img = tf.norm(img_embeddings, axis=-1)
+        norm_labels = tf.norm(labels_embeddings, axis=-1)
+
+        cosine_similarity = similarities / (norm_img * norm_labels)
+
+        name_to_score = {}
+
+        names = self.title_list
+        static_names = self.unique_title_list
+
+        for i in range(len(names)):
+            name = names[i]
+            score = cosine_similarity[i]
+
+            if name not in name_to_score or score > name_to_score[name]:
+                name_to_score[name] = score
+
+        result = [name_to_score.get(name, 0) for name in static_names]
+        return tf.convert_to_tensor(result, dtype=tf.float32)
+
+
+class OpenAiClipModel:
+
+    def __init__(self, payload):
+        self.payload = payload
+        self.processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+        self.model = TFCLIPModel.from_pretrained("openai/clip-vit-base-patch32")
+        self.df = pd.DataFrame([(key, value) for key, values in self.payload.items() for value in values], columns=['image_category', 'image_path'])
+        self.labels_embeddings = []
+        self.image_data = []
+        self.title_list = self.df["image_category"].tolist()
+        self.unique_title_list = list(set(self.title_list))
+
+
+    def generate_text_embedding(self, text):
+        inputs = self.processor(text=[text], return_tensors="pt", padding=True)
+        with torch.no_grad():
+            text_embedding = self.model.get_text_features(**inputs)
+        return text_embedding
+
+    def generate_image_embeddings(self):
+        for index, row in self.df.iterrows():
+            image_info = row.to_dict()
+            image_path = image_info.get("image_path")
+
+            if os.path.exists(image_path):
+                try:
+                    image = Image.open(image_path)
+                    inputs = self.processor(images=image, return_tensors="tf")
+                    outputs = self.model.get_image_features(**inputs)
+                    image_embedding = outputs.numpy().flatten()
+                    image_info['embeddings'] = image_embedding.tolist()
+                    self.labels_embeddings.append(outputs)
+                except Exception as e:
+                    image_info['embeddings'] = None
+            else:
+                image_info['embeddings'] = None
+
+            self.image_data.append(image_info)
+
+
+    def build_model(self):
+        self.generate_image_embeddings()
+        self.title_list = self.df["image_category"].tolist()
+
+        tflite_client = TfliteConverter(
+            model=self.model,
+            image_data=self.image_data,
+            title_list=self.title_list,
+            unique_title_list=self.unique_title_list
+        )
+        model_file = tflite_client.build_model()
+        return model_file
+
+
+