Update app.py

app.py

@@ -10,7 +10,80 @@ from langchain.chains import ConversationalRetrievalChain
 from langchain.memory import ConversationBufferMemory
 from langchain_community.llms import HuggingFaceEndpoint
 from langchain_huggingface import HuggingFaceEmbeddings
-
+import torch
+from PIL import Image
+from torchvision import transforms
+from torchvision.models import resnet50, ResNet50_Weights
+from torchvision import transforms, models
+
+
+class GeometryImageClassifier:
+    def __init__(self):
+        # Load ResNet50 but only use it for feature extraction
+        self.model = models.resnet50(weights='DEFAULT')
+        # Remove the final classification layer
+        self.model.fc = torch.nn.Identity()
+        self.model.eval()
+        
+        self.transform = transforms.Compose([
+            transforms.Resize((224, 224)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        ])
+        
+        # Pre-computed embeddings for our 3 reference images with manual labels
+        self.reference_embeddings = {
+            "flat.png": {
+                "embedding": None,  # Will be computed on first run
+                "label": "Flat or Sheet-Based"
+            },
+            "cylindrical.png": {
+                "embedding": None,
+                "label": "Cylindrical"
+            },
+            "complex.png": {
+                "embedding": None,
+                "label": "Complex Multi Axis Geometry"
+            }
+        }
+        
+    def compute_embedding(self, images):
+        img = Image.open(images).convert('RGB')
+        img_tensor = self.transform(img).unsqueeze(0)
+        
+        with torch.no_grad():
+            embedding = self.model(img_tensor)
+        return embedding.squeeze().numpy()
+    
+    def initialize_reference_embeddings(self, reference_folder):
+        for image_name in self.reference_embeddings.keys():
+            images = f"{reference_folder}/{image_name}"
+            self.reference_embeddings[image_name]["embedding"] = self.compute_embedding(images)
+    
+    def find_closest_geometry(self, query_embedding):
+        best_similarity = -1
+        best_label = None
+        
+        for ref_data in self.reference_embeddings.values():
+            similarity = cosine_similarity(
+                query_embedding.reshape(1, -1),
+                ref_data["embedding"].reshape(1, -1)
+            )[0][0]
+            
+            if similarity > best_similarity:
+                best_similarity = similarity
+                best_label = ref_data["label"]
+        
+        return best_label
+    
+    def process_image(self, images):
+        # Compute embedding for the input image
+        query_embedding = self.compute_embedding(images)
+        
+        # Find the closest matching reference geometry
+        return self.find_closest_geometry(query_embedding)
+
+    
 # ✅ Use a strong sentence embedding model
 semantic_model = SentenceTransformer("all-MiniLM-L6-v2")
 
@@ -103,7 +176,7 @@ def validate_query_semantically(query, retrieved_docs):
 
     print(f"🔍 Semantic Similarity Score: {similarity_score}")
 
-    return similarity_score >= 0.4  # 🔥 Stricter threshold to ensure correctness
+    return similarity_score >= 0.3  # 🔥 Stricter threshold to ensure correctness
 
 
 def handle_query(query, history, retriever, qa_chain, embeddings):
@@ -163,24 +236,64 @@ def initialize_chatbot(vector_db):
     return retriever, qa_chain, embeddings
 
 
+def process_image_and_generate_query(image):
+    classifier = GeometryImageClassifier()
+    geometry_type = classifier.process_image(image)
+    
+    query = f"I have a {geometry_type} geometry, which screw should I use and what is the best machine to use for {geometry_type} geometry?"
+    return geometry_type, query
+
 def demo():
-    """ ✅ Starts the chatbot application using Gradio. """
+    # Initialize classifier once at startup
+    classifier = GeometryImageClassifier()
+    classifier.initialize_reference_embeddings("images")
+    
+    # Initialize chatbot components
     retriever, qa_chain, embeddings = initialize_chatbot(create_db(load_documents()))
-
+    
     with gr.Blocks() as app:
-        gr.Markdown("### 🤖 **Fastener Agent** 📚")
-        chatbot = gr.Chatbot()
-        query_input = gr.Textbox(label="Ask me a question")
-        query_btn = gr.Button("Ask")
+        gr.Markdown("### 🤖 **Fastener Agent with Image Recognition** 📚")
+        
+        with gr.Row():
+            with gr.Column(scale=1):
+                image_input = gr.Image(type="filepath", label="Upload Geometry Image")
+                geometry_label = gr.Textbox(label="Detected Geometry Type", interactive=False)
+                
+            with gr.Column(scale=2):
+                chatbot = gr.Chatbot()
+                query_input = gr.Textbox(label="Ask me a question")
+                query_btn = gr.Button("Ask")
+
+        def image_upload_handler(image):
+            if image is None:
+                return "", ""
+            # Use the initialized classifier
+            geometry_type = classifier.process_image(image)
+            suggested_query = f"I have a {geometry_type} geometry, which screw should I use and what is the best machine to use for {geometry_type} geometry?"
+            return geometry_type, suggested_query
 
         def user_query_handler(query, history):
             return handle_query(query, history, retriever, qa_chain, embeddings)
 
-        query_btn.click(user_query_handler, inputs=[query_input, chatbot], outputs=[chatbot, query_input])
-        query_input.submit(user_query_handler, inputs=[query_input, chatbot], outputs=[chatbot, query_input])
+        image_input.change(
+            image_upload_handler,
+            inputs=[image_input],
+            outputs=[geometry_label, query_input]
+        )
+        
+        query_btn.click(
+            user_query_handler,
+            inputs=[query_input, chatbot],
+            outputs=[chatbot, query_input]
+        )
+        
+        query_input.submit(
+            user_query_handler,
+            inputs=[query_input, chatbot],
+            outputs=[chatbot, query_input]
+        )
 
     app.launch()
 
-
 if __name__ == "__main__":
     demo()