import gradio as gr import cv2 import numpy as np from PIL import Image import torch from transformers import ( CLIPProcessor, CLIPModel, LlamaForCausalLM, LlamaTokenizer, pipeline ) import requests from io import BytesIO import os class ImageStoryteller: def __init__(self): print("Initializing Image Storyteller with CLIP-ViT + LLaMA...") # Load CLIP model for image understanding try: self.clip_model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") self.clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32") print("CLIP-ViT model loaded successfully!") except Exception as e: print(f"CLIP loading failed: {e}") self.clip_model = None self.clip_processor = None # Initialize LLaMA for text generation try: # Using a smaller LLaMA variant that works on Hugging Face Spaces self.llama_model = LlamaForCausalLM.from_pretrained( "huggyllama/llama-7b", # Using a smaller variant torch_dtype=torch.float16, device_map="auto", load_in_8bit=True # For memory efficiency ) self.llama_tokenizer = LlamaTokenizer.from_pretrained("huggyllama/llama-7b") print("LLaMA model loaded successfully!") except Exception as e: print(f"LLaMA loading failed: {e}") # Fallback to smaller model try: self.story_pipeline = pipeline( "text-generation", model="microsoft/DialoGPT-medium", torch_dtype=torch.float32 ) print("Fallback story pipeline initialized!") except Exception as e: print(f"Fallback pipeline failed: {e}") self.story_pipeline = None # Common objects for scene understanding self.common_objects = [ 'person', 'people', 'human', 'man', 'woman', 'child', 'baby', 'dog', 'cat', 'animal', 'bird', 'horse', 'cow', 'sheep', 'car', 'vehicle', 'bus', 'truck', 'bicycle', 'motorcycle', 'building', 'house', 'skyscraper', 'architecture', 'tree', 'forest', 'nature', 'mountain', 'sky', 'clouds', 'water', 'ocean', 'river', 'lake', 'beach', 'food', 'fruit', 'vegetable', 'meal', 'indoor', 'outdoor', 'urban', 'rural' ] # Scene categories for classification self.scene_categories = [ "portrait", "landscape", "cityscape", "indoor scene", "outdoor scene", "nature", "urban", "beach", "mountain", "forest", "street", "party", "celebration", "sports", "action", "still life", "abstract", "art", "architecture", "wildlife", "pet" ] def analyze_image_with_clip(self, image): """Analyze image using CLIP to understand content and scene""" if self.clip_model is None or self.clip_processor is None: return self.fallback_image_analysis(image) try: # Convert PIL to RGB image_rgb = image.convert('RGB') # Analyze objects in the image object_inputs = self.clip_processor( text=self.common_objects, images=image_rgb, return_tensors="pt", padding=True ) with torch.no_grad(): object_outputs = self.clip_model(**object_inputs) object_logits = object_outputs.logits_per_image object_probs = object_logits.softmax(dim=1) # Get top objects top_object_indices = torch.topk(object_probs, 5, dim=1).indices[0] detected_objects = [] for idx in top_object_indices: obj_name = self.common_objects[idx] confidence = object_probs[0][idx].item() if confidence > 0.1: # Confidence threshold detected_objects.append({ 'name': obj_name, 'confidence': confidence }) # Analyze scene type scene_inputs = self.clip_processor( text=self.scene_categories, images=image_rgb, return_tensors="pt", padding=True ) with torch.no_grad(): scene_outputs = self.clip_model(**scene_inputs) scene_logits = scene_outputs.logits_per_image scene_probs = scene_logits.softmax(dim=1) top_scene_indices = torch.topk(scene_probs, 3, dim=1).indices[0] scene_types = [] for idx in top_scene_indices: scene_name = self.scene_categories[idx] confidence = scene_probs[0][idx].item() scene_types.append({ 'type': scene_name, 'confidence': confidence }) return { 'objects': detected_objects, 'scenes': scene_types, 'success': True } except Exception as e: print(f"CLIP analysis failed: {e}") return self.fallback_image_analysis(image) def fallback_image_analysis(self, image): """Fallback image analysis when CLIP fails""" img_np = np.array(image) height, width = img_np.shape[:2] # Simple color-based analysis hsv = cv2.cvtColor(img_np, cv2.COLOR_RGB2HSV) objects = [] scenes = [] # Detect blue areas (sky/water) blue_mask = cv2.inRange(hsv, (100, 50, 50), (130, 255, 255)) if np.sum(blue_mask) > height * width * 0.1: objects.append({'name': 'sky', 'confidence': 0.6}) scenes.append({'type': 'outdoor scene', 'confidence': 0.7}) # Detect green areas (nature) green_mask = cv2.inRange(hsv, (35, 50, 50), (85, 255, 255)) if np.sum(green_mask) > height * width * 0.1: objects.append({'name': 'nature', 'confidence': 0.6}) scenes.append({'type': 'nature', 'confidence': 0.7}) # Detect skin tones (people) skin_mask = cv2.inRange(hsv, (0, 30, 60), (20, 150, 255)) if np.sum(skin_mask) > 1000: objects.append({'name': 'person', 'confidence': 0.5}) scenes.append({'type': 'portrait', 'confidence': 0.6}) # Detect edges (buildings/structures) gray = cv2.cvtColor(img_np, cv2.COLOR_RGB2GRAY) edges = cv2.Canny(gray, 50, 150) if np.sum(edges) > height * width * 0.05: objects.append({'name': 'building', 'confidence': 0.5}) scenes.append({'type': 'urban', 'confidence': 0.6}) return { 'objects': objects, 'scenes': scenes, 'success': False } def create_visualization(self, image, analysis_result): """Create a visualization showing detected elements""" img_np = np.array(image) viz_image = img_np.copy() height, width = img_np.shape[:2] # Add text overlay with analysis results font = cv2.FONT_HERSHEY_SIMPLEX font_scale = 0.6 font_color = (255, 255, 255) background_color = (0, 0, 0) thickness = 2 # Prepare text lines text_lines = ["CLIP-ViT Analysis:"] # Add objects if analysis_result['objects']: text_lines.append("Objects:") for obj in analysis_result['objects'][:3]: # Top 3 objects text_lines.append(f" {obj['name']} ({obj['confidence']:.2f})") # Add scenes if analysis_result['scenes']: text_lines.append("Scene:") for scene in analysis_result['scenes'][:2]: # Top 2 scenes text_lines.append(f" {scene['type']} ({scene['confidence']:.2f})") # Add text to image y_offset = 30 for i, line in enumerate(text_lines): text_size = cv2.getTextSize(line, font, font_scale, thickness)[0] # Add background for text cv2.rectangle(viz_image, (10, y_offset - text_size[1] - 5), (10 + text_size[0] + 10, y_offset + 5), background_color, -1) # Add text cv2.putText(viz_image, line, (15, y_offset), font, font_scale, font_color, thickness) y_offset += 25 return Image.fromarray(viz_image) def generate_narrative_with_llama(self, analysis_result, image_size): """Generate narrative using LLaMA based on CLIP analysis""" # Prepare context from analysis objects_text = ", ".join([obj['name'] for obj in analysis_result['objects'][:5]]) scenes_text = analysis_result['scenes'][0]['type'] if analysis_result['scenes'] else "unknown scene" width, height = image_size # Create prompt for LLaMA prompt = f"""Based on this image analysis: Image Size: {width}x{height} Detected Objects: {objects_text} Scene Type: {scenes_text} Please write a beautiful, descriptive narrative story about this image. Focus on the emotional and visual elements, creating a compelling story that brings the scene to life.""" try: if hasattr(self, 'llama_model') and self.llama_model is not None: # Tokenize input inputs = self.llama_tokenizer(prompt, return_tensors="pt") # Generate with LLaMA with torch.no_grad(): outputs = self.llama_model.generate( inputs.input_ids, max_length=300, temperature=0.7, do_sample=True, top_p=0.9, pad_token_id=self.llama_tokenizer.eos_token_id ) # Decode response narrative = self.llama_tokenizer.decode(outputs[0], skip_special_tokens=True) # Extract only the generated part (after prompt) if narrative.startswith(prompt): narrative = narrative[len(prompt):].strip() return narrative elif hasattr(self, 'story_pipeline') and self.story_pipeline is not None: # Use fallback pipeline result = self.story_pipeline( prompt, max_length=250, temperature=0.7, do_sample=True, pad_token_id=50256 ) return result[0]['generated_text'] except Exception as e: print(f"LLaMA narrative generation failed: {e}") # Fallback narrative return self.fallback_narrative(analysis_result, image_size) def fallback_narrative(self, analysis_result, image_size): """Fallback narrative generation""" width, height = image_size objects = [obj['name'] for obj in analysis_result['objects']] scene = analysis_result['scenes'][0]['type'] if analysis_result['scenes'] else "scene" if 'person' in objects: return f"In this captivating {width}x{height} {scene}, human presence tells a story of connection and experience. " \ f"The composition speaks of moments frozen in time, where light and shadow dance together to reveal " \ f"the beauty of ordinary moments made extraordinary through the lens of perception." elif 'nature' in objects: return f"This breathtaking {width}x{height} natural landscape captures the essence of Earth's timeless beauty. " \ f"Each element harmonizes with the next, creating a symphony of visual poetry that whispers " \ f"ancient stories of growth, change, and the enduring power of the natural world." elif 'building' in objects: return f"Architectural elegance defines this {width}x{height} {scene}, where human ingenuity meets artistic vision. " \ f"The structures stand as silent witnesses to countless stories, their forms telling tales " \ f"of aspiration, community, and the relentless march of progress through time." else: return f"In this compelling {width}x{height} composition, visual elements converge to create a unique narrative. " \ f"The scene invites contemplation, asking viewers to explore the relationships between forms, " \ f"colors, and spaces that together tell a story beyond words." def generate_poetry(self, narrative): """Generate poetic verses based on the narrative""" prompt = f"""Based on this image description: "{narrative}" Create a beautiful 6-line poem that captures the essence and emotion of the scene:""" try: if hasattr(self, 'llama_model') and self.llama_model is not None: inputs = self.llama_tokenizer(prompt, return_tensors="pt") with torch.no_grad(): outputs = self.llama_model.generate( inputs.input_ids, max_length=200, temperature=0.8, do_sample=True, top_p=0.9, pad_token_id=self.llama_tokenizer.eos_token_id ) poetry = self.llama_tokenizer.decode(outputs[0], skip_special_tokens=True) if poetry.startswith(prompt): poetry = poetry[len(prompt):].strip() # Format as 6 lines lines = [line.strip() for line in poetry.split('.') if line.strip()] if len(lines) >= 4: return '\n'.join(lines[:6]) return poetry elif hasattr(self, 'story_pipeline') and self.story_pipeline is not None: result = self.story_pipeline( prompt, max_length=150, temperature=0.8, do_sample=True ) poetry = result[0]['generated_text'] lines = [line.strip() for line in poetry.split('.') if line.strip()] if len(lines) >= 4: return '\n'.join(lines[:6]) return poetry except Exception as e: print(f"Poetry generation failed: {e}") # Fallback poetry return self.fallback_poetry(narrative) def fallback_poetry(self, narrative): """Fallback poetry generation""" if 'person' in narrative.lower(): return """A figure stands where light does fall Their silent story captures all In moments caught by lens and eye Where truth and beauty never die Each breath a verse, each glance a call To understand, to stand in awe""" elif 'nature' in narrative.lower(): return """Where trees reach up to touch the sky And gentle streams go flowing by The earth reveals her ancient art In every leaf, in every part Nature's truth will never die In landscape's soul, we learn to fly""" elif 'building' in narrative.lower(): return """Stone and glass against the blue Tell stories old and stories new Where human hands have shaped the space With vision, time, and careful grace Each structure holds a different view Of dreams that humans can pursue""" else: return """In frames of light and color bold A thousand stories wait untold Each element with voice unique In visual language they all speak Of mysteries that unfold More precious than the purest gold""" def process_image(self, image): """Main processing function""" try: # Analyze image with CLIP-ViT analysis_result = self.analyze_image_with_clip(image) # Generate narrative with LLaMA narrative = self.generate_narrative_with_llama(analysis_result, image.size) # Generate poetry poetry = self.generate_poetry(narrative) # Create visualization viz_image = self.create_visualization(image, analysis_result) return narrative, poetry, viz_image except Exception as e: error_msg = f"An error occurred while processing the image: {str(e)}" return error_msg, "Unable to generate poetry due to processing error.", image # Initialize the storyteller storyteller = ImageStoryteller() # Check for local example images example_images = [] for i in range(1, 10): filename = f"obj_{i:02d}.jpg" if os.path.exists(filename): example_images.append([filename]) print(f"Found example image: {filename}") if not example_images: print("No local example images found, using placeholder") # Create a placeholder if no local images example_images = [[np.ones((300, 300, 3), dtype=np.uint8) * 100]] # Create Gradio interface with gr.Blocks(title="CLIP + LLaMA Image Storyteller", theme="soft") as demo: gr.Markdown("# 🎨 CLIP + LLaMA Image Storyteller") gr.Markdown("**Upload any image and watch AI understand the scene using CLIP-ViT and create beautiful stories with LLaMA!**") with gr.Row(): with gr.Column(): input_image = gr.Image( type="pil", label="🖼️ Upload Your Image", height=300 ) process_btn = gr.Button("✨ Analyze Image & Create Story", variant="primary", size="lg") with gr.Column(): analysis_output = gr.Image( label="🔍 CLIP-ViT Analysis", height=300, show_download_button=True ) with gr.Row(): with gr.Column(): with gr.Tab("📖 Narrative Story"): narrative_output = gr.Textbox( label="Image Narrative", lines=5, max_lines=8, placeholder="Your image's story will appear here...", show_copy_button=True ) with gr.Tab("🎭 Poetic Verses"): poetry_output = gr.Textbox( label="6-Line Poetry", lines=6, max_lines=7, placeholder="Poetic interpretation will appear here...", show_copy_button=True ) # Examples section gr.Markdown("### 🎯 Try These Examples") gr.Examples( examples=example_images, inputs=input_image, outputs=[narrative_output, poetry_output, analysis_output], fn=storyteller.process_image, cache_examples=True ) # How it works section with gr.Accordion("🔍 How It Works", open=False): gr.Markdown(""" **The Magic Behind the Stories:** 1. **CLIP-ViT Analysis**: OpenAI's CLIP model understands image content and scene types 2. **Object Recognition**: Identifies objects, people, scenery with confidence scores 3. **Scene Classification**: Determines the overall scene type (portrait, landscape, urban, etc.) 4. **LLaMA Storytelling**: Meta's LLaMA model generates compelling narratives 5. **Poetic Creation**: Transforms analysis into beautiful 6-line verses **Technical Stack:** - **CLIP-ViT**: Vision transformer for image understanding - **LLaMA**: Large language model for text generation - **Transformers**: Hugging Face library for model inference **Features:** - Semantic image understanding - Context-aware storytelling - Emotional narrative generation - Beautiful poetic interpretations - Real-time analysis visualization **Perfect for:** - Personal photography - Landscape and nature scenes - Urban and architectural photography - Artistic compositions - Memory preservation """) # Set up the processing process_btn.click( fn=storyteller.process_image, inputs=input_image, outputs=[narrative_output, poetry_output, analysis_output] ) # Launch the application if __name__ == "__main__": demo.launch( server_name="0.0.0.0", server_port=7860, share=False )