Spaces:
Sleeping
Sleeping
| from PIL import Image | |
| from pathlib import Path | |
| from typing import Dict, List, Tuple | |
| import os | |
| from modules.image_captioner import ImageCaptioner | |
| from modules.knowledge_graph import KnowledgeGraph | |
| class ImageSearchSystem: | |
| """Integration class for image captioning and knowledge graph""" | |
| def __init__(self): | |
| self.captioner = ImageCaptioner() | |
| self.knowledge_graph = KnowledgeGraph() | |
| def process_image(self, image_path_or_img) -> Dict[str, str]: | |
| """Process a single image and add to knowledge graph""" | |
| try: | |
| # Handle both string paths and PIL Image objects | |
| if isinstance(image_path_or_img, str): | |
| image = Image.open(image_path_or_img) | |
| else: | |
| image = image_path_or_img # Already a PIL Image | |
| # Generate caption | |
| caption = self.captioner.generate_caption(image) | |
| # Add to knowledge graph | |
| image_id = str(hash(str(image))) # Generate unique ID for the image | |
| self.knowledge_graph.add_image_node(image_id, caption) | |
| return {"image_id": image_id, "caption": caption} | |
| except Exception as e: | |
| print(f"Error processing image: {e}") | |
| return {"error": str(e)} | |
| def process_directory(self, dir_path: str) -> List[Dict[str, str]]: | |
| """Process all images in a directory""" | |
| results = [] | |
| image_extensions = {'.jpg', '.jpeg', '.png'} | |
| for path in Path(dir_path).iterdir(): | |
| if path.suffix.lower() in image_extensions: | |
| result = self.process_image(str(path)) | |
| results.append(result) | |
| # Expand concepts from caption | |
| words = result['caption'].split() | |
| for word in words: | |
| relationships = self.knowledge_graph.expand_concept(word) | |
| self.knowledge_graph.add_weighted_relationships(word, relationships) | |
| return results | |
| def search_images(self, query: str, max_depth: int = 3, top_k: int = 5) -> List[Tuple[str, float]]: | |
| """ | |
| Search for images using semantic understanding | |
| Args: | |
| query: Search query string | |
| max_depth: Maximum path length to consider | |
| top_k: Number of top results to return | |
| """ | |
| return self.knowledge_graph.search_with_depth(query, max_depth=max_depth, limit=top_k) | |
| # def search_images(self, query: str, max_depth: int = 3) -> List[Tuple[str, float]]: | |
| # """Search for images using semantic understanding""" | |
| # return self.knowledge_graph.search_with_depth(query, max_depth=max_depth) | |