app_imagen.py

import gradio as gr
import google.generativeai as genai
from google import genai as google_genai
from google.genai import types
import cv2
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import os
from dotenv import load_dotenv
import requests
import base64
import io
import json
import time
import logging
from typing import Optional, Tuple, List
import warnings
warnings.filterwarnings("ignore")

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

load_dotenv()

try:
    genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
    logger.info("Gemini API configured successfully")
except Exception as e:
    logger.warning(f"Gemini API configuration failed: {e}")

try:
    google_genai.configure(api_key=os.getenv("GEMINI_API_KEY"))
    imagen_client = google_genai.Client()
    logger.info("Imagen client configured successfully")
except Exception as e:
    logger.warning(f"Imagen client configuration failed: {e}")
    imagen_client = None

try:
    from elevenlabs import generate, set_api_key
    if os.getenv("ELEVENLABS_API_KEY"):
        set_api_key(os.getenv("ELEVENLABS_API_KEY"))
        logger.info("ElevenLabs API configured successfully")
except ImportError:
    logger.warning("ElevenLabs not available - voice features disabled")
except Exception as e:
    logger.warning(f"ElevenLabs API configuration failed: {e}")

class BuildTheFutureImagen:
    def __init__(self):
        self.gemini_model = None
        self.yolo_model = None
        self.imagen_client = imagen_client
        self._initialize_models()
        
    def _initialize_models(self):
        try:
            self.gemini_model = genai.GenerativeModel('gemini-2.0-flash-exp')
            logger.info("Gemini model initialized successfully")
        except Exception as e:
            logger.error(f"Failed to initialize Gemini model: {e}")
            self.gemini_model = None
        
    def load_yolo_model(self):
        try:
            from ultralytics import YOLO
            model_path = 'best.pt' if os.path.exists('best.pt') else 'yolov11n.pt'
            self.yolo_model = YOLO(model_path)
            logger.info(f"YOLO model loaded: {model_path}")
            return True
        except Exception as e:
            logger.error(f"Error loading YOLO model: {e}")
            return False
    
    def detect_structures(self, image):
        if self.yolo_model is None:
            if not self.load_yolo_model():
                return image, "YOLO model not available"
        
        try:
            img_cv = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
            results = self.yolo_model(img_cv)
            annotated_img = results[0].plot()
            annotated_pil = Image.fromarray(cv2.cvtColor(annotated_img, cv2.COLOR_BGR2RGB))
            
            return annotated_pil, "Structures detected successfully"
        except Exception as e:
            return image, f"Detection error: {str(e)}"
    
    def create_imagen_prompt(self, style, construction_type="building"):
        construction_contexts = {
            "building": "unfinished construction site with steel beams, concrete foundations, and incomplete walls",
            "bridge": "incomplete bridge construction with support pillars and missing deck sections",
            "road": "road construction site with incomplete pavement, construction equipment, and barriers",
            "general": "unfinished construction project with exposed structural elements"
        }
        
        context = construction_contexts.get(construction_type, construction_contexts["general"])
        
        style_prompts = {
            "realistic": f"Professional architectural photography of a completed {context.replace('unfinished', 'finished').replace('incomplete', 'completed')}. High-quality construction with proper materials, realistic lighting, and professional finishing. 4K HDR photo, architectural photography style, detailed construction work, natural lighting, professional construction standards.",
            
            "futuristic": f"Futuristic high-tech building completion of a {context}. Modern glass facades, smart building technology, solar panels, LED lighting systems, and advanced architectural elements. Sci-fi architecture, 2050 technology, innovative design, high-tech materials, futuristic cityscape, digital art style, cutting-edge construction.",
            
            "artistic": f"Creative artistic completion of a {context}. Unique architectural design with creative materials, colorful elements, artistic touches, and innovative construction techniques. Artistic architecture, creative design, unique materials, colorful construction, innovative building techniques, artistic interpretation, creative engineering."
        }
        
        return style_prompts.get(style, style_prompts["realistic"])
    
    def complete_construction_imagen(self, image, style="realistic", aspect_ratio="1:1", image_size="1K"):
        if self.imagen_client is None:
            return image, "Imagen client not available. Please check your API key."
        
        try:
            construction_type = self.analyze_construction_type(image)
            prompt = self.create_imagen_prompt(style, construction_type)
            config = types.GenerateImagesConfig(
                number_of_images=1,
                sample_image_size=image_size,
                aspect_ratio=aspect_ratio,
                person_generation="dont_allow"
            )
            response = self.imagen_client.models.generate_images(
                model='imagen-3.0-generate-002',
                prompt=prompt,
                config=config
            )
            if response.generated_images:
                generated_image = response.generated_images[0]
                completed_image = generated_image.image
                return completed_image, f"Construction completed successfully with {style} style using Imagen"
            
            return image, "No image generated - please try again"
            
        except Exception as e:
            logger.error(f"Imagen completion error: {e}")
            return image, f"Imagen completion error: {str(e)}"
    
    def complete_construction_gemini(self, image, style="realistic"):
        if self.gemini_model is None:
            return image, "Gemini model not available. Please check your API key."
        
        try:
            style_prompts = {
                "realistic": "Complete this unfinished construction site realistically. Fill in the missing parts with appropriate building materials, proper architectural details, and realistic finishing touches. Make it look like a completed, functional building. Ensure the completion looks natural and follows proper construction practices.",
                "futuristic": "Transform this unfinished construction site into a futuristic, high-tech building. Add modern architectural elements, glass facades, smart building features, solar panels, and futuristic design elements. Make it look like a building from the year 2050 with advanced technology integration.",
                "artistic": "Complete this construction site with artistic and creative architectural elements. Add unique design features, creative materials, colorful elements, and artistic touches that make it visually striking and memorable. Think outside the box with creative architecture."
            }
            
            prompt = style_prompts.get(style, style_prompts["realistic"])
            enhanced_prompt = f"{prompt} The image should maintain the same perspective and lighting as the original. Focus on completing the unfinished parts while maintaining architectural coherence."
            buffered = io.BytesIO()
            image.save(buffered, format='PNG')
            image_bytes = buffered.getvalue()

            gemini_image_model = genai.GenerativeModel('gemini-1.5-flash')
            response = gemini_image_model.generate_content([
                enhanced_prompt,
                {
                    'mime_type': 'image/png',
                    'data': base64.b64encode(image_bytes).decode('utf-8')
                }
            ])

            if getattr(response, 'candidates', None):
                for part in response.candidates[0].content.parts:
                    if hasattr(part, 'inline_data') and part.inline_data and getattr(part.inline_data, 'data', None):
                        image_data = base64.b64decode(part.inline_data.data)
                        completed_image = Image.open(io.BytesIO(image_data)).convert('RGB')
                        return completed_image, f"Construction completed successfully with {style} style using Gemini"

            return image, "No image generated - please try again"
            
        except Exception as e:
            logger.error(f"Gemini completion error: {e}")
            return image, f"Gemini completion error: {str(e)}"
    
    def analyze_construction_type(self, image):
        width, height = image.size
        
        if height > width * 1.5:
            return "bridge"
        elif width > height * 1.5:
            return "road"
        else:
            return "building"
    
    def generate_voice_narration(self, style, model_used="Imagen"):
        if not os.getenv("ELEVENLABS_API_KEY"):
            logger.info("ElevenLabs API key not found - voice narration disabled")
            return None
        
        try:
            narration_texts = {
                "realistic": f"Here's how your construction project will look when completed with realistic finishing touches and proper architectural details. The {model_used} AI has filled in the missing parts with appropriate materials and construction techniques.",
                "futuristic": f"Behold the future! Your construction site has been transformed into a cutting-edge, high-tech building of tomorrow using {model_used}. Notice the modern architectural elements and smart building features.",
                "artistic": f"Watch as your construction site becomes a masterpiece of creative architecture and artistic design with {model_used}. The AI has added unique design features and creative elements that make it visually striking."
            }
            
            text = narration_texts.get(style, narration_texts["realistic"])
            audio = generate(text=text, voice="Rachel", model="eleven_monolingual_v1")
            return audio
        except Exception as e:
            logger.error(f"Voice generation error: {e}")
            return None
    
    def create_comparison_image(self, original, completed):
        if original is None or completed is None:
            return None
        
        try:
            height = min(original.height, completed.height)
            width = min(original.width, completed.width)
            
            original_resized = original.resize((width, height), Image.Resampling.LANCZOS)
            completed_resized = completed.resize((width, height), Image.Resampling.LANCZOS)
            
            comparison_width = width * 2 + 20
            comparison_height = height + 60
            
            comparison_img = Image.new('RGB', (comparison_width, comparison_height), 'white')
            
            comparison_img.paste(original_resized, (0, 30))
            comparison_img.paste(completed_resized, (width + 20, 30))
            
            draw = ImageDraw.Draw(comparison_img)
            try:
                font = ImageFont.truetype("arial.ttf", 20)
            except:
                font = ImageFont.load_default()
            
            draw.text((width//2 - 50, 5), "BEFORE", fill='black', font=font)
            draw.text((width + 20 + width//2 - 50, 5), "AFTER", fill='black', font=font)
            
            return comparison_img
        except Exception as e:
            logger.error(f"Comparison image creation error: {e}")
            return None

app = BuildTheFutureImagen()

def process_image_imagen(uploaded_image, completion_style, model_choice, aspect_ratio, image_size):
    if uploaded_image is None:
        return None, None, None, None, "Please upload an image first", None
    
    try:
        detected_image, detection_msg = app.detect_structures(uploaded_image)
        if model_choice == "Imagen (Recommended)":
            completed_image, completion_msg = app.complete_construction_imagen(
                uploaded_image, completion_style, aspect_ratio, image_size
            )
            model_used = "Imagen"
        else:
            completed_image, completion_msg = app.complete_construction_gemini(
                uploaded_image, completion_style
            )
            model_used = "Gemini"
        
        comparison_image = app.create_comparison_image(uploaded_image, completed_image)
        audio = app.generate_voice_narration(completion_style, model_used)
        
        status_msg = f"✅ Detection: {detection_msg}\n✅ Completion: {completion_msg}"
        if audio is None:
            status_msg += "\n💡 Voice narration is optional (ElevenLabs API key not configured)"
        
        return uploaded_image, detected_image, completed_image, comparison_image, status_msg, audio
        
    except Exception as e:
        logger.error(f"Processing error: {e}")
        return uploaded_image, uploaded_image, uploaded_image, None, f"❌ Error: {str(e)}", None

with gr.Blocks(title="BuildTheFuture: AI-Powered Construction Completion with Imagen", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # BuildTheFuture: AI Construction Completion with Imagen
    
    Upload unfinished construction images and complete them with AI using Imagen models.
    """)
    
    with gr.Row():
        with gr.Column(scale=1):
            image_input = gr.Image(
                label="Upload Unfinished Construction Photo",
                type="pil",
                height=300
            )
            
            with gr.Group():
                gr.Markdown("### Settings")
                
                style_selector = gr.Radio(
                    choices=["realistic", "futuristic", "artistic"],
                    value="realistic",
                    label="Completion Style",
                    info="Choose how you want the construction to be completed"
                )
                
                model_selector = gr.Radio(
                    choices=["Imagen (Recommended)", "Gemini 2.5 Flash"],
                    value="Imagen (Recommended)",
                    label="AI Model",
                    info="Imagen provides higher quality results"
                )
                
                aspect_ratio = gr.Dropdown(
                    choices=["1:1", "3:4", "4:3", "9:16", "16:9"],
                    value="1:1",
                    label="Aspect Ratio",
                    info="Choose the aspect ratio for the completed image"
                )
                
                image_size = gr.Radio(
                    choices=["1K", "2K"],
                    value="1K",
                    label="Image Quality",
                    info="Higher quality takes longer to generate"
                )
            
            process_btn = gr.Button("🚀 Complete Construction", variant="primary", size="lg")
            
            status_text = gr.Textbox(
                label="Status",
                interactive=False,
                lines=4
            )
        
        with gr.Column(scale=2):
            with gr.Tabs():
                with gr.Tab("Original"):
                    original_output = gr.Image(label="Original Unfinished Site", height=400)
                
                with gr.Tab("Detected Structures"):
                    detected_output = gr.Image(label="YOLO Detection Overlay", height=400)
                
                with gr.Tab("AI Completed"):
                    completed_output = gr.Image(label="AI-Completed Construction", height=400)
                
                with gr.Tab("Before vs After"):
                    gr.Markdown("### Interactive Comparison")
                    comparison = gr.Image(
                        label="Drag slider to compare before and after",
                        height=400,
                        show_download_button=True
                    )
    
    with gr.Row():
        audio_output = gr.Audio(
            label="Voice Narration",
            visible=True
        )
    
    process_btn.click(
        fn=process_image_imagen,
        inputs=[image_input, style_selector, model_selector, aspect_ratio, image_size],
        outputs=[original_output, detected_output, completed_output, comparison, status_text, audio_output]
    )
    
    with gr.Row():
        gr.Examples(
            examples=[
                ["samples_imagen/skyscraper_construction.jpg", "realistic", "Imagen (Recommended)", "4:3", "2K"],
                ["samples_imagen/suspension_bridge.jpg", "futuristic", "Imagen (Recommended)", "16:9", "2K"],
                ["samples_imagen/highway_construction.jpg", "realistic", "Imagen (Recommended)", "16:9", "2K"],
                ["samples_imagen/residential_construction.jpg", "artistic", "Imagen (Recommended)", "1:1", "1K"]
            ],
            inputs=[image_input, style_selector, model_selector, aspect_ratio, image_size],
            label="Sample Images"
        )

if __name__ == "__main__":
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=True,
        show_error=True
    )