Update app.py

app.py

@@ -4,6 +4,8 @@ import requests
 import gradio as gr
 from huggingface_hub import InferenceClient
 from dataclasses import dataclass
+import pytesseract
+from PIL import Image
 
 @dataclass
 class ChatMessage:
@@ -71,10 +73,8 @@ class XylariaChat:
     def caption_image(self, image):
         """
         Caption an uploaded image using Hugging Face API
-
         Args:
             image (str): Base64 encoded image or file path
-
         Returns:
             str: Image caption or error message
         """
@@ -110,14 +110,35 @@ class XylariaChat:
         except Exception as e:
             return f"Error processing image: {str(e)}"
 
+    def perform_math_ocr(self, image_path):
+        """
+        Perform OCR on an image and return the extracted text.
+
+        Args:
+            image_path (str): Path to the image file.
+
+        Returns:
+            str: Extracted text from the image, or an error message.
+        """
+        try:
+            # Open the image using Pillow library
+            img = Image.open(image_path)
+
+            # Use Tesseract to do OCR on the image
+            text = pytesseract.image_to_string(img)
+
+            # Remove leading/trailing whitespace and return
+            return text.strip()
+
+        except Exception as e:
+            return f"Error during Math OCR: {e}"
+        
     def get_response(self, user_input, image=None):
         """
         Generate a response using chat completions with improved error handling
-
         Args:
             user_input (str): User's message
             image (optional): Uploaded image
-
         Returns:
             Stream of chat completions or error message
         """
@@ -202,18 +223,31 @@ class XylariaChat:
         
     
     def create_interface(self):
-        def streaming_response(message, chat_history, image_filepath):
+        def streaming_response(message, chat_history, image_filepath, math_ocr_image_path):
+            
+            ocr_text = ""
+            if math_ocr_image_path:
+                ocr_text = self.perform_math_ocr(math_ocr_image_path)
+                if ocr_text.startswith("Error"):
+                    # Handle OCR error
+                    updated_history = chat_history + [[message, ocr_text]]
+                    yield "", updated_history, None, None
+                    return
+                else:
+                    message = f"Math OCR Result: {ocr_text}\n\nUser's message: {message}"
+
             # Check if an image was actually uploaded
             if image_filepath:
                 response_stream = self.get_response(message, image_filepath)
             else:
                 response_stream = self.get_response(message)
+                
 
             # Handle errors in get_response
             if isinstance(response_stream, str):
                 # Return immediately with the error message
                 updated_history = chat_history + [[message, response_stream]]
-                yield "", updated_history, None
+                yield "", updated_history, None, None
                 return
 
             # Prepare for streaming response
@@ -229,12 +263,12 @@ class XylariaChat:
                         
                         # Update the last message in chat history with partial response
                         updated_history[-1][1] = full_response
-                        yield "", updated_history, None
+                        yield "", updated_history, None, None
             except Exception as e:
                 print(f"Streaming error: {e}")
                 # Display error in the chat interface
                 updated_history[-1][1] = f"Error during response: {e}"
-                yield "", updated_history, None
+                yield "", updated_history, None, None
                 return
 
             # Update conversation history
@@ -317,6 +351,17 @@ class XylariaChat:
                         )
                         with gr.Row():
                             clear_image_btn = gr.Button("Clear Image")
+                
+                with gr.Accordion("Math Input", open=False):
+                    with gr.Column():
+                        math_ocr_img = gr.Image(
+                            sources=["upload", "webcam"],
+                            type="filepath",
+                            label="Upload Image for math",
+                            elem_classes="image-preview"
+                        )
+                        with gr.Row():
+                            clear_math_ocr_btn = gr.Button("Clear Math Image")
 
                 # Input row with improved layout
                 with gr.Row():
@@ -341,16 +386,24 @@ class XylariaChat:
                     queue=False
                 )
 
+                # Clear Math OCR image functionality
+                clear_math_ocr_btn.click(
+                    fn=lambda: None,
+                    inputs=None,
+                    outputs=[math_ocr_img],
+                    queue=False
+                )
+
                 # Submit functionality with streaming and image support
                 btn.click(
                     fn=streaming_response,
-                    inputs=[txt, chatbot, img],
-                    outputs=[txt, chatbot, img]
+                    inputs=[txt, chatbot, img, math_ocr_img],
+                    outputs=[txt, chatbot, img, math_ocr_img]
                 )
                 txt.submit(
                     fn=streaming_response,
-                    inputs=[txt, chatbot, img],
-                    outputs=[txt, chatbot, img]
+                    inputs=[txt, chatbot, img, math_ocr_img],
+                    outputs=[txt, chatbot, img, math_ocr_img]
                 )
 
                 # Clear conversation history