import gradio as gr
from transformers import BlipProcessor, BlipForConditionalGeneration
import torch
from PIL import Image
from deep_translator import GoogleTranslator
from gradio.themes import Base

# Load BLIP model and processor
caption_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-large")
caption_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-large")

# Translator
translator = GoogleTranslator(source='en', target='hi')

# Device setup
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
caption_model.to(device)

def generate_caption(image):
    try:
        # Preprocess the image
        inputs = caption_processor(images=image, return_tensors="pt")
        pixel_values = inputs.pixel_values.to(device)
        
        # Generate caption
        output_ids = caption_model.generate(
            pixel_values,
            max_length=30,
            num_beams=10,
            num_beam_groups=2,
            diversity_penalty=0.5,
            repetition_penalty=2.0,
            temperature=0.6,
            top_k=50,
            top_p=0.95,
            no_repeat_ngram_size=3
        )
        
        # Decode English caption
        english_caption = caption_processor.decode(output_ids[0], skip_special_tokens=True)
        
        # Translate to Hindi
        hindi_caption = translator.translate(english_caption)
        
        return english_caption, hindi_caption
    
    except Exception as e:
        return f"Error: {str(e)}", f"Error: {str(e)}"

# Custom theme with a blue and white color scheme
custom_theme = gr.themes.Default(
    primary_hue="blue",         # Main color (buttons, highlights)
    secondary_hue="gray",       # Secondary elements
    neutral_hue="slate",        # Backgrounds, borders
    text_size="lg",             # Larger text for readability
    radius_size="md",           # Rounded corners
    font=[gr.themes.GoogleFont("Roboto"), "sans-serif"]  # Modern font
)

# Gradio interface with improved visuals
interface = gr.Interface(
    fn=generate_caption,
    inputs=gr.Image(type="pil", label="Upload an Image"),
    outputs=[
        gr.Textbox(label="English Caption", lines=2, placeholder="English caption will appear here..."),
        gr.Textbox(label="Hindi Caption", lines=2, placeholder="हिंदी कैप्शन यहाँ दिखाई देगा...")
    ],
    title="Image Caption Generator (English & Hindi)",
    description="Upload an image to generate captions in English and Hindi with a sleek, modern interface.",
    theme=custom_theme,
    css="""
        .gradio-container { max-width: 800px; margin: auto; }
        h1 { text-align: center; color: #1E40AF; }
        .label { font-weight: bold; }
        input, output { border-radius: 8px; }
    """
)

interface.launch()