Update app.py

app.py

@@ -24,11 +24,126 @@ from io import BytesIO
 
 class ImageScraper:
     def __init__(self):
-        self.PIXABAY_API_KEY = "48069976-37e20099248207cee12385560"  # Your API key
+        self.PIXABAY_API_KEY = "48069976-37e20099248207cee12385560"
         self.headers = {
             'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36'
         }
         self.temp_dir = Path(tempfile.mkdtemp())
+        
+        # Initialize keyword extractor model
+        try:
+            self.keyword_model = pipeline(
+                "text-classification",
+                model="facebook/bart-large-mnli",
+                device=0 if torch.cuda.is_available() else -1
+            )
+        except Exception as e:
+            print(f"Failed to load keyword model: {e}")
+            self.keyword_model = None
+
+    def extract_keywords(self, text: str) -> List[Dict[str, str]]:
+        """Extract relevant keywords and categories from text using AI"""
+        keywords = []
+        
+        try:
+            # Define candidate labels for classification
+            candidate_labels = [
+                "technology", "science", "education", "business",
+                "health", "nature", "people", "urban", "abstract",
+                "sports", "food", "travel", "architecture", "art",
+                "music", "fashion", "medical", "industrial", "space",
+                "environmental", "historical", "cultural", "professional"
+            ]
+            
+            # Use model to classify text against each label
+            if self.keyword_model:
+                results = self.keyword_model(text, candidate_labels, multi_label=True)
+                
+                # Filter results with high confidence
+                for score, label in zip(results['scores'], results['labels']):
+                    if score > 0.3:  # Confidence threshold
+                        keywords.append({
+                            'keyword': label,
+                            'confidence': score,
+                            'category': self.categorize_keyword(label)
+                        })
+            
+            # Extract additional keywords using NLP
+            additional_keywords = self.extract_noun_phrases(text)
+            for keyword in additional_keywords:
+                keywords.append({
+                    'keyword': keyword,
+                    'confidence': 0.5,
+                    'category': 'content_specific'
+                })
+            
+            # Sort by confidence
+            keywords = sorted(keywords, key=lambda x: x['confidence'], reverse=True)
+            
+            return keywords
+            
+        except Exception as e:
+            print(f"Keyword extraction error: {e}")
+            return self.get_fallback_keywords()
+
+    def extract_noun_phrases(self, text: str) -> List[str]:
+        """Extract important noun phrases from text"""
+        words = text.lower().split()
+        phrases = []
+        
+        # Common adjectives that might indicate important concepts
+        adjectives = {'digital', 'smart', 'modern', 'advanced', 'innovative', 
+                     'technical', 'professional', 'creative', 'strategic'}
+                     
+        for i in range(len(words)-1):
+            if words[i] in adjectives:
+                phrases.append(f"{words[i]} {words[i+1]}")
+                
+        return list(set(phrases))
+
+    def categorize_keyword(self, keyword: str) -> str:
+        """Categorize keyword into general themes"""
+        categories = {
+            'technical': {'technology', 'digital', 'software', 'computer', 'cyber'},
+            'scientific': {'science', 'research', 'laboratory', 'experiment'},
+            'business': {'business', 'professional', 'corporate', 'office'},
+            'educational': {'education', 'learning', 'teaching', 'academic'},
+            'creative': {'art', 'design', 'creative', 'innovation'},
+        }
+        
+        for category, terms in categories.items():
+            if any(term in keyword.lower() for term in terms):
+                return category
+        return 'general'
+
+    def get_images_for_keyword(self, keyword: str) -> List[Dict[str, str]]:
+        """Get images for a specific keyword with metadata"""
+        try:
+            base_url = "https://pixabay.com/api/"
+            params = {
+                'key': self.PIXABAY_API_KEY,
+                'q': keyword,
+                'image_type': 'photo',
+                'per_page': 5,
+                'safesearch': True,
+                'lang': 'en'
+            }
+            
+            response = requests.get(base_url, params=params, headers=self.headers)
+            
+            if response.status_code == 200:
+                data = response.json()
+                if 'hits' in data and data['hits']:
+                    return [{
+                        'url': img['largeImageURL'],
+                        'keyword': keyword,
+                        'relevance': 'Primary match' if keyword in img['tags'] else 'Related',
+                        'tags': img['tags']
+                    } for img in data['hits']]
+            return []
+        except Exception as e:
+            print(f"Error fetching images for keyword {keyword}: {e}")
+            return []
 
     def get_pixabay_images(self, query: str) -> List[str]:
         """Get images from Pixabay API with enhanced error handling"""
@@ -85,29 +200,57 @@ class ImageScraper:
             "https://images.pexels.com/photos/5473950/pexels-photo-5473950.jpeg"
         ]
 
-    def get_images(self, query: str, num_images: int = 15) -> List[str]:
-        """Main method to get images with fallback options"""
-        # First try Pixabay
-        images = self.get_pixabay_images(query)
-        
-        # If no Pixabay images, try with technology-related terms
-        if not images:
-            print("No images found for original query, trying tech terms...")
-            tech_terms = ["digital security", "technology", "cyber security", "data protection"]
-            for term in tech_terms:
-                images.extend(self.get_pixabay_images(term))
-                if images:
-                    break
-        
-        # If still no images, use stock images
-        if not images:
-            print("Using fallback stock images...")
-            images = self.get_stock_images()
-        
-        # Remove duplicates and limit to num_images
-        unique_images = list(dict.fromkeys(images))[:num_images]
-        print(f"Returning {len(unique_images)} images")
-        return unique_images
+     def get_images(self, prompt: str, num_images: int = 15) -> Dict[str, List[Dict[str, str]]]:
+        """Get images with AI-powered keyword extraction and categorization"""
+        try:
+            # Extract keywords from prompt
+            keywords = self.extract_keywords(prompt)
+            
+            # Dictionary to store categorized images
+            categorized_images = {
+                'primary': [],      # Most relevant images
+                'secondary': [],    # Related images
+                'general': []       # Generic/fallback images
+            }
+            
+            # Get images for each keyword
+            for kw_data in keywords:
+                keyword = kw_data['keyword']
+                images = self.get_images_for_keyword(keyword)
+                
+                # Categorize images based on confidence
+                for img in images:
+                    if kw_data['confidence'] > 0.7:
+                        categorized_images['primary'].extend(images)
+                    elif kw_data['confidence'] > 0.4:
+                        categorized_images['secondary'].extend(images)
+                    else:
+                        categorized_images['general'].extend(images)
+            
+            # If no images found, use fallback
+            if not any(categorized_images.values()):
+                categorized_images['general'] = [{
+                    'url': url,
+                    'keyword': 'generic',
+                    'relevance': 'Fallback',
+                    'tags': ''
+                } for url in self.get_stock_images()]
+            
+            return categorized_images
+            
+        except Exception as e:
+            print(f"Error in get_images: {e}")
+            return {'general': [{'url': url, 'keyword': 'fallback', 'relevance': 'Fallback', 'tags': ''} 
+                              for url in self.get_stock_images()]}
+
+    def get_fallback_keywords(self) -> List[Dict[str, str]]:
+        """Return fallback keywords if AI extraction fails"""
+        return [
+            {'keyword': 'technology', 'confidence': 1.0, 'category': 'technical'},
+            {'keyword': 'business', 'confidence': 0.8, 'category': 'business'},
+            {'keyword': 'professional', 'confidence': 0.8, 'category': 'business'},
+            {'keyword': 'digital', 'confidence': 0.7, 'category': 'technical'}
+        ]
 
     def verify_image_url(self, url: str) -> bool:
         """Verify if an image URL is accessible"""
@@ -741,102 +884,140 @@ class VideoGeneratorUI:
                 max-width: 1200px;
                 margin: 0 auto;
             }
-            .image-grid {
-                display: grid;
-                grid-template-columns: repeat(auto-fill, minmax(250px, 1fr));
-                gap: 1rem;
+            .image-category {
+                margin-top: 2rem;
+                padding: 1rem;
+                border-radius: 0.5rem;
+                background: #f8f9fa;
             }
-            .stButton>button {
-                width: 100%;
+            .image-metadata {
+                font-size: 0.8rem;
+                color: #666;
+                margin-top: 0.5rem;
             }
             </style>
         """, unsafe_allow_html=True)
 
-        # Header
         st.title("VaultGenix Video Generator")
         st.markdown("Create professional videos for your digital legacy management platform")
         
-        # Input Section
         with st.container():
             prompt = st.text_area("Enter your video script", height=200)
             
             if prompt:
-                with st.spinner("Fetching images..."):
+                with st.spinner("Analyzing prompt and fetching relevant images..."):
                     try:
-                        st.subheader("Select Images for Your Video")
-                        images = self.generator.image_scraper.get_images(prompt)
-                        
-                        if not images:
-                            st.warning("No images found. Loading default technology images...")
-                            images = self.generator.image_scraper.get_stock_images()
-                        
-                        # Image Selection Grid
-                        selected_images = []
-                        cols = st.columns(3)
-                        
-                        for idx, img_url in enumerate(images):
-                            with cols[idx % 3]:
-                                try:
-                                    # Verify image URL before displaying
-                                    if self.generator.image_scraper.verify_image_url(img_url):
-                                        st.image(img_url, use_container_width=True)
-                                        if st.checkbox(f"Select Image {idx + 1}", key=f"img_{idx}"):
-                                            selected_images.append(img_url)
-                                except Exception as e:
-                                    print(f"Error displaying image {idx}: {str(e)}")
-                                    continue
+                        # Get categorized images
+                        image_categories = self.generator.image_scraper.get_images(prompt)
                         
-                        # Only show video settings if images are available
-                        if selected_images:
-                            st.subheader("Video Settings")
-                            col1, col2 = st.columns(2)
-                            with col1:
-                                style = st.selectbox(
-                                    "Choose style",
-                                    options=["Professional", "Creative", "Educational"],
-                                    index=0
-                                )
-                            with col2:
-                                duration = st.slider(
-                                    "Video duration (seconds)",
-                                    min_value=30,
-                                    max_value=180,
-                                    value=60,
-                                    step=30
-                                )
-
-                            if st.button("Generate Video", type="primary"):
-                                with st.spinner("Generating your video..."):
-                                    try:
-                                        output_path = f"vaultgenix_video_{int(time.time())}.mp4"
-                                        video_path = self.generator.create_video(
-                                            prompt,
-                                            style,
-                                            duration,
-                                            output_path,
-                                            selected_images
-                                        )
-                                        
-                                        if os.path.exists(video_path):
-                                            st.success("✨ Video generated successfully!")
-                                            st.video(video_path)
-                                            
-                                            with open(video_path, 'rb') as video_file:
-                                                st.download_button(
-                                                    "⬇️ Download Video",
-                                                    video_file.read(),
-                                                    file_name=output_path,
-                                                    mime="video/mp4"
-                                                )
-                                    except Exception as e:
-                                        st.error(f"Failed to generate video: {str(e)}")
-                                        print(f"Video generation error: {str(e)}")
+                        if any(image_categories.values()):
+                            # Display primary matches
+                            if image_categories['primary']:
+                                st.subheader("Most Relevant Images")
+                                self.display_image_grid(image_categories['primary'])
+                            
+                            # Display secondary matches
+                            if image_categories['secondary']:
+                                st.subheader("Related Images")
+                                self.display_image_grid(image_categories['secondary'])
+                            
+                            # Display general/fallback images
+                            if image_categories['general']:
+                                st.subheader("Additional Suggested Images")
+                                self.display_image_grid(image_categories['general'])
+                            
+                            # Collect selected images
+                            selected_images = []
+                            for category in image_categories.values():
+                                for img in category:
+                                    key = f"img_{img['url']}"
+                                    if st.session_state.get(key, False):
+                                        selected_images.append(img['url'])
+                            
+                            # Video generation section
+                            if selected_images:
+                                self.show_video_settings(prompt, selected_images)
                         else:
-                            st.error("No valid images found. Please try a different search term.")
-                    
+                            st.warning("No images found. Please try a different prompt.")
+                            
                     except Exception as e:
                         st.error(f"An error occurred: {str(e)}")
-                        print(f"Setup error: {str(e)}")
+                        print(f"Error in UI: {str(e)}")
+
+    def display_image_grid(self, images: List[Dict[str, str]], cols: int = 3):
+        """Display images in a grid with metadata"""
+        for i in range(0, len(images), cols):
+            cols = st.columns(cols)
+            for j, col in enumerate(cols):
+                if i + j < len(images):
+                    img = images[i + j]
+                    with col:
+                        try:
+                            st.image(img['url'], use_container_width=True)
+                            st.checkbox(
+                                "Select",
+                                key=f"img_{img['url']}",
+                                help=f"Keywords: {img['keyword']}\nTags: {img['tags']}"
+                            )
+                            st.markdown(
+                                f"<div class='image-metadata'>"
+                                f"Relevance: {img['relevance']}<br>"
+                                f"Keywords: {img['keyword']}"
+                                f"</div>",
+                                unsafe_allow_html=True
+                            )
+                        except Exception as e:
+                            print(f"Error displaying image: {e}")
+
+    def show_video_settings(self, prompt: str, selected_images: List[str]):
+        """Show video generation settings and controls"""
+        st.subheader("Video Settings")
+        col1, col2 = st.columns(2)
+        with col1:
+            style = st.selectbox(
+                "Choose style",
+                options=["Professional", "Creative", "Educational"],
+                index=0
+            )
+        with col2:
+            duration = st.slider(
+                "Video duration (seconds)",
+                min_value=30,
+                max_value=180,
+                value=60,
+                step=30
+            )
+
+        if st.button("Generate Video", type="primary"):
+            self.generate_video(prompt, style, duration, selected_images)
+
+    def generate_video(self, prompt: str, style: str, duration: int, selected_images: List[str]):
+        """Handle video generation"""
+        with st.spinner("Generating your video..."):
+            try:
+                output_path = f"vaultgenix_video_{int(time.time())}.mp4"
+                video_path = self.generator.create_video(
+                    prompt,
+                    style,
+                    duration,
+                    output_path,
+                    selected_images
+                )
+                
+                if os.path.exists(video_path):
+                    st.success("✨ Video generated successfully!")
+                    st.video(video_path)
+                    
+                    with open(video_path, 'rb') as video_file:
+                        st.download_button(
+                            "⬇️ Download Video",
+                            video_file.read(),
+                            file_name=output_path,
+                            mime="video/mp4"
+                        )
+            except Exception as e:
+                st.error(f"Failed to generate video: {str(e)}")
+                print(f"Video generation error: {str(e)}")
 
 if __name__ == "__main__":
     ui = VideoGeneratorUI()