Upload processor/upscale.py with huggingface_hub

processor/upscale.py

@@ -0,0 +1,135 @@
+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