import json import os import time import uuid import tempfile from PIL import Image, ImageDraw, ImageFont import gradio as gr import base64 import mimetypes from google import genai from google.genai import types def save_binary_file(file_name, data): with open(file_name, "wb") as f: f.write(data) def generate(text, file_name, api_key, model="gemini-2.0-flash-exp"): # Initialize client using provided api_key (or fallback to env variable) client = genai.Client(api_key=(api_key.strip() if api_key and api_key.strip() != "" else os.environ.get("GEMINI_API_KEY"))) try: print("Uploading file to Gemini API...") files = [ client.files.upload(file=file_name) ] contents = [ types.Content( role="user", parts=[ types.Part.from_uri( file_uri=files[0].uri, mime_type=files[0].mime_type, ), types.Part.from_text(text=text), ], ), ] generate_content_config = types.GenerateContentConfig( temperature=0, # Lower temperature for more consistent, conservative results top_p=0.92, max_output_tokens=8192, response_modalities=["image", "text"], response_mime_type="text/plain", # Additional parameters to encourage subtle, natural results safety_settings=[ { "category": "HARM_CATEGORY_DANGEROUS_CONTENT", "threshold": "BLOCK_MEDIUM_AND_ABOVE" } ] ) text_response = "" image_path = None # Create a temporary file to potentially store image data with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp: temp_path = tmp.name print("Sending request to Gemini API...") # Add a timeout to prevent indefinite waiting start_time = time.time() max_wait_time = 60 # Maximum wait time in seconds try: stream = client.models.generate_content_stream( model=model, contents=contents, config=generate_content_config, ) for chunk in stream: # Check for timeout if time.time() - start_time > max_wait_time: print("Gemini API request timed out after", max_wait_time, "seconds") break if not chunk.candidates or not chunk.candidates[0].content or not chunk.candidates[0].content.parts: continue candidate = chunk.candidates[0].content.parts[0] # Check for inline image data if candidate.inline_data: save_binary_file(temp_path, candidate.inline_data.data) print(f"Smile enhancement image generated: {temp_path}") image_path = temp_path # If an image is found, we assume that is the desired output. break else: # Accumulate text response if no inline_data is present. text_response += chunk.text + "\n" print("Received text response from Gemini API") except Exception as e: print(f"Error during content generation: {str(e)}") # Continue with the function, returning empty responses except Exception as e: print(f"Error in Gemini API setup: {str(e)}") return None, f"Error: {str(e)}" finally: # Always clean up files try: if 'files' in locals() and files: del files except: pass return image_path, text_response def assess_image_quality(original_image, enhanced_image): """ Assesses the quality of the enhanced image based on specific criteria. Returns a tuple of (is_acceptable, feedback_message) """ try: # Check if enhanced image exists if enhanced_image is None: return False, "No enhanced image generated" # Image dimension checks if enhanced_image.size[0] < 100 or enhanced_image.size[1] < 100: return False, "Enhanced image appears to be too small or improperly sized" # Check that the enhanced image has similar dimensions to the original # This helps ensure facial proportions are maintained width_diff = abs(original_image.size[0] - enhanced_image.size[0]) height_diff = abs(original_image.size[1] - enhanced_image.size[1]) # If dimensions are significantly different, it suggests the image proportions changed if width_diff > 20 or height_diff > 20: return False, "Enhanced image dimensions differ significantly from original, suggesting facial proportions may have changed" # Check image has proper RGB channels for natural skin tones if enhanced_image.mode != 'RGB': return False, "Enhanced image does not have the correct color mode" # For now, we'll do basic checks and assume the model follows guidelines return True, "Image passes quality assessment criteria" except Exception as e: print(f"Error in quality assessment: {str(e)}") # Default to not accepting the image if assessment fails return False, f"Assessment error: {str(e)}" def process_smile_enhancement(input_image, max_attempts=2): try: if input_image is None: return None, "", "" # Get API key from environment variable gemini_api_key = "AIzaSyCVzRDxkuvtaS1B22F_F-zl0ehhXR0nuU8" if not gemini_api_key: print("Error: GEMINI_API_KEY not found in environment variables") return [input_image], "", "API key not configured" # Save the input image to a temporary file with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as tmp: input_path = tmp.name input_image.save(input_path) print(f"Input image saved to {input_path}") # Initialize attempt counter and result variables current_attempt = 0 result_img = None feedback_history = [] max_processing_time = 120 # Maximum time in seconds for overall processing start_processing_time = time.time() while current_attempt < max_attempts: # Check if overall processing time exceeded if time.time() - start_processing_time > max_processing_time: print(f"Overall processing time exceeded {max_processing_time} seconds") break current_attempt += 1 print(f"Starting processing attempt {current_attempt}/{max_attempts}...") # Create a comprehensive prompt for true smile enhancement that affects facial features naturally # Adjust prompt based on previous attempts if needed prompt = """ Create a naturally enhanced smile that suits this specific person's face and character. Make the following personalized improvements: - Slightly enhance the existing teeth while PRESERVING their natural color, spacing, and individual characteristics - DO NOT make teeth perfectly white or perfectly aligned - keep some natural variation and character - Create subtle, natural smile lines around the eyes (crow's feet) appropriate for this person's age and face - Slightly raise the cheeks WITHOUT widening the face - Add a slight narrowing of the eyes that happens in genuine smiles - Create subtle dimples ONLY if they already exist in the original image - Enhance the overall joyful expression while maintaining the person's unique facial structure IMPORTANT GUIDELINES: - PRESERVE THE PERSON'S NATURAL DENTAL CHARACTERISTICS - teeth should still look like THEIR teeth, just slightly enhanced - Keep teeth coloration natural and appropriate for the person - avoid unnaturally white teeth - Maintain slight natural imperfections in tooth alignment that give character to the smile - Create a genuine, authentic-looking smile that affects the entire face naturally - ABSOLUTELY CRITICAL: DO NOT widen the face or change face width at all - Preserve the person's identity completely (extremely important) - Preserve exact facial proportions and face width of the original image - Maintain natural-looking results appropriate for the person's age and face structure - Keep teeth proportionate to the face - avoid making them too large or prominent - Maintain proper tooth-to-face ratio and ensure teeth fit naturally within the mouth - Keep the original background, lighting, and image quality intact - Ensure the enhanced smile looks natural, genuine, and believable - Create a smile that looks like a moment of true happiness for THIS specific person - Remember that not everyone has or wants perfect white teeth - the enhancements should SUIT THE INDIVIDUAL - If teeth are enhanced, maintain their natural characteristics while making subtle improvements """ # If not the first attempt, add previous feedback to the prompt if current_attempt > 1 and feedback_history: prompt += """ IMPORTANT FEEDBACK FROM PREVIOUS ATTEMPT: """ + " ".join(feedback_history) + """ Please address these issues in this new attempt. """ # Process silently print(f"Processing attempt {current_attempt}/{max_attempts}...") # Set timeout for individual API call api_call_timeout = time.time() + 45 # 45 second timeout for API call try: # Process the image using Google's Gemini model with timeout image_path, text_response = generate(text=prompt, file_name=input_path, api_key=gemini_api_key) # Check if API call timeout occurred if time.time() > api_call_timeout: print("API call timeout occurred") feedback_history.append("API call timed out, trying again with simplified request.") continue print(f"API response received: Image path: {image_path is not None}, Text length: {len(text_response)}") if image_path: # Load and convert the image if needed try: result_img = Image.open(image_path) if result_img.mode == "RGBA": result_img = result_img.convert("RGB") print("Successfully loaded generated image") # Assess the quality of the enhanced image is_acceptable, assessment_feedback = assess_image_quality(input_image, result_img) print(f"Image quality assessment: {is_acceptable}, {assessment_feedback}") if is_acceptable: # Return the enhanced image with success message success_message = "Successfully loaded generated image\nImage quality assessment: True, Image passes quality assessment criteria" return [result_img], "", success_message else: # Image didn't pass quality assessment, add feedback for next attempt feedback_history.append(assessment_feedback) # If we've reached max attempts, return the best result so far if current_attempt >= max_attempts: print("Max attempts reached, returning best result") return [result_img], "", "" except Exception as img_error: print(f"Error processing the generated image: {str(img_error)}") feedback_history.append(f"Error with image: {str(img_error)}") else: # No image was generated, only text response print("No image was generated, only text response") feedback_history.append("No image was generated in the previous attempt.") # If we've reached max attempts, return the original image if current_attempt >= max_attempts: print("Max attempts reached, returning original image") return [input_image], "", "" except Exception as gen_error: print(f"Error during generation attempt {current_attempt}: {str(gen_error)}") feedback_history.append(f"Error during processing: {str(gen_error)}") # If we've reached max attempts, return the original image if current_attempt >= max_attempts: return [input_image], "", "" # Return the original image as a fallback without messages print("Returning original image as fallback") return [input_image], "", "" except Exception as e: # Return the original image silently on error print(f"Overall error in process_smile_enhancement: {str(e)}") return [input_image], "", "" # Create a clean interface with minimal UI elements with gr.Blocks(title="Smile Enhancement") as demo: with gr.Row(): with gr.Column(): image_input = gr.Image( type="pil", label=None, image_mode="RGB", elem_classes="upload-box" ) submit_btn = gr.Button("Enhance Smile with Natural Expressions", elem_classes="generate-btn") with gr.Column(): output_gallery = gr.Gallery(label=None) # Show feedback text to display assessment results feedback_text = gr.Textbox(label="Status", visible=True) # Hidden element for structure output_text = gr.Textbox(visible=False) submit_btn.click( fn=process_smile_enhancement, inputs=[image_input], outputs=[output_gallery, output_text, feedback_text] ) demo.queue(max_size=50).launch(share=True)