Update modules/search_system.py

modules/search_system.py

@@ -13,17 +13,28 @@ class ImageSearchSystem:
         self.captioner = ImageCaptioner()
         self.knowledge_graph = KnowledgeGraph()
         
-    def process_image(self, image_path: str) -> Dict[str, str]:
+    def process_image(self, image_path_or_img) -> Dict[str, str]:
         """Process a single image and add to knowledge graph"""
-        # Load and caption image
-        image = Image.open(image_path)
-        caption = self.captioner.generate_caption(image)
-        
-        # Add to knowledge graph
-        image_id = os.path.basename(image_path)
-        self.knowledge_graph.add_image_node(image_id, caption)
+        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}
         
-        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"""