Update custom handler

handler.py

@@ -14,6 +14,39 @@ from clip.clip import _transform, tokenize
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+def preprocess_image(image_base64, transformer):
+    """Convert base64 encoded sketch to tensor."""
+    image = Image.open(BytesIO(base64.b64decode(image_base64))).convert("RGB")
+    image = transformer(image).unsqueeze(0).to(device)
+    return image
+
+def preprocess_text(text):
+    """Tokenize text query."""
+    return tokenize([str(text)])[0].unsqueeze(0).to(device)
+
+def get_fused_embedding(sketch_base64, text, model):
+    """Fuse sketch and text features into a single embedding."""
+    with torch.no_grad():
+        sketch_tensor = preprocess_image(sketch_base64)
+        text_tensor = preprocess_text(text)
+
+        sketch_feature = model.encode_sketch(sketch_tensor)
+        text_feature = model.encode_text(text_tensor)
+        
+        sketch_feature = sketch_feature / sketch_feature.norm(dim=-1, keepdim=True)
+        text_feature = text_feature / text_feature.norm(dim=-1, keepdim=True)
+
+        fused_embedding = model.feature_fuse(sketch_feature, text_feature)
+    return fused_embedding.cpu().numpy().tolist()
+
+def get_image_embedding(image_base64, model):
+    """Convert base64 encoded image to tensor."""
+    image_tensor = preprocess_image(image_base64)
+    with torch.no_grad():
+        image_feature = model.encode_image(image_tensor)
+        image_feature = image_feature / image_feature.norm(dim=-1, keepdim=True)
+    return image_feature.cpu().numpy().tolist()
+
 class EndpointHandler:
     def __init__(self, path: str = ""):
         """
@@ -43,35 +76,26 @@ class EndpointHandler:
         """
         Process the request and return the fused embedding.
         Args:
-            data (dict): Includes 'image' (base64) and 'text' (str) inputs.
+            data (dict): Includes 'sketch' (base64) and 'text' (str) inputs, or 'image' (base64)
         Returns:
-            dict: {"fused_embedding": [float, float, ...]}
+            dict: {"embedding": [float, float, ...]}
         """
         # Parse inputs
         inputs = data.pop("inputs", data)
-        image_base64 = inputs.get("image", "")
-        text_query = inputs.get("text", "")
-
-        if not image_base64 or not text_query:
-            return {"error": "Both 'image' (base64) and 'text' are required inputs."}
-
-        # Preprocess the image
-        image = Image.open(BytesIO(base64.b64decode(image_base64))).convert("RGB")
-        image_tensor = self.transform(image).unsqueeze(0).to(device)
-
-        # Preprocess the text
-        text_tensor = tokenize([str(text_query)])[0].unsqueeze(0).to(device)
-
-        # Generate features
-        with torch.no_grad():
-            sketch_feature = self.model.encode_sketch(image_tensor)
-            text_feature = self.model.encode_text(text_tensor)
-
-            # Normalize features
-            sketch_feature = sketch_feature / sketch_feature.norm(dim=-1, keepdim=True)
-            text_feature = text_feature / text_feature.norm(dim=-1, keepdim=True)
-
-            # Fuse features
-            fused_embedding = self.model.feature_fuse(sketch_feature, text_feature)
+        if "sketch" in inputs:
+            sketch_base64 = inputs.get("sketch", "")
+            text_query = inputs.get("text", "")
+            if not sketch_base64 or not text_query:
+                return {"error": "Both 'sketch' (base64) and 'text' are required inputs."}
 
-        return {"fused_embedding": fused_embedding.cpu().numpy().tolist()}
+            # Generate Fused Embedding
+            fused_embedding = get_fused_embedding(sketch_base64, text_query)
+            return {"embedding": fused_embedding}
+        elif "image" in inputs:
+            image_base64 = inputs.get("image", "")
+            if not image_base64:
+                return {"error": "Image 'image' (base64) is required input."}
+            embedding = get_image_embedding(image_base64)
+            return {"embedding": embedding}
+        else:
+            return {"error": "Input 'sketch' or 'image' is required."}

test.py

@@ -10,13 +10,13 @@ def encode_image_to_base64(image_path):
 handler = EndpointHandler(path=".")
 
 # Prepare sample inputs
-image_path = "path_to_your_sketch_image.jpg"  # Replace with your image path
+image_path = "sketches/COCO_val2014_000000163852.jpg"
 base64_image = encode_image_to_base64(image_path)
 text_query = "A pink flower"
 
 # Create payload
 payload = {
-    "image": base64_image,
+    "sketch": base64_image,
     "text": text_query
 }