Update app.py

app.py

@@ -289,7 +289,7 @@ class ImageScraper:
         ]
 
     def get_images(self, query: str, num_images: int = 15) -> Dict[str, List[Dict[str, str]]]:
-        """Get images for either single word queries or extract keywords from long prompts"""
+        """Get images with AI-driven selection and ranking"""
         try:
             # Initialize result structure
             result = {
@@ -298,46 +298,33 @@ class ImageScraper:
                 'general': []
             }
             
-            # 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]
-            
-            # 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'
-                }
-                
-                response = requests.get(base_url, params=params, headers=self.headers)
-                
-                if response.status_code == 200:
-                    data = response.json()
-                    hits = data.get('hits', [])
-                    
-                    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)
+            # Extract and analyze keywords using AI
+            keywords = self.extract_key_topics(query)
+            print(f"AI extracted keywords: {keywords}")
+            
+            # Score and rank keywords based on relevance to query
+            keyword_scores = self.score_keywords(query, keywords)
+            ranked_keywords = sorted(keyword_scores.items(), key=lambda x: x[1], reverse=True)
+            
+            # Fetch and analyze images for each keyword
+            all_images = []
+            for keyword, score in ranked_keywords:
+                images = self.get_images_for_keyword(keyword)
+                for img in images:
+                    img['relevance_score'] = score * self.analyze_image_relevance(img, query)
+                    all_images.append(img)
+            
+            # Sort images by relevance score
+            sorted_images = sorted(all_images, key=lambda x: x['relevance_score'], reverse=True)
+            
+            # Distribute images across categories
+            total_images = min(len(sorted_images), num_images)
+            primary_count = total_images // 2
+            secondary_count = total_images // 3
+            
+            result['primary'] = sorted_images[:primary_count]
+            result['secondary'] = sorted_images[primary_count:primary_count + secondary_count]
+            result['general'] = sorted_images[primary_count + secondary_count:total_images]
             
             # If no images found, use stock images
             if not any(result.values()):
@@ -346,23 +333,75 @@ class ImageScraper:
                     'url': url,
                     'keyword': 'technology',
                     'relevance': 'Fallback',
-                    'tags': 'technology'
+                    'tags': 'technology',
+                    'relevance_score': 0.5
                 } 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[:num_images]]
+            return self.get_fallback_images(num_images)
+    
+    def score_keywords(self, query: str, keywords: List[str]) -> Dict[str, float]:
+        """Score keywords based on relevance to query"""
+        scores = {}
+        query_words = set(query.lower().split())
+        
+        for keyword in keywords:
+            score = 0.0
+            keyword_words = set(keyword.lower().split())
+            
+            # Direct word match
+            word_matches = len(keyword_words.intersection(query_words))
+            score += word_matches * 0.3
+            
+            # Contextual relevance
+            context_terms = {
+                'digital': 0.8,
+                'security': 0.7,
+                'legacy': 0.9,
+                'protection': 0.6,
+                'management': 0.5,
+                'AI': 0.8,
+                'technology': 0.6
             }
+            
+            for term, weight in context_terms.items():
+                if term in keyword.lower():
+                    score += weight
+            
+            scores[keyword] = min(score, 1.0)  # Normalize to 0-1
+        
+        return scores
+    
+    def analyze_image_relevance(self, image: Dict[str, str], query: str) -> float:
+        """Analyze image relevance based on tags and metadata"""
+        score = 0.0
+        
+        # Analyze tags
+        tags = set(image['tags'].lower().split(','))
+        query_words = set(query.lower().split())
+        
+        # Tag matching
+        matching_tags = len(tags.intersection(query_words))
+        score += matching_tags * 0.2
+        
+        # Context relevance
+        relevant_terms = {
+            'technology': 0.3,
+            'digital': 0.3,
+            'security': 0.3,
+            'business': 0.2,
+            'professional': 0.2,
+            'modern': 0.1
+        }
+        
+        for term, weight in relevant_terms.items():
+            if term in tags:
+                score += weight
+        
+        return min(score, 1.0)  # Normalize to 0-1
                 
                 
     def get_fallback_keywords(self) -> List[Dict[str, str]]: