"""
AI Image Caption Generator - Main Application
Gradio-based web interface for generating image captions using BLIP and GIT models
with customizable styling via Groq API.
"""
import gradio as gr
import time
import numpy as np
from pathlib import Path
from PIL import Image
from typing import Tuple, Optional
# Import our modules
from config import ui_config, performance_config
from src.utils import (
get_image_processor,
get_caption_cache,
get_analytics_manager,
ImageProcessingError
)
from src.models import (
get_model_manager,
get_style_model,
CaptionModelError,
StyleModelError
)
class CaptionGeneratorApp:
"""
Main application class for the caption generator
Manages the Gradio interface and coordinates all components
"""
def __init__(self):
"""Initialize the application"""
print("=" * 60)
print("š INITIALIZING AI IMAGE CAPTION GENERATOR")
print("=" * 60)
# Initialize components
self.image_processor = get_image_processor()
self.model_manager = get_model_manager()
self.style_model = get_style_model()
self.cache = get_caption_cache()
self.analytics = get_analytics_manager()
print("\nā Components initialized")
# Load models
print("\nš¦ Loading AI models (this may take a few minutes on first run)...")
blip_success, git_success = self.model_manager.load_all_models()
if not (blip_success and git_success):
print("\nā ļø Warning: Some models failed to load")
print(f" BLIP: {'ā' if blip_success else 'ā'}")
print(f" GIT: {'ā' if git_success else 'ā'}")
else:
print("\nā All models loaded successfully")
# Check style model
if self.style_model.is_api_available():
print("ā Groq API connected")
else:
print("ā ļø Groq API not available - using fallback styling")
print("\n" + "=" * 60)
print("ā INITIALIZATION COMPLETE")
print("=" * 60 + "\n")
def generate_captions(
self,
image,
style: str,
progress=gr.Progress()
) -> Tuple[str, str, str]:
"""
Generate captions for an image
Args:
image: Image path (str) or PIL Image
style: Style to apply
progress: Gradio progress tracker
Returns:
Tuple[str, str, str]: (blip_caption, git_caption, stats_text)
"""
start_time = time.time()
try:
# Step 1: Validate and preprocess image
progress(0.1, desc="Validating image...")
if image is None:
return (
"ā Error: No image provided",
"ā Error: No image provided",
"ā ļø Please upload an image"
)
# Convert to PIL Image from various formats
try:
if isinstance(image, str):
pil_image = Image.open(image)
elif isinstance(image, Image.Image):
pil_image = image
elif hasattr(image, 'shape'):
import numpy as np
if isinstance(image, np.ndarray):
pil_image = Image.fromarray(image.astype('uint8'))
else:
raise ValueError("Unsupported array type")
else:
return (
f"ā Error: Unsupported image type: {type(image)}",
f"ā Error: Unsupported image type: {type(image)}",
"ā ļø Image format not supported"
)
except Exception as e:
return (
f"ā Error: Cannot load image - {str(e)}",
f"ā Error: Cannot load image - {str(e)}",
"ā ļø Image loading failed"
)
# Validate image
is_valid, error_msg = self.image_processor.validate_image(pil_image)
if not is_valid:
return (
f"ā Error: {error_msg}",
f"ā Error: {error_msg}",
"ā ļø Image validation failed"
)
# Preprocess image
progress(0.2, desc="Processing image...")
processed_img, metadata = self.image_processor.preprocess_image(pil_image)
# Generate image hash for caching
image_hash = self.image_processor.generate_image_hash(processed_img)
# Step 2: Check cache
progress(0.3, desc="Checking cache...")
blip_cached = self.cache.get_caption(image_hash, "blip", style)
git_cached = self.cache.get_caption(image_hash, "git", style)
# Step 3: Generate captions if not cached
raw_captions = {}
if blip_cached is None or git_cached is None:
progress(0.4, desc="Generating captions...")
raw_captions = self.model_manager.generate_captions(processed_img)
# Step 4: Apply styling
progress(0.6, desc=f"Applying {style} style...")
styled_captions = {}
# BLIP caption
if blip_cached:
styled_captions["blip"] = blip_cached
else:
blip_raw = raw_captions.get("blip", "Error generating caption")
styled_captions["blip"] = self.style_model.style_caption(blip_raw, style)
self.cache.set_caption(image_hash, "blip", style, styled_captions["blip"])
# GIT caption
if git_cached:
styled_captions["git"] = git_cached
else:
git_raw = raw_captions.get("git", "Error generating caption")
styled_captions["git"] = self.style_model.style_caption(git_raw, style)
self.cache.set_caption(image_hash, "git", style, styled_captions["git"])
# Step 5: Record analytics
progress(0.9, desc="Finalizing...")
processing_time = time.time() - start_time
# Record for each model
self.analytics.record_caption_generation("blip", style, processing_time / 2, True)
self.analytics.record_caption_generation("git", style, processing_time / 2, True)
# Get stats
stats_text = self.analytics.get_display_stats()
stats_text += f" | ā±ļø This generation: {processing_time:.2f}s"
progress(1.0, desc="Complete!")
return (
styled_captions.get("blip", "Error"),
styled_captions.get("git", "Error"),
stats_text
)
except ImageProcessingError as e:
error_msg = f"ā Image Error: {str(e)}"
return error_msg, error_msg, "ā ļø Image processing failed"
except CaptionModelError as e:
error_msg = f"ā Model Error: {str(e)}"
return error_msg, error_msg, "ā ļø Caption generation failed"
except Exception as e:
error_msg = f"ā Unexpected Error: {str(e)}"
print(f"Error in generate_captions: {e}")
# Record error
self.analytics.record_caption_generation("unknown", style, 0, False)
return error_msg, error_msg, "ā ļø An error occurred"
def create_interface(self) -> gr.Blocks:
"""
Create Gradio interface with enhanced UI and mobile responsiveness
Returns:
gr.Blocks: Configured Gradio interface
"""
with gr.Blocks(
theme=gr.themes.Soft(
primary_hue="purple",
secondary_hue="blue",
neutral_hue="slate",
),
title=ui_config.TITLE,
css=self._get_custom_css()
) as interface:
# Header with enhanced styling
with gr.Row(elem_classes="header-container"):
with gr.Column():
gr.HTML("""
šØ
AI Image Caption Generator
Transform your images into compelling captions with AI-powered models
""")
# Main content area with mobile-responsive layout
with gr.Row(elem_classes="main-content"):
# Left column - Input section
with gr.Column(scale=1, min_width=300, elem_classes="input-column"):
gr.HTML('