Update app.py

app.py

@@ -8,17 +8,34 @@ import uuid
 from openai import OpenAI
 import requests
 from io import BytesIO
-from PIL import Image
+import os
+#from PIL import Image
 
 client = OpenAI(
     base_url="https://a7g1ajqixo23revq.us-east-1.aws.endpoints.huggingface.cloud/v1/",
     api_key="hf_XXXXX"
 )
 
-def format_math(text):
-    text = re.sub(r"\[(.*?)\]", r"$$\1$$", text, flags=re.DOTALL)
-    text = text.replace(r"\(", "$").replace(r"\)", "$")
-    return text
+from transformers import VisionEncoderDecoderModel, TrOCRProcessor
+from PIL import Image
+
+# Load the model and processor
+processor = TrOCRProcessor.from_pretrained("breezedeus/pix2text-mfr")
+model = VisionEncoderDecoderModel.from_pretrained("breezedeus/pix2text-mfr")
+model.eval()
+
+
+def image_to_latex(image_path):
+    image = Image.open(image_path).convert("RGB")
+    pixel_values = processor(images=image, return_tensors="pt").pixel_values
+    with torch.no_grad():
+        generated_ids = model.generate(pixel_values)
+    generated_text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+    return generated_text
+def process_image(image):
+    latex_code = image_to_latex(image)
+    return latex_code
+    
 
 # Global dictionary to store all conversations: {id: {"title": str, "messages": list}}
 conversations = {}
@@ -93,7 +110,7 @@ def generate_response(user_message,
     yield new_history, new_history
 
 # OCR function using HF hosted model
-def extract_latex_from_image(image):
+'''def extract_latex_from_image(image):
     if image is None:
         return gr.update(value="")
 
@@ -112,7 +129,7 @@ def extract_latex_from_image(image):
         return gr.update(value=text)
     else:
         return gr.update(value="⚠️ OCR failed. Please try again.")
-
+'''
 
 
 example_messages = {
@@ -175,16 +192,22 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         with gr.Column(scale=4):
             #chatbot = gr.Chatbot(label="Chat", type="messages")
             chatbot = gr.Chatbot(label="Chat", type="messages", height=520)
+            
             with gr.Row():
                 user_input = gr.Textbox(label="User Input", placeholder="Type your question here...", lines=3, scale=8)
                 with gr.Column():
                     submit_button = gr.Button("Send", variant="primary", scale=1)
                     clear_button = gr.Button("Clear", scale=1)
 
-            gr.Markdown("### Upload Image of a Math Question")
+            '''gr.Markdown("### Upload Image of a Math Question")
             image_input = gr.Image(type="pil", label="Upload Image (PNG/JPG)", height=120)
             ocr_button = gr.Button("Extract LaTeX from Image 🧠")
-            ocr_button.click(fn=extract_latex_from_image, inputs=image_input, outputs=user_input)
+            ocr_button.click(fn=extract_latex_from_image, inputs=image_input, outputs=user_input)'''
+            
+            with gr.Row():
+                image_input = gr.Image(type="filepath", label="Upload Image")
+                output_text = gr.Textbox(label="LaTeX Output")
+            image_input.change(fn=process_image, inputs=image_input, outputs=output_text)
 
                     
             gr.Markdown("**Try these examples:**")