demeter

  1# encoders/clip_encoder.py
  2import torch
  3import clip
  4from PIL import Image
  5import io
  6import base64
  7from pathlib import Path
  8
  9class VisionEncoder:
 10    def __init__(self, model_name="ViT-B/32"):
 11        self.device = "cuda" if torch.cuda.is_available() else "cpu"
 12        self.model, self.preprocess = clip.load(model_name, device=self.device)
 13        self.model.eval()
 14
 15    def encode(self, image_input):
 16        """
 17        Encode an image from multiple input types:
 18        - File path (str or Path)
 19        - Base64 string
 20        - BytesIO object
 21        - PIL Image object
 22        
 23        Args:
 24            image_input: File path, base64 string, BytesIO, or PIL Image
 25            
 26        Returns:
 27            numpy array: Normalized image embedding vector
 28        """
 29        # Convert input to PIL Image
 30        pil_image = self._to_pil_image(image_input)
 31        
 32        # Preprocess and encode
 33        image = self.preprocess(pil_image.convert("RGB")) \
 34                    .unsqueeze(0).to(self.device)
 35
 36        with torch.no_grad():
 37            vec = self.model.encode_image(image)
 38            vec = vec / vec.norm(dim=-1, keepdim=True)
 39
 40        return vec.cpu().numpy().flatten()
 41    
 42    def _to_pil_image(self, image_input):
 43        """
 44        Convert various input types to PIL Image.
 45        """
 46        # If already a PIL Image
 47        if isinstance(image_input, Image.Image):
 48            return image_input
 49        
 50        # If BytesIO object
 51        if isinstance(image_input, io.BytesIO):
 52            image_input.seek(0)  # Reset to beginning
 53            return Image.open(image_input)
 54        
 55        # If it's a string, determine if it's a path or base64
 56        if isinstance(image_input, (str, Path)):
 57            # Check if it's a file path
 58            if isinstance(image_input, Path) or Path(image_input).exists():
 59                return Image.open(image_input)
 60            
 61            # Otherwise, treat as base64
 62            return self._base64_to_pil(image_input)
 63        
 64        # If bytes object
 65        if isinstance(image_input, bytes):
 66            return Image.open(io.BytesIO(image_input))
 67        
 68        raise TypeError(f"Unsupported image input type: {type(image_input)}")
 69    
 70    def _base64_to_pil(self, base64_string):
 71        """
 72        Convert base64 string to PIL Image.
 73        """
 74        # Remove header if present (e.g., "data:image/png;base64,...")
 75        if "," in base64_string:
 76            base64_string = base64_string.split(",")[1]
 77        
 78        # Add padding if necessary
 79        missing_padding = len(base64_string) % 4
 80        if missing_padding:
 81            base64_string += '=' * (4 - missing_padding)
 82        
 83        # Decode and open
 84        image_bytes = base64.b64decode(base64_string)
 85        return Image.open(io.BytesIO(image_bytes))
 86
 87
 88# Example usage:
 89if __name__ == "__main__":
 90    encoder = VisionEncoder()
 91    
 92    # Test with file path
 93    # vec1 = encoder.encode("path/to/image.jpg")
 94    
 95    # Test with base64
 96    sample_base64 = "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+ip1sAAAAASUVORK5CYII="
 97    vec2 = encoder.encode(sample_base64)
 98    print(f"✅ Encoded base64 image. Vector shape: {vec2.shape}")
 99    
100    # Test with BytesIO
101    # image_stream = io.BytesIO(image_bytes)
102    # vec3 = encoder.encode(image_stream)