import os # Fix for permissions on Hugging Face Spaces os.environ['TRANSFORMERS_CACHE'] = '/tmp/hf' os.environ['HF_HOME'] = '/tmp/hf' os.environ['XDG_CACHE_HOME'] = '/tmp' os.environ['STREAMLIT_HOME'] = '/tmp' os.makedirs('/tmp/hf', exist_ok=True) import streamlit as st import torch import numpy as np from PIL import Image, ImageEnhance import io import requests from transformers import ( BlipForConditionalGeneration, BlipProcessor, VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer, CLIPProcessor, CLIPModel, AutoModelForCausalLM, AutoProcessor ) from deep_translator import GoogleTranslator from scipy.ndimage import variance from concurrent.futures import ThreadPoolExecutor # ......................... PAGE CONFIGURATION .......................... st.set_page_config( page_title="đŸ–ŧī¸ AI Image Caption Generator", layout="wide", initial_sidebar_state="expanded" ) # .......................... MODEL CONFIGURATION .................... MODEL_CONFIGS = { "BLIP": { "name": "BLIP", "icon": "⭐", "description": "BLIP (Bootstrapping Language-Image Pre-training) is designed to learn vision-language representation from noisy web data. It excels at generating detailed and accurate image descriptions.", "generate_params": {"max_length": 50, "num_beams": 5, "min_length": 10, "do_sample": True, "top_p": 0.9, "repetition_penalty": 1.5} # Added do_sample=True }, "ViT-GPT2": { "name": "ViT-GPT2", "icon": "⭐", "description": "ViT-GPT2 combines Vision Transformer for image encoding with GPT2 for text generation. It's effective at capturing visual details and creating fluent natural language descriptions.", "generate_params": {"max_length": 50, "num_beams": 5, "min_length": 10, "repetition_penalty": 1.5} }, "GIT": { "name": "GIT-base", "icon": "⭐", "description": "GIT (Generative Image-to-text Transformer) is designed specifically for image captioning tasks, focusing on generating coherent and contextually relevant descriptions.", "generate_params": {"max_length": 50, "num_beams": 4, "min_length": 8, "repetition_penalty": 1.5} }, "CLIP": { "name": "CLIP", "icon": "⭐", "description": "CLIP (Contrastive Language-Image Pre-training) analyzes images across multiple dimensions including content type, scene attributes, and photographic style.", } } # ......................... LOADING FUNCTIONS ..................................... @st.cache_resource def load_blip_model(): model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base") # Changed to base model processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base") return model, processor @st.cache_resource def load_vit_gpt2_model(): model = VisionEncoderDecoderModel.from_pretrained("nlpconnect/vit-gpt2-image-captioning") feature_extractor = ViTImageProcessor.from_pretrained("nlpconnect/vit-gpt2-image-captioning") tokenizer = AutoTokenizer.from_pretrained("nlpconnect/vit-gpt2-image-captioning") return model, feature_extractor, tokenizer @st.cache_resource def load_git_model(): processor = AutoProcessor.from_pretrained("microsoft/git-base") model = AutoModelForCausalLM.from_pretrained("microsoft/git-base") return model, processor @st.cache_resource def load_clip_model(): processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32") # Changed to smaller model model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") return model, processor # ......................... IMAGE PROCESSING ............................... def preprocess_image(image): max_size = 1024 if max(image.size) > max_size: ratio = max_size / max(image.size) new_size = (int(image.size[0] * ratio), int(image.size[1] * ratio)) image = image.resize(new_size, Image.LANCZOS) enhancer = ImageEnhance.Contrast(image) image = enhancer.enhance(1.2) img_array = np.array(image.convert('L')) if np.mean(img_array) < 100: brightness_enhancer = ImageEnhance.Brightness(image) image = brightness_enhancer.enhance(1.3) return image def check_image_quality(image): if image.width < 200 or image.height < 200: return False, "Image is too small. Please use a bigger image." img_array = np.array(image.convert('L')) if variance(img_array) < 100: return False, "Image might be too blurry. Please use a clearer image." return True, "Image looks good for captioning." # .................... CAPTION GENERATION FUNCTIONS .................. def generate_caption(image, model_name, models_data): if model_name == "BLIP": model, processor = models_data[model_name] return get_blip_caption(image, model, processor) elif model_name == "ViT-GPT2": model, feature_extractor, tokenizer = models_data[model_name] return get_vit_gpt2_caption(image, model, feature_extractor, tokenizer) elif model_name == "GIT": model, processor = models_data[model_name] return get_git_caption(image, model, processor) elif model_name == "CLIP": model, processor = models_data[model_name] return get_clip_caption(image, model, processor) return "Model not supported" def get_blip_caption(image, model, processor): try: inputs = processor(images=image, return_tensors="pt", padding=True, truncation=True) output = model.generate(**inputs, **MODEL_CONFIGS["BLIP"]["generate_params"]) caption = processor.decode(output[0], skip_special_tokens=True) return caption except Exception as e: return f"BLIP model error: {str(e)}" def get_vit_gpt2_caption(image, model, feature_extractor, tokenizer): try: inputs = feature_extractor(images=image, return_tensors="pt", padding=True) output = model.generate( pixel_values=inputs.pixel_values, **MODEL_CONFIGS["ViT-GPT2"]["generate_params"], attention_mask=inputs.attention_mask if hasattr(inputs, "attention_mask") else None ) caption = tokenizer.decode(output[0], skip_special_tokens=True) return caption except Exception as e: return f"ViT-GPT2 model error: {str(e)}" def get_git_caption(image, model, processor): try: inputs = processor(images=image, return_tensors="pt", padding=True) output = model.generate(**inputs, **MODEL_CONFIGS["GIT"]["generate_params"]) caption = processor.decode(output[0], skip_special_tokens=True) return caption except Exception as e: return f"GIT model error: {str(e)}" # .................... CLIP CATEGORIES ................ CONTENT_CATEGORIES = [ "a portrait photograph", "a landscape photograph", "a wildlife photograph", "an architectural photograph", "a street photograph", "a food photograph", "a fashion photograph", "a sports photograph", "a macro photograph", "a night photograph", "an aerial photograph", "an underwater photograph", "a product photograph", "a documentary photograph", "a travel photograph", "a black and white photograph", "an abstract photograph", "a concert photograph", "a wedding photograph", "a nature photograph" ] SCENE_ATTRIBUTES = [ "indoors", "outdoors", "daytime", "nighttime", "urban", "rural", "beach", "mountains", "forest", "desert", "snowy", "rainy", "foggy", "sunny", "crowded", "empty", "modern", "vintage", "colorful", "minimalist" ] STYLE_ATTRIBUTES = [ "professional", "casual", "artistic", "documentary", "aerial view", "close-up", "wide-angle", "telephoto", "panoramic", "HDR", "long exposure", "shallow depth of field", "silhouette", "motion blur" ] def get_clip_caption(image, model, processor): try: content_inputs = processor(text=CONTENT_CATEGORIES, images=image, return_tensors="pt", padding=True) content_outputs = model(**content_inputs) content_probs = content_outputs.logits_per_image.softmax(dim=1)[0] top_content_probs, top_content_indices = torch.topk(content_probs, 2) scene_inputs = processor(text=SCENE_ATTRIBUTES, images=image, return_tensors="pt", padding=True) scene_outputs = model(**scene_inputs) scene_probs = scene_outputs.logits_per_image.softmax(dim=1)[0] top_scene_probs, top_scene_indices = torch.topk(scene_probs, 2) style_inputs = processor(text=STYLE_ATTRIBUTES, images=image, return_tensors="pt", padding=True) style_outputs = model(**style_inputs) style_probs = style_outputs.logits_per_image.softmax(dim=1)[0] top_style_probs, top_style_indices = torch.topk(style_probs, 1) primary_content = CONTENT_CATEGORIES[top_content_indices[0].item()].replace("a ", "") primary_scene = SCENE_ATTRIBUTES[top_scene_indices[0].item()] primary_style = STYLE_ATTRIBUTES[top_style_indices[0].item()] secondary_elements = [] if top_content_probs[1].item() > 0.15: secondary_content = CONTENT_CATEGORIES[top_content_indices[1].item()].replace("a ", "") secondary_elements.append(f"with elements of {secondary_content}") if top_scene_probs[1].item() > 0.15: secondary_scene = SCENE_ATTRIBUTES[top_scene_indices[1].item()] secondary_elements.append(f"also showing {secondary_scene} characteristics") detailed_caption = f"This looks like {CONTENT_CATEGORIES[top_content_indices[0].item()]} in a {primary_scene} setting" if secondary_elements: detailed_caption += ", " + " ".join(secondary_elements) detailed_caption += f". The image has a {primary_style} look." detailed_caption += f" (Main type: {top_content_probs[0].item()*100:.1f}% sure)" return detailed_caption except Exception as e: return f"CLIP model error: {str(e)}" # ......................... TRANSLATION FUNCTION ....................... def batch_translate(texts, target_lang): try: translator = GoogleTranslator(source='en', target=target_lang) return {key: translator.translate(value) for key, value in texts.items()} except Exception as e: return {key: f"Translation error: {str(e)}" for key in texts} # .......................... GUI STYLE ............................. def apply_styles(): st.markdown(""" """, unsafe_allow_html=True) # ............................ COMPONENTS ....................... def display_sidebar(): with st.sidebar: st.markdown('

📘 About This App

', unsafe_allow_html=True) st.markdown('
This application is part of our NLP project, focused on generating captions for images using four pretrained models to generate captions, each offering a different approach and style to describe the image content.
', unsafe_allow_html=True) st.markdown('

🤖 AI Models

', unsafe_allow_html=True) st.markdown('''
⭐BLIP: Detailed descriptions
⭐ViT-GPT2: Smooth & concise
⭐GIT: Scene understanding
⭐CLIP: Image categorization
''', unsafe_allow_html=True) st.markdown('

🔧 Tech Used

', unsafe_allow_html=True) st.markdown('''
Streamlit
Hugging Face
PyTorch
Google Translator
''', unsafe_allow_html=True) with st.expander("📊 Model Comparison", expanded=False): st.markdown(''' | Model | Strength | Best Use | |----------|--------------------|-------------------| | BLIP | Detailed captions | Deep analysis | | ViT-GPT2 | Smooth captions | Quick summaries | | GIT | Scene storytelling | Context insight | | CLIP | Categorization | Filtering & tags | ''', unsafe_allow_html=True) def image_input_section(): with st.container(): st.markdown('

🌄 Image Input

', unsafe_allow_html=True) with st.container(): st.markdown('
', unsafe_allow_html=True) input_option = st.radio("Choose input method:", ["Upload Image", "Image URL"], horizontal=True) image = None if input_option == "Upload Image": uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"], label_visibility="collapsed") if uploaded_file is not None: try: image = Image.open(uploaded_file).convert("RGB") except Exception as e: st.error(f"Error opening image: {e}") else: url = st.text_input("Enter Image URL", placeholder="https://example.com/image.jpg", label_visibility="collapsed") if url: try: response = requests.get(url) if response.status_code == 200 and 'image' in response.headers.get('Content-Type', ''): image = Image.open(io.BytesIO(response.content)).convert("RGB") else: st.error("Not a valid image URL.") except Exception as e: st.error(f"Error loading image from URL: {e}") st.markdown('
', unsafe_allow_html=True) return image def model_selection_and_display(image): if image: with st.container(): col_models, col_image = st.columns([2, 3]) with col_models: st.markdown('

⚙ī¸ Select Models

', unsafe_allow_html=True) with st.container(): st.markdown('
', unsafe_allow_html=True) use_blip = st.checkbox("BLIP (Bootstrapping Language-Image Pre-training)", value=True) use_vit_gpt2 = st.checkbox("ViT-GPT2 (ViT-GPT2 combines Vision Transformer)", value=True) use_git = st.checkbox("GIT (Generative Image-to-text Transformer)", value=True) use_clip = st.checkbox("CLIP (Contrastive Language-Image Pre-training)", value=True) with st.expander("TRANSLATION LANGUAGES"): translation_language = st.selectbox( "Translation Language", ["Arabic", "French", "Spanish", "Chinese", "Russian", "German"], index=0 ) language_code_map = { "Arabic": "ar", "French": "fr", "Spanish": "es", "Chinese": "zh", "Russian": "ru", "German": "de" } selected_lang_code = language_code_map[translation_language] st.markdown("
", unsafe_allow_html=True) generate_button = st.button("✨ Generate Captions", type="primary") st.markdown('
', unsafe_allow_html=True) with col_image: with st.spinner("Processing image..."): quality_ok, quality_message = check_image_quality(image) if not quality_ok: st.warning(quality_message) processed_image = preprocess_image(image) st.markdown('
', unsafe_allow_html=True) st.image(processed_image, caption="PROCESSED IMAGE ", use_container_width=True) st.markdown('
', unsafe_allow_html=True) return { "generate_button": generate_button, "use_blip": use_blip, "use_vit_gpt2": use_vit_gpt2, "use_git": use_git, "use_clip": use_clip, "selected_lang_code": selected_lang_code, "translation_language": translation_language, "processed_image": processed_image } return None def generate_and_display_captions(config): if not config or not config["generate_button"]: return selected_models = [] if config["use_blip"]: selected_models.append("BLIP") if config["use_vit_gpt2"]: selected_models.append("ViT-GPT2") if config["use_git"]: selected_models.append("GIT") if config["use_clip"]: selected_models.append("CLIP") if not selected_models: st.warning("Please CHOOSE at least one model.") return with st.spinner("Loading models..."): models_data = {} if config["use_blip"]: models_data["BLIP"] = load_blip_model() if config["use_vit_gpt2"]: models_data["ViT-GPT2"] = load_vit_gpt2_model() if config["use_git"]: models_data["GIT"] = load_git_model() if config["use_clip"]: models_data["CLIP"] = load_clip_model() with st.spinner("Creating captions... Please wait... "): captions = {} with ThreadPoolExecutor(max_workers=min(len(selected_models), 4)) as executor: future_to_model = { executor.submit(generate_caption, config["processed_image"], model_name, models_data): model_name for model_name in selected_models } for future in future_to_model: model_name = future_to_model[future] try: caption = future.result() captions[model_name] = caption except Exception as e: captions[model_name] = f"Error with caption: {str(e)}" with st.spinner(f"Translating to {config['translation_language']}..."): translations = batch_translate(captions, config["selected_lang_code"]) display_captions_in_tabs(captions, translations, config["selected_lang_code"], config["translation_language"]) # Caption Comparison if len(captions) > 1: display_caption_comparison(captions) def display_captions_in_tabs(captions, translations, selected_lang_code, translation_language): st.markdown('

📝 Generated Captions

', unsafe_allow_html=True) model_colors = { "BLIP": "#2d3d7d", "ViT-GPT2": "#2d3d7d", "GIT": "#2d3d7d", "CLIP": "#2d3d7d" } tabs = st.tabs([f"{MODEL_CONFIGS[model_name]['icon']} {model_name}" for model_name in captions]) rtl_languages = ["ar"] text_dir = "rtl" if selected_lang_code in rtl_languages else "ltr" for i, model_name in enumerate(captions): with tabs[i]: st.markdown('
', unsafe_allow_html=True) eng_col, trans_col = st.columns(2) with eng_col: st.markdown(f"**đŸ‡Ŧ🇧 English Caption:**") st.markdown(f"""
{captions[model_name]}
""", unsafe_allow_html=True) with trans_col: lang_flags = { "ar": "đŸ‡ĒđŸ‡Ŧ", "fr": "đŸ‡Ģ🇷", "es": "đŸ‡Ē🇸", "zh": "🇨đŸ‡ŗ", "ru": "🇷đŸ‡ē", "de": "🇩đŸ‡Ē" } st.markdown(f"**{lang_flags.get(selected_lang_code, '🌐')} {translation_language} Translation:**") st.markdown(f"""
{translations[model_name]}
""", unsafe_allow_html=True) st.markdown('
', unsafe_allow_html=True) with st.expander("ℹī¸ About this model"): st.markdown(MODEL_CONFIGS[model_name]["description"]) def display_caption_comparison(captions): st.markdown('

⭐ Compare All Captions

', unsafe_allow_html=True) st.markdown('
', unsafe_allow_html=True) model_colors = { "BLIP": "#2d3d7d", "ViT-GPT2": "#2d3d7d", "GIT": "#2d3d7d", "CLIP": "#2d3d7d" } for model_name, caption in captions.items(): st.markdown(f"""
{MODEL_CONFIGS[model_name]['icon']} {model_name}
{caption}
""", unsafe_allow_html=True) st.markdown('
', unsafe_allow_html=True) # ................................... MAIN APPLICATION ............................ def main(): # Apply CSS apply_styles() # Display title and info st.markdown('

🌌 AI Image Caption Generator

', unsafe_allow_html=True) st.markdown('
Generate, translate, and compare image captions easily.
', unsafe_allow_html=True) # Sidebar display_sidebar() # Image input image = image_input_section() # Model selection and image display if image: config = model_selection_and_display(image) # Generate and display captions if config: generate_and_display_captions(config) if __name__ == "__main__": main()