Update app.py

app.py

@@ -8,9 +8,115 @@ import rembg
 from io import BytesIO
 import os
 
+# class ImageStoryteller:
+#     def __init__(self):
+#         print("Initializing Image Storyteller with CLIP-ViT...")
+        
+#         # 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
+        
+#         # 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)
+import torch
+from transformers import CLIPModel, CLIPProcessor, AutoTokenizer, AutoModelForCausalLM
+from huggingface_hub import login
+from PIL import Image
+import gradio as gr
+
 class ImageStoryteller:
-    def __init__(self):
-        print("Initializing Image Storyteller with CLIP-ViT...")
+    def __init__(self, llm_model_id="microsoft/phi-2"):
+        print("Initializing Image Storyteller with CLIP-ViT and LLM...")
         
         # Load CLIP model for image understanding
         try:
@@ -22,6 +128,26 @@ class ImageStoryteller:
             self.clip_model = None
             self.clip_processor = None
         
+        # Load LLM for story generation
+        try:
+            # For Gemma, you need to login first (uncomment if using Gemma)
+            # login()  # Only for Gemma models
+            
+            # Choose your LLM (phi-2 doesn't require login)
+            self.llm_model_id = llm_model_id
+            self.tokenizer = AutoTokenizer.from_pretrained(llm_model_id)
+            self.llm_model = AutoModelForCausalLM.from_pretrained(
+                llm_model_id,
+                torch_dtype=torch.float16,
+                device_map="auto",
+                trust_remote_code=True if "phi" in llm_model_id else False
+            )
+            print(f"LLM model {llm_model_id} loaded successfully!")
+        except Exception as e:
+            print(f"LLM loading failed: {e}")
+            self.llm_model = None
+            self.tokenizer = None
+        
         # Common objects for scene understanding
         self.common_objects = [
             'person', 'people', 'human', 'man', 'woman', 'child', 'baby',
@@ -110,46 +236,136 @@ class ImageStoryteller:
             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]
+        """Fallback analysis when CLIP fails"""
+        return {
+            'objects': [{'name': 'scene', 'confidence': 1.0}],
+            'scenes': [{'type': 'general image', 'confidence': 1.0}],
+            'success': False
+        }
+    
+    def generate_story_from_analysis(self, analysis_result, creativity_level=0.7):
+        """Generate a story based on detected objects and scene"""
+        if self.llm_model is None:
+            return "Story generation model not available."
         
-        # 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})
+        try:
+            # Extract detected objects and scene
+            objects = [obj['name'] for obj in analysis_result['objects']]
+            scenes = [scene['type'] for scene in analysis_result['scenes']]
+            
+            # Create a prompt for the LLM
+            objects_str = ", ".join(objects[:3])  # Use top 3 objects
+            scene_str = scenes[0] if scenes else "general scene"
+            
+            # Different prompt templates for creativity
+            if creativity_level > 0.8:
+                prompt = f"""Based on this image containing {objects_str} in a {scene_str}, write a creative and imaginative short story (3-4 paragraphs). 
+                Make it engaging and add interesting details about the scene."""
+            elif creativity_level > 0.5:
+                prompt = f"""Create a short story about an image with {objects_str} in a {scene_str}. 
+                Write 2-3 paragraphs that describe what might be happening in this scene."""
+            else:
+                prompt = f"""Describe what you see in an image containing {objects_str} in a {scene_str}. 
+                Write a simple 1-2 paragraph description."""
+            
+            # Format for the specific LLM
+            if "phi" in self.llm_model_id:
+                # Phi-2 specific formatting
+                formatted_prompt = f"Instruct: {prompt}\nOutput:"
+            elif "gemma" in self.llm_model_id:
+                # Gemma specific formatting
+                formatted_prompt = f"<start_of_turn>user\n{prompt}<end_of_turn>\n<start_of_turn>model\n"
+            else:
+                # Generic formatting
+                formatted_prompt = f"Write a story: {prompt}\n\nStory:"
+            
+            # Tokenize and generate
+            inputs = self.tokenizer(formatted_prompt, return_tensors="pt").to(self.llm_model.device)
+            
+            with torch.no_grad():
+                outputs = self.llm_model.generate(
+                    **inputs,
+                    max_new_tokens=250,  # Shorter for faster generation
+                    temperature=creativity_level,
+                    do_sample=True,
+                    top_p=0.9,
+                    repetition_penalty=1.1,
+                    pad_token_id=self.tokenizer.eos_token_id
+                )
+            
+            # Decode and clean up
+            story = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
+            
+            # Remove the prompt from the beginning if present
+            if story.startswith(formatted_prompt):
+                story = story[len(formatted_prompt):].strip()
+            
+            return story
+            
+        except Exception as e:
+            print(f"Story generation failed: {e}")
+            return f"Failed to generate story. Detected objects: {objects_str} in a {scene_str}."
+    
+    def process_image_and_generate_story(self, image, creativity_level=0.7):
+        """Complete pipeline: analyze image and generate story"""
+        print("Analyzing image...")
+        analysis = self.analyze_image_with_clip(image)
+        
+        print("Generating story...")
+        story = self.generate_story_from_analysis(analysis, creativity_level)
+        
+        # Return both analysis and story
+        detected_objects = [obj['name'] for obj in analysis['objects']]
+        scene_type = analysis['scenes'][0]['type'] if analysis['scenes'] else "unknown"
         
         return {
-            'objects': objects,
-            'scenes': scenes,
-            'success': False
+            'detected_objects': detected_objects,
+            'scene_type': scene_type,
+            'story': story,
+            'analysis_success': analysis['success']
         }
+
+    # 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_analysis_overlay(self, image, analysis_result):
         """Create analysis overlay in bottom left with white text on black background"""
@@ -201,28 +417,28 @@ class ImageStoryteller:
         
         return Image.fromarray(overlay)
     
-    def generate_story(self, analysis_result, image_size):
-        """Generate story 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"
+    # def generate_story(self, analysis_result, image_size):
+    #     """Generate story 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
+    #     width, height = image_size
         
-        # Create story based on analysis
-        if 'person' in objects_text.lower():
-            story = f"In this captivating {width}x{height} {scenes_text}, we see {objects_text}. A story unfolds where human presence meets the environment, creating moments of connection and experience that speak to the heart of what it means to be alive in this visual narrative."
+    #     # Create story based on analysis
+    #     if 'person' in objects_text.lower():
+    #         story = f"In this captivating {width}x{height} {scenes_text}, we see {objects_text}. A story unfolds where human presence meets the environment, creating moments of connection and experience that speak to the heart of what it means to be alive in this visual narrative."
         
-        elif 'nature' in objects_text.lower():
-            story = f"This breathtaking {width}x{height} {scenes_text} reveals {objects_text}. Nature's timeless beauty tells a story of growth, change, and the enduring power of the natural world, where every element harmonizes to create a symphony of visual poetry."
+    #     elif 'nature' in objects_text.lower():
+    #         story = f"This breathtaking {width}x{height} {scenes_text} reveals {objects_text}. Nature's timeless beauty tells a story of growth, change, and the enduring power of the natural world, where every element harmonizes to create a symphony of visual poetry."
         
-        elif 'building' in objects_text.lower() or 'urban' in scenes_text.lower():
-            story = f"Architectural elegance defines this {width}x{height} {scenes_text} featuring {objects_text}. The structures stand as silent witnesses to countless stories, their forms telling tales of human ingenuity, community, and the relentless march of progress through time."
+    #     elif 'building' in objects_text.lower() or 'urban' in scenes_text.lower():
+    #         story = f"Architectural elegance defines this {width}x{height} {scenes_text} featuring {objects_text}. The structures stand as silent witnesses to countless stories, their forms telling tales of human ingenuity, community, and the relentless march of progress through time."
         
-        else:
-            story = f"In this compelling {width}x{height} composition showing {objects_text}, visual elements converge to create a unique narrative. The {scenes_text} invites contemplation, asking viewers to explore the relationships between forms, colors, and spaces that together tell a story beyond words."
+    #     else:
+    #         story = f"In this compelling {width}x{height} composition showing {objects_text}, visual elements converge to create a unique narrative. The {scenes_text} invites contemplation, asking viewers to explore the relationships between forms, colors, and spaces that together tell a story beyond words."
         
-        return story
+    #     return story
     
     # def create_story_overlay(self, image, story):
     #     """Create story overlay in bottom left with white text on black background"""
@@ -435,9 +651,9 @@ def load_selected_example(evt: gr.SelectData):
     return None
 
 # Create Gradio interface
-with gr.Blocks(title="CLIP-ViT Image Analyzer", theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 🎨 CLIP-ViT Image Analyzer")
-    gr.Markdown("**Upload an image to analyze content and generate stories**")
+with gr.Blocks(title="Who says AI isn’t creative? Watch it turn a single image into a beautifully written story", theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# Image Story Teller")
+    gr.Markdown("**Upload an image to analyse content and generate stories**")
     
     # Load example images
     example_images_list = get_example_images()