Spaces:
Build error
Build error
File size: 4,533 Bytes
821a664 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 | 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
|