Update app.py

app.py

@@ -201,110 +201,82 @@ class ImageScraper:
         ]
 
     def get_images(self, query: str, num_images: int = 15) -> Dict[str, List[Dict[str, str]]]:
-        """
-        Get images from Pixabay API with enhanced error handling and categorization
-        Returns a dictionary with categorized image results
-        """
+        """Get images for either single word queries or extract keywords from long prompts"""
         try:
-            # Ensure query is not empty
-            if not query.strip():
-                query = "digital security"
-    
-            base_url = "https://pixabay.com/api/"
-            params = {
-                'key': self.PIXABAY_API_KEY,
-                'q': query,
-                'image_type': 'photo',
-                'per_page': num_images,
-                'safesearch': True,
-                'lang': 'en'
+            # Initialize result structure
+            result = {
+                'primary': [],
+                'secondary': [],
+                'general': []
             }
             
-            print(f"Querying Pixabay API with params: {params}")
-            response = requests.get(base_url, params=params, headers=self.headers)
-            print(f"API Response status: {response.status_code}")
+            # Extract keywords if query is long
+            if len(query.split()) > 3:
+                keywords = self.extract_key_topics(query)
+                print(f"Extracted keywords: {keywords}")  # Debug log
+            else:
+                keywords = [query]
             
-            if response.status_code == 200:
-                data = response.json()
-                hits = data.get('hits', [])
-                print(f"Found {len(hits)} images")
-                
-                if not hits:
-                    # If no results, try backup keywords
-                    backup_keywords = ["digital security", "technology", "cyber security", "business technology"]
-                    for keyword in backup_keywords:
-                        params['q'] = keyword
-                        backup_response = requests.get(base_url, params=params, headers=self.headers)
-                        if backup_response.status_code == 200:
-                            backup_data = backup_response.json()
-                            hits.extend(backup_data.get('hits', []))
-                            if hits:
-                                break
-                
-                # Structure response with categories
-                result = {
-                    'primary': [],
-                    'secondary': [],
-                    'general': []
+            # Fetch images for each keyword
+            for keyword in keywords:
+                base_url = "https://pixabay.com/api/"
+                params = {
+                    'key': self.PIXABAY_API_KEY,
+                    'q': keyword,
+                    'image_type': 'photo',
+                    'per_page': max(3, num_images // len(keywords)),  # Distribute images among keywords
+                    'safesearch': True,
+                    'lang': 'en'
                 }
                 
-                # If no hits found even after backup keywords, return stock images
-                if not hits:
-                    stock_images = self.get_stock_images()
-                    result['general'] = [
-                        {
-                            'url': url,
-                            'keyword': 'technology',
-                            'relevance': 'Fallback',
-                            'tags': 'technology'
-                        } for url in stock_images
-                    ]
-                    return result
+                response = requests.get(base_url, params=params, headers=self.headers)
                 
-                # Process and categorize the found images
-                for hit in hits[:num_images]:
-                    image_data = {
-                        'url': hit['largeImageURL'],
-                        'keyword': query,
-                        'relevance': 'Primary match' if query.lower() in hit.get('tags', '').lower() else 'Related',
-                        'tags': hit.get('tags', '')
-                    }
+                if response.status_code == 200:
+                    data = response.json()
+                    hits = data.get('hits', [])
                     
-                    # Distribute images across categories
-                    if len(result['primary']) < num_images // 3 and query.lower() in hit.get('tags', '').lower():
-                        result['primary'].append(image_data)
-                    elif len(result['secondary']) < num_images // 3:
-                        result['secondary'].append(image_data)
-                    else:
-                        result['general'].append(image_data)
-    
-                # Ensure we have some results in each category
-                if not result['primary'] and result['secondary']:
-                    result['primary'] = result['secondary'][:num_images//3]
-                    result['secondary'] = result['secondary'][num_images//3:]
-                
-                if not result['secondary'] and result['general']:
-                    result['secondary'] = result['general'][:num_images//3]
-                    result['general'] = result['general'][num_images//3:]
-    
-                return result
-    
+                    for hit in hits:
+                        image_data = {
+                            'url': hit['largeImageURL'],
+                            'keyword': keyword,
+                            'relevance': 'Primary match',
+                            'tags': hit.get('tags', '')
+                        }
+                        
+                        # Distribute images across categories
+                        if len(result['primary']) < num_images // 3:
+                            result['primary'].append(image_data)
+                        elif len(result['secondary']) < num_images // 3:
+                            result['secondary'].append(image_data)
+                        else:
+                            result['general'].append(image_data)
+            
+            # If no images found, use stock images
+            if not any(result.values()):
+                stock_images = self.get_stock_images()
+                result['general'] = [{
+                    'url': url,
+                    'keyword': 'technology',
+                    'relevance': 'Fallback',
+                    'tags': 'technology'
+                } for url in stock_images[:num_images]]
+            
+            return result
+            
         except Exception as e:
             print(f"Error in get_images: {str(e)}")
             # Return stock images as fallback
             stock_images = self.get_stock_images()
             return {
-                'general': [
-                    {
-                        'url': url,
-                        'keyword': 'technology',
-                        'relevance': 'Fallback',
-                        'tags': 'technology'
-                    } for url in stock_images
-                ]
+                'general': [{
+                    'url': url,
+                    'keyword': 'technology',
+                    'relevance': 'Fallback',
+                    'tags': 'technology'
+                } for url in stock_images[:num_images]]
             }
-            
-            
+                
+                
     def get_fallback_keywords(self) -> List[Dict[str, str]]:
         """Return fallback keywords if AI extraction fails"""
         return [
@@ -820,25 +792,39 @@ class EnhancedVideoGenerator:
         return text.strip()
 
     def extract_key_topics(self, script: str) -> List[str]:
-        """Extract main topics from the script for visual asset generation"""
+        """Extract key topics from a long text prompt"""
         try:
-            # Simple keyword extraction based on noun phrases
-            # In a production environment, you might want to use a proper NLP library
-            sentences = script.split('.')
-            topics = []
+            # Common tech-related terms to look for
+            tech_terms = {
+                'security', 'digital', 'technology', 'data', 'encryption', 
+                'privacy', 'protection', 'software', 'platform', 'system',
+                'AI', 'artificial intelligence', 'computing', 'cyber'
+            }
+            
+            # Split into words and get unique terms
+            words = script.lower().split()
+            topics = set()
+            
+            # Extract single word topics
+            for word in words:
+                if word in tech_terms:
+                    topics.add(word)
+            
+            # Extract two-word phrases
+            for i in range(len(words) - 1):
+                phrase = f"{words[i]} {words[i+1]}"
+                if any(term in phrase for term in tech_terms):
+                    topics.add(phrase)
             
-            for sentence in sentences:
-                words = sentence.strip().split()
-                if len(words) >= 2:
-                    # Extract potential noun phrases (pairs of words)
-                    topics.append(' '.join(words[:2]))
+            # Add some default topics if none found
+            if not topics:
+                topics = {'digital security', 'data protection', 'technology'}
             
-            # Remove duplicates and limit to top 5 topics
-            return list(dict.fromkeys(topics))[:5]
+            return list(topics)[:5]  # Return top 5 topics
             
         except Exception as e:
-            self.logger.error(f"Topic extraction failed: {str(e)}")
-            return ["default topic"]
+            print(f"Topic extraction error: {e}")
+            return ['digital security']
 
     def generate_ai_image(self, prompt: str, style: str) -> Optional[Image.Image]:
         """Generate an AI image using Stability AI"""
@@ -976,29 +962,40 @@ class VideoGeneratorUI:
     
     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}")
-
+        # Ensure images is a list and not empty
+        if not images or not isinstance(images, list):
+            return
+        
+        # Calculate number of rows needed
+        n_images = len(images)
+        n_rows = (n_images + cols - 1) // cols
+        
+        # Create grid
+        for row in range(n_rows):
+            with st.container():
+                columns = st.columns(cols)
+                for col in range(cols):
+                    idx = row * cols + col
+                    if idx < n_images:
+                        img = images[idx]
+                        with columns[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")