Update app.py

app.py

@@ -2,6 +2,8 @@ import gradio as gr
 import os
 import docx
 import numpy as np
+import cv2
+from pathlib import Path
 from sentence_transformers import SentenceTransformer
 from sklearn.metrics.pairwise import cosine_similarity
 from langchain.text_splitter import RecursiveCharacterTextSplitter
@@ -11,6 +13,66 @@ from langchain.memory import ConversationBufferMemory
 from langchain_community.llms import HuggingFaceEndpoint
 from langchain_huggingface import HuggingFaceEmbeddings
 
+class SimpleGeometryClassifier:
+    def __init__(self):
+        self.reference_embeddings = {
+            "flat.png": {
+                "embedding": None,
+                "label": "Flat or Sheet-Based"
+            },
+            "cylindrical.png": {
+                "embedding": None,
+                "label": "Cylindrical"
+            },
+            "complex.png": {
+                "embedding": None,
+                "label": "Complex Multi Axis Geometry"
+            }
+        }
+
+    def compute_embedding(self, image_path):
+        img = cv2.imread(image_path)
+        img = cv2.resize(img, (224, 224))
+        img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
+        
+        win_size = (224, 224)
+        cell_size = (8, 8)
+        block_size = (16, 16)
+        block_stride = (8, 8)
+        num_bins = 9
+        
+        hog = cv2.HOGDescriptor(win_size, block_size, block_stride, cell_size, num_bins)
+        embedding = hog.compute(img)
+        
+        return embedding.flatten()
+
+    def initialize_reference_embeddings(self, reference_folder):
+        for image_name in self.reference_embeddings.keys():
+            image_path = str(Path(reference_folder) / image_name)
+            if Path(image_path).exists():
+                self.reference_embeddings[image_name]["embedding"] = self.compute_embedding(image_path)
+            else:
+                print(f"Warning: Reference image {image_path} not found")
+
+    def find_closest_geometry(self, query_embedding):
+        best_similarity = -1
+        best_label = None
+        
+        for ref_data in self.reference_embeddings.values():
+            if ref_data["embedding"] is not None:
+                similarity = np.dot(query_embedding, ref_data["embedding"]) / (
+                    np.linalg.norm(query_embedding) * np.linalg.norm(ref_data["embedding"])
+                )
+                if similarity > best_similarity:
+                    best_similarity = similarity
+                    best_label = ref_data["label"]
+        
+        return best_label or "Unknown Geometry"
+
+    def process_image(self, image_path):
+        query_embedding = self.compute_embedding(image_path)
+        return self.find_closest_geometry(query_embedding)
+
 # Initialize semantic model
 semantic_model = SentenceTransformer("all-MiniLM-L6-v2")
 
@@ -138,20 +200,51 @@ def handle_query(query, history, retriever, qa_chain, embeddings):
     history.append((query, assistant_response))
     return history, ""
 
+def process_image_and_generate_query(image):
+    classifier = SimpleGeometryClassifier()
+    classifier.initialize_reference_embeddings("images")
+    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():
+    # Initialize classifier
+    classifier = SimpleGeometryClassifier()
+    classifier.initialize_reference_embeddings("images")
+    
+    # Initialize chatbot components
     documents = load_documents()
     vector_db, embeddings = create_db(documents)
     retriever, qa_chain = initialize_chatbot(vector_db, embeddings)
     
     with gr.Blocks() as app:
-        gr.Markdown("### 🤖 Document Question Answering System")
-        
-        chatbot = gr.Chatbot()
-        query_input = gr.Textbox(label="Ask a question about the documents")
-        query_btn = gr.Button("Submit")
+        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 a question about the documents")
+                query_btn = gr.Button("Submit")
+
+        def image_upload_handler(image):
+            if image is None:
+                return "", ""
+            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)
+
+        image_input.change(
+            image_upload_handler,
+            inputs=[image_input],
+            outputs=[geometry_label, query_input]
+        )
         
         query_btn.click(
             user_query_handler,
@@ -164,7 +257,7 @@ def demo():
             inputs=[query_input, chatbot],
             outputs=[chatbot, query_input]
         )
-    
+
     app.launch()
 
 if __name__ == "__main__":