| import json |
| import chromadb |
| from chromadb.utils.embedding_functions import OpenCLIPEmbeddingFunction |
| from chromadb.utils.data_loaders import ImageLoader |
| import ollama |
| import base64 |
| import io |
| from io import BytesIO |
| import openai |
| from google import genai |
| import os |
| from PIL import Image |
| import numpy as np |
| from google.genai.types import Part, Content |
|
|
| client = genai.Client(api_key=os.environ.get("GENAI_API_KEY")) |
|
|
| class PlatinumPipeline: |
|
|
| def __init__(self, json_file="pricelist.json"): |
| self.json_file = json_file |
| self.collection_name = "pricetag" |
| self.TOP_K = 3 |
| self.chroma_client = chromadb.Client() |
| self.image_loader = ImageLoader() |
| self.multimodal_ef = OpenCLIPEmbeddingFunction() |
| self.model_name = "gemma3:4b" |
| self.multimodal_db = self.chroma_client.create_collection( |
| name=self.collection_name, |
| embedding_function=self.multimodal_ef, |
| data_loader=self.image_loader |
| ) |
| with open(self.json_file, "r") as f: |
| self.items = json.load(f) |
|
|
| self.valid_items, self.valid_ids, self.valid_uris, self.valid_metadatas = [], [], [], [] |
|
|
| for i, item in enumerate(self.items): |
| image_path = item["image"] |
| try: |
| with open(image_path, "rb"): |
| self.valid_items.append(item) |
| self.valid_ids.append(str(i)) |
| self.valid_uris.append(image_path) |
| self.valid_metadatas.append({ |
| "name": item["item_name"], |
| "price": item["price"], |
| "original_index": i |
| }) |
| except FileNotFoundError: |
| continue |
|
|
| if self.valid_items: |
| try: |
| self.multimodal_db.add( |
| ids=self.valid_ids, |
| uris=self.valid_uris, |
| metadatas=self.valid_metadatas |
| ) |
| except Exception as e: |
| print(e) |
| else: |
| print("Records are empty") |
|
|
| def query(self, image_bytes: bytes): |
| context = [] |
| if self.multimodal_db.count() == 0: |
| return "No records found in database!!!" |
|
|
| try: |
| pil_image = Image.open(io.BytesIO(image_bytes)).convert("RGB") |
| image_np = np.array(pil_image) |
| |
| query_results = self.multimodal_db.query( |
| query_images=[image_np], |
| n_results=min(self.TOP_K, self.multimodal_db.count()), |
| include=['uris', 'metadatas', 'distances', 'data'] |
| ) |
|
|
| if query_results['ids'] and len(query_results['ids'][0]) > 0: |
| for i in range(len(query_results['ids'][0])): |
| context.append(query_results['metadatas'][0][i]) |
| |
| |
| prompt = ( |
| f"Context:\n{context}\n\n" |
| "You have to write the estimated price of the item in the image " |
| "that has been provided based on the context which contains " |
| "the names and prices of similar items in the database." |
| ) |
|
|
| text_part = Part(text=prompt) |
| image_part = Part( |
| inline_data={ |
| "mime_type": "image/jpeg", |
| "data": base64.b64encode(image_bytes).decode('utf-8') |
| } |
| ) |
| |
| content = Content(parts=[text_part, image_part]) |
| response = client.models.generate_content( |
| model="gemini-2.0-flash-exp", |
| contents=content |
| ) |
| |
| return response.text |
| else: |
| return "No context found" |
| except Exception as e: |
| return f"Error: {str(e)}" |
