processor/upscale.py

import cv2
import numpy as np
import os
import io
import base64
from PIL import Image
from dotenv import load_dotenv
from google import genai
from google.genai import types

load_dotenv()

# Initialize the Gemini client
_client = None

def get_client():
    global _client
    if _client is None:
        api_key = os.getenv("GEMINI_API_KEY")
        if not api_key or api_key == "your_api_key_here":
            raise ValueError("GEMINI_API_KEY not set. Please add it to your .env file.")
        _client = genai.Client(api_key=api_key)
    return _client


def upscale_image(img: np.ndarray) -> np.ndarray:
    """
    Upscales the image using Google Gemini's image generation API.
    Sends the cropped card image to Gemini with a prompt to upscale it,
    then returns the AI-enhanced result.
    
    Handles both BGR and BGRA (transparent) images.
    Falls back to local upscaling if Gemini API fails.
    """
    has_alpha = len(img.shape) == 3 and img.shape[2] == 4
    
    if has_alpha:
        bgr = img[:, :, :3]
        alpha = img[:, :, 3]
    else:
        bgr = img
        alpha = None
    
    try:
        # Convert BGR (OpenCV) to RGB (PIL)
        rgb = cv2.cvtColor(bgr, cv2.COLOR_BGR2RGB)
        pil_image = Image.fromarray(rgb)
        
        # Call Gemini API to upscale
        upscaled_pil = _gemini_upscale(pil_image)
        
        # Convert back to OpenCV BGR
        upscaled_rgb = np.array(upscaled_pil)
        upscaled_bgr = cv2.cvtColor(upscaled_rgb, cv2.COLOR_RGB2BGR)
        
        if alpha is not None:
            # Resize alpha to match the upscaled image
            h, w = upscaled_bgr.shape[:2]
            upscaled_alpha = cv2.resize(alpha, (w, h), interpolation=cv2.INTER_LANCZOS4)
            _, upscaled_alpha = cv2.threshold(upscaled_alpha, 127, 255, cv2.THRESH_BINARY)
            
            return cv2.merge((
                upscaled_bgr[:, :, 0],
                upscaled_bgr[:, :, 1],
                upscaled_bgr[:, :, 2],
                upscaled_alpha
            ))
        else:
            return upscaled_bgr
            
    except Exception as e:
        print(f"Gemini upscale failed: {e}")
        print("Falling back to local upscaling...")
        return _local_fallback_upscale(img)


def _gemini_upscale(pil_image: Image.Image) -> Image.Image:
    """
    Uses the Gemini API to upscale/enhance an image.
    """
    client = get_client()

    response = client.models.generate_content(
        model="gemini-2.0-flash-exp",
        contents=[
            "Upscale this credit card image to high resolution. "
            "Make the text sharp, crisp, and readable. "
            "Preserve all colors, logos, textures, and details exactly. "
            "Do not add any watermarks, borders, or extra elements. "
            "Do not change the content of the image in any way. "
            "Output only the enhanced image.",
            pil_image,
        ],
        config=types.GenerateContentConfig(
            response_modalities=["IMAGE", "TEXT"],
        ),
    )

    # Extract the image from the response
    for part in response.candidates[0].content.parts:
        if part.inline_data is not None:
            img_bytes = part.inline_data.data
            return Image.open(io.BytesIO(img_bytes))
    
    raise ValueError("Gemini did not return an image in the response")


def _local_fallback_upscale(img: np.ndarray) -> np.ndarray:
    """
    Fallback: local multi-pass Lanczos + sharpening if Gemini API is unavailable.
    """
    has_alpha = len(img.shape) == 3 and img.shape[2] == 4
    
    if has_alpha:
        bgr = img[:, :, :3]
        alpha = img[:, :, 3]
    else:
        bgr = img
        alpha = None

    h, w = bgr.shape[:2]
    upscaled = cv2.resize(bgr, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)
    upscaled = cv2.bilateralFilter(upscaled, d=5, sigmaColor=40, sigmaSpace=40)
    
    # Unsharp mask
    blurred = cv2.GaussianBlur(upscaled, (0, 0), 2.0)
    upscaled = cv2.addWeighted(upscaled, 2.0, blurred, -1.0, 0)
    
    if alpha is not None:
        uh, uw = upscaled.shape[:2]
        upscaled_alpha = cv2.resize(alpha, (uw, uh), interpolation=cv2.INTER_LANCZOS4)
        _, upscaled_alpha = cv2.threshold(upscaled_alpha, 127, 255, cv2.THRESH_BINARY)
        return cv2.merge((upscaled[:,:,0], upscaled[:,:,1], upscaled[:,:,2], upscaled_alpha))
    
    return upscaled