""" Utility functions for image processing and validation. """ import base64 import io import os from pathlib import Path from PIL import Image SUPPORTED_FORMATS = {".png", ".jpg", ".jpeg", ".gif", ".webp", ".bmp"} MAX_IMAGE_SIZE_MB = 10 MAX_DIMENSION = 4096 # Max width/height for API submission def validate_image(image_bytes: bytes, filename: str = "image.png") -> tuple[bool, str]: """Validate an image file for chart analysis. Returns: (is_valid, error_message) """ ext = Path(filename).suffix.lower() if ext not in SUPPORTED_FORMATS: return False, f"Unsupported format: {ext}. Use: {', '.join(SUPPORTED_FORMATS)}" size_mb = len(image_bytes) / (1024 * 1024) if size_mb > MAX_IMAGE_SIZE_MB: return False, f"Image too large: {size_mb:.1f}MB (max {MAX_IMAGE_SIZE_MB}MB)" try: img = Image.open(io.BytesIO(image_bytes)) img.verify() except Exception as e: return False, f"Invalid image file: {e}" return True, "" def prepare_image_for_api( image_bytes: bytes, max_dimension: int = MAX_DIMENSION, ) -> tuple[bytes, str]: """Resize image if needed and return (processed_bytes, media_type). Ensures the image fits within API size limits while preserving quality. """ img = Image.open(io.BytesIO(image_bytes)) # Convert RGBA to RGB if needed (for JPEG output) if img.mode == "RGBA": background = Image.new("RGB", img.size, (255, 255, 255)) background.paste(img, mask=img.split()[3]) img = background elif img.mode != "RGB": img = img.convert("RGB") # Resize if too large w, h = img.size if w > max_dimension or h > max_dimension: ratio = min(max_dimension / w, max_dimension / h) new_size = (int(w * ratio), int(h * ratio)) img = img.resize(new_size, Image.LANCZOS) # Save as PNG for best quality buffer = io.BytesIO() img.save(buffer, format="PNG", optimize=True) return buffer.getvalue(), "image/png" def image_to_base64(image_bytes: bytes) -> str: """Encode image bytes to base64 string.""" return base64.standard_b64encode(image_bytes).decode("utf-8") def load_image_from_path(path: str) -> bytes: """Load image bytes from a file path.""" with open(path, "rb") as f: return f.read() def get_image_info(image_bytes: bytes) -> dict: """Get basic image metadata.""" img = Image.open(io.BytesIO(image_bytes)) return { "width": img.size[0], "height": img.size[1], "format": img.format or "unknown", "mode": img.mode, "size_kb": len(image_bytes) / 1024, } def detect_ticker_from_image(image_bytes: bytes, provider: str, model: str, api_key: str) -> str: """Quick vision call to extract the stock ticker from a chart screenshot. Returns the ticker string (e.g., 'NVDA', 'AAPL', 'D05.SI') or empty string. """ processed_bytes, media_type = prepare_image_for_api(image_bytes) b64 = image_to_base64(processed_bytes) prompt = ( "Look at this stock chart image. What is the stock ticker symbol shown? " "Reply with ONLY the ticker symbol (e.g., NVDA, AAPL, D05.SI). " "If you cannot determine it, reply with UNKNOWN." ) try: if provider == "google": from google import genai from google.genai import types client = genai.Client(api_key=api_key) image_part = types.Part.from_bytes(data=processed_bytes, mime_type=media_type) response = client.models.generate_content(model=model, contents=[image_part, prompt]) result = response.text.strip() elif provider in ("huggingface", "openai", "openrouter"): import openai if provider == "huggingface": base_url = "https://router.huggingface.co/v1" elif provider == "openrouter": base_url = "https://openrouter.ai/api/v1" else: base_url = None client = openai.OpenAI(base_url=base_url, api_key=api_key) if base_url else openai.OpenAI(api_key=api_key) response = client.chat.completions.create( model=model, max_tokens=20, messages=[{ "role": "user", "content": [ {"type": "image_url", "image_url": {"url": f"data:{media_type};base64,{b64}"}}, {"type": "text", "text": prompt}, ], }], ) result = response.choices[0].message.content.strip() elif provider == "anthropic": import anthropic client = anthropic.Anthropic(api_key=api_key) response = client.messages.create( model=model, max_tokens=20, messages=[{"role": "user", "content": [ {"type": "image", "source": {"type": "base64", "media_type": media_type, "data": b64}}, {"type": "text", "text": prompt}, ]}], ) result = response.content[0].text.strip() else: return "" # Clean up — extract just the ticker from the response result = result.replace("*", "").replace("`", "").strip() if result and result != "UNKNOWN" and len(result) <= 15: return result except Exception: pass return ""