import logging import os from io import BytesIO # Load environment variables from .env if present (helps local dev) try: from dotenv import load_dotenv load_dotenv() except Exception: pass import base64 import cv2 import numpy as np from PIL import Image import google.generativeai as genai log = logging.getLogger(__name__) # Remote inference configuration (Gemini API key only; no Vertex required) DEFAULT_MODEL_ID = os.environ.get("GEMINI_IMAGE_MODEL", "gemini-2.5-flash-image") DEFAULT_PROMPT = os.environ.get( "GEMINI_IMAGE_PROMPT", ( "TASK TYPE: STRICT IMAGE INPAINTING — OBJECT REMOVAL ONLY\n\n" "You are given:\n" "1) An original image\n" "2) A binary mask image\n\n" "MASK RULE (MANDATORY):\n" "• White pixels (#FFFFFF) indicate the exact region to be REMOVED.\n" "• Black pixels (#000000) indicate regions that MUST remain completely unchanged.\n\n" "PRIMARY OBJECTIVE:\n" "Completely delete everything inside the white masked area.\n" "The object in the white region must be fully removed with no visible remnants,\n" "no partial shapes, no outlines, no shadows, and no color traces.\n\n" "INPAINTING INSTRUCTIONS:\n" "Ignore the content of the white masked area entirely.\n" "Reconstruct that region using ONLY surrounding background information.\n" "Extend nearby background textures, patterns, and structures naturally.\n" "Match lighting direction, brightness, contrast, color temperature, and noise.\n" "Continue edges, lines, and surfaces realistically across the removed area.\n" "Blend boundaries smoothly so the edit is visually undetectable.\n\n" "STRICT CONSTRAINTS:\n" "• Do NOT generate or keep any part of the removed object.\n" "• Do NOT invent new objects or details.\n" "• Do NOT repaint, modify, blur, or enhance any black (unmasked) area.\n" "• Do NOT change the original image composition.\n" "• Do NOT change camera angle, perspective, or scale.\n\n" "QUALITY REQUIREMENTS:\n" "• No ghosting or transparent object remains.\n" "• No edge halos or smearing.\n" "• No repeated textures or patchy fills.\n" "• Result must look like the object never existed.\n\n" "FAILURE CONDITIONS (MUST BE AVOIDED):\n" "If any object fragment, outline, shadow, or color from the removed object\n" "is still visible, the result is incorrect and must be re-generated." ), ) _GENAI_MODEL: genai.GenerativeModel | None = None def _resize_mask(mask: np.ndarray, target_hw: tuple[int, int]) -> np.ndarray: """Resize mask to match the target height/width.""" target_h, target_w = target_hw if mask.shape[:2] == (target_h, target_w): return mask return cv2.resize(mask, (target_w, target_h), interpolation=cv2.INTER_NEAREST) def _binary_mask_from_rgba(mask: np.ndarray, invert_mask: bool) -> np.ndarray: """ Normalize incoming RGBA masks to a 0/255 binary mask. - Transparent alpha (0) is treated as "remove" - White/bright RGB is treated as "remove" when alpha is mostly opaque """ if mask.shape[2] == 3: alpha_channel = np.ones(mask.shape[:2], dtype=np.uint8) * 255 rgb_channels = mask else: alpha_channel = mask[:, :, 3] rgb_channels = mask[:, :, :3] # If alpha carries information, prefer it if alpha_channel.mean() < 240: mask_bw = np.where(alpha_channel < 128, 255, 0).astype(np.uint8) else: gray = cv2.cvtColor(rgb_channels, cv2.COLOR_RGB2GRAY) mask_bw = np.where(gray > 128, 255, 0).astype(np.uint8) if not invert_mask: mask_bw = 255 - mask_bw return mask_bw def _pil_to_png_bytes(img: Image.Image) -> bytes: """Encode a PIL image to PNG bytes for Gemini edit endpoints.""" buf = BytesIO() img.save(buf, format="PNG") buf.seek(0) return buf.getvalue() def _get_gemini_model() -> genai.GenerativeModel: global _GENAI_MODEL if _GENAI_MODEL is None: api_key = ( os.environ.get("GEMINI_API_KEY") or os.environ.get("GOOGLE_API_KEY") or os.environ.get("GOOGLE_GENAI_API_KEY") ) if not api_key: raise RuntimeError("Set Gemini API key via GEMINI_API_KEY / GOOGLE_API_KEY / GOOGLE_GENAI_API_KEY") genai.configure(api_key=api_key) model_id = os.environ.get("GEMINI_IMAGE_MODEL", DEFAULT_MODEL_ID) _GENAI_MODEL = genai.GenerativeModel(model_id) return _GENAI_MODEL def _call_gemini_edit( image_rgb: np.ndarray, mask_bw: np.ndarray, prompt: str | None, target_size: tuple[int, int], ) -> Image.Image: """ Send source image + binary mask to Gemini via API-key-only generate_content. We include both the base image and the mask as separate parts and instruct the model to remove masked regions. """ model = _get_gemini_model() base_image = Image.fromarray(image_rgb).convert("RGB") mask_image = Image.fromarray(mask_bw).convert("L") # Build a guidance image where the removal region is painted white for clarity guidance_rgb = image_rgb.copy() guidance_rgb[mask_bw > 0] = 255 guidance_image = Image.fromarray(guidance_rgb).convert("RGB") base_bytes = _pil_to_png_bytes(base_image) mask_bytes = _pil_to_png_bytes(mask_image) guidance_bytes = _pil_to_png_bytes(guidance_image) # Enrich prompt to explicitly describe the two images being sent effective_prompt = ( (prompt or DEFAULT_PROMPT).strip() + "\nIMAGE ORDER:\n" + "Image A: Original photo with the removal region painted white.\n" + "Image B: Binary mask (white=remove, black=keep). Use this mask to decide what to remove.\n" ) log.info( "Calling Gemini generate_content model=%s (mask-guided remove) mask_pixels=%d", model.model_name, int((mask_bw > 0).sum()), ) # Build content parts: prompt + guidance image + mask image (explicit order) content = [ effective_prompt, {"mime_type": "image/png", "data": guidance_bytes}, {"mime_type": "image/png", "data": mask_bytes}, ] response = model.generate_content(content, stream=False) output_img: Image.Image | None = None # Extract first image from response parts try: for candidate in getattr(response, "candidates", []): parts = getattr(candidate, "content", None) if not parts or not getattr(parts, "parts", None): continue for part in parts.parts: inline = getattr(part, "inline_data", None) if inline and inline.data: data = inline.data if isinstance(data, str): data = base64.b64decode(data) output_img = Image.open(BytesIO(data)).convert("RGB") break if output_img: break except Exception as err: log.warning("Failed to parse Gemini response image: %s", err) if output_img is None: raise RuntimeError("Gemini generate_content returned no image") # Ensure output matches original dimensions if Gemini rescaled if output_img.size != target_size: output_img = output_img.resize(target_size, Image.Resampling.LANCZOS) return output_img def process_inpaint( image: np.ndarray, mask: np.ndarray, invert_mask: bool = True, prompt: str | None = None, ) -> np.ndarray: """ Forward inpainting to Gemini edit API using source image + mask. """ image_rgba = Image.fromarray(image).convert("RGBA") image_rgb = np.array(image_rgba.convert("RGB")) mask_rgba = np.array(Image.fromarray(mask).convert("RGBA")) mask_bw = _binary_mask_from_rgba(mask_rgba, invert_mask) mask_bw = _resize_mask(mask_bw, image_rgb.shape[:2]) target_size = (image_rgb.shape[1], image_rgb.shape[0]) # (width, height) edited_image = _call_gemini_edit(image_rgb, mask_bw, prompt, target_size) return np.array(edited_image)