Update app.py

app.py

@@ -290,53 +290,43 @@ class ImageScraper:
         ]
 
     def get_images(self, query: str, num_images: int = 15) -> Dict[str, List[Dict[str, str]]]:
-        """Get images with AI-driven selection and ranking"""
+        """Get images with enhanced AI-driven selection and ranking"""
         try:
-            # Initialize result structure
+            # Extract key topics and their importance
+            topics = self.extract_key_topics(query)
+            topic_scores = {topic: score for score, topic in 
+                           zip(np.linspace(1.0, 0.6, len(topics)), topics)}
+            
+            # Initialize categories
             result = {
                 'primary': [],
                 'secondary': [],
                 'general': []
             }
             
-            # 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)
+            # Fetch and analyze images for each topic
+            for topic, base_score in topic_scores.items():
+                images = self.get_images_for_keyword(topic)
                 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()):
-                stock_images = self.get_stock_images()
-                result['general'] = [{
-                    'url': url,
-                    'keyword': 'technology',
-                    'relevance': 'Fallback',
-                    'tags': 'technology',
-                    'relevance_score': 0.5
-                } for url in stock_images[:num_images]]
+                    # Enhanced relevance scoring
+                    relevance_score = self.calculate_relevance_score(img, topic, base_score, query)
+                    img['relevance_score'] = relevance_score
+                    
+                    # Categorize based on relevance score
+                    if relevance_score > 0.8:
+                        result['primary'].append(img)
+                    elif relevance_score > 0.6:
+                        result['secondary'].append(img)
+                    else:
+                        result['general'].append(img)
+            
+            # Sort each category by relevance score
+            for category in result:
+                result[category] = sorted(
+                    result[category],
+                    key=lambda x: x['relevance_score'],
+                    reverse=True
+                )[:num_images // 3]  # Limit images per category
             
             return result
             
@@ -344,6 +334,37 @@ class ImageScraper:
             print(f"Error in get_images: {str(e)}")
             return self.get_fallback_images(num_images)
     
+    def calculate_relevance_score(self, image: Dict[str, str], topic: str, base_score: float, query: str) -> float:
+        """Calculate enhanced relevance score for an image"""
+        score = base_score
+        
+        # Analyze image tags
+        tags = set(image['tags'].lower().split(','))
+        query_words = set(query.lower().split())
+        
+        # Direct matches with query
+        query_matches = len(tags.intersection(query_words))
+        score += query_matches * 0.1
+        
+        # Topic relevance
+        if topic.lower() in tags:
+            score += 0.2
+        
+        # Context relevance
+        relevant_terms = {
+            'digital': 0.15,
+            'security': 0.15,
+            'technology': 0.1,
+            'professional': 0.1,
+            'modern': 0.05
+        }
+        
+        for term, weight in relevant_terms.items():
+            if term in tags:
+                score += weight
+        
+        return min(score, 1.0)  # Normalize to 0-1
+        
     def score_keywords(self, query: str, keywords: List[str]) -> Dict[str, float]:
         """Score keywords based on relevance to query"""
         scores = {}