add image_as_payload_to_embeddings

app.py

@@ -9,13 +9,13 @@ import math
 # from transformers import CLIPTextModel, CLIPTokenizer
 import os
 
-
 # clip_model_id = "openai/clip-vit-large-patch14-336"
 # clip_retrieval_indice_name, clip_model_id ="laion5B-L-14", "/laion/CLIP-ViT-L-14-laion2B-s32B-b82K"
 clip_retrieval_service_url = "https://knn.laion.ai/knn-service"
 # available models = ['RN50', 'RN101', 'RN50x4', 'RN50x16', 'RN50x64', 'ViT-B/32', 'ViT-B/16', 'ViT-L/14', 'ViT-L/14@336px']
 # clip_model="ViT-B/32"
 clip_model="ViT-L/14"
+clip_image_size = 224
 clip_model_id ="laion5B-L-14"
 
 
@@ -35,6 +35,62 @@ def debug_print(*args, **kwargs):
     if debug_print_on:
         print(*args, **kwargs)
 
+# support sending images as base64
+
+def encode_numpy_array(image_np):
+    import base64
+    import json
+    # Flatten the numpy array and convert it to bytes
+    image_bytes = image_np.tobytes()
+
+    # Encode the byte data as base64
+    encoded_image = base64.b64encode(image_bytes).decode()
+    payload = {
+        "encoded_image": encoded_image,
+        "width": image_np.shape[1],
+        "height": image_np.shape[0],
+        "channels": image_np.shape[2],
+    }
+    payload_json = json.dumps(payload)
+    return payload_json
+
+def decode_numpy_array(payload):
+    import base64
+    import json
+    payload_json = json.loads(payload)
+    # payload_json = payload.json()
+    encoded_image = payload_json["encoded_image"]
+    width = payload_json["width"]
+    height = payload_json["height"]
+    channels = payload_json["channels"]
+    # Decode the base64 data
+    decoded_image = base64.b64decode(encoded_image)
+
+    # Convert the byte data back to a NumPy array
+    image_np = np.frombuffer(decoded_image, dtype=np.uint8).reshape(height, width, channels)
+
+    return image_np
+
+
+def preprocess_image(image_np, max_size=224):
+    from torchvision.transforms import Compose, Resize, CenterCrop
+    # Convert the numpy array to a PIL image
+    image = Image.fromarray(image_np)
+
+    # Define the transformation pipeline
+    transforms = Compose([
+        Resize(max_size, interpolation=Image.BICUBIC),
+        CenterCrop(max_size),
+    ])
+
+    # Apply the transformations to the image
+    image = transforms(image)
+
+    # Convert the PIL image back to a numpy array
+    image_np = np.array(image)
+
+    return image_np
+
 def image_to_embedding(input_im):
     # debug_print("image_to_embedding")
     input_im = Image.fromarray(input_im)
@@ -181,6 +237,16 @@ def on_image_load_update_embeddings(image_data):
     # return gr.Text.update(embeddings_b64)
     return embeddings_b64
 
+def on_image_as_payload_update_embeddings(payload):
+    debug_print("on_image_as_payload_update_embeddings")
+    if payload is None or payload == "":
+        return ''
+    image_data = decode_numpy_array(payload)
+    embeddings = image_to_embedding(image_data)
+    embeddings_b64 = embedding_to_base64(embeddings)
+    # return gr.Text.update(embeddings_b64)
+    return embeddings_b64
+
 def on_prompt_change_update_embeddings(prompt):
     debug_print("on_prompt_change_update_embeddings")
     # prompt to embeddings
@@ -578,6 +644,23 @@ UI elements to mock out the API
             _output = gr.Textbox(value="", lines=2, label="Output")
         _btn = gr.Button(value="Submit")
         _btn.click(on_image_load_update_embeddings, inputs=_input, outputs=[_output], api_name="image_to_embeddings")
+    with gr.Row():
+        def _on_image_load(input_image):
+            debug_print("_on_image_load")
+            # resize if size is bigger than clip_image_size
+            if input_image.shape[0] > clip_image_size or input_image.shape[1] > clip_image_size:
+                input_image = preprocess_image(input_image, clip_image_size)
+            payload = encode_numpy_array(input_image)
+            return payload
+
+        _input = gr.Image(label="Image Prompt", show_label=True)
+        with gr.Accordion(f"Image (base64)", open=False):
+            _input_as_texct = gr.Textbox(value="", lines=2, label="Output")
+        _input.change(_on_image_load, inputs=_input, outputs=[_input_as_texct])
+        with gr.Accordion(f"Embeddings (base64)", open=False):
+            _output = gr.Textbox(value="", lines=2, label="Output")
+        _btn = gr.Button(value="Submit")
+        _btn.click(on_image_as_payload_update_embeddings, inputs=_input_as_texct, outputs=[_output], api_name="image_as_payload_to_embeddings")
 
 # ![Alt Text](file/pup1.jpg)