Update app.py

app.py

@@ -2,12 +2,14 @@ import gradio as gr
 import os
 import requests
 import json
-import pytesseract
-from PIL import Image, ImageFilter, ImageEnhance
+import base64
+from PIL import Image
+from io import BytesIO
 
 GROQ_API_URL = "https://api.groq.com/openai/v1/chat/completions"
 SPACE_GROQ_KEY = os.environ.get("GROQ_API_KEY", "")
-MODEL = "llama-3.3-70b-versatile"
+TEXT_MODEL = "llama-3.3-70b-versatile"
+VISION_MODEL = "llama-3.2-90b-vision-preview"
 
 SYSTEM_PROMPT = """You are a warm, patient math tutor for students from Grade 1 through High School.
 Rules you ALWAYS follow:
@@ -36,22 +38,10 @@ Always format your response exactly like this:
 💡 [One encouraging sentence]"""
 
 
-def ocr_image(pil_image):
-    try:
-        w, h = pil_image.size
-        scale = max(1, 3000 // max(w, h))
-        img = pil_image.resize((w * scale, h * scale), Image.LANCZOS)
-        img = img.convert("L")
-        img = ImageEnhance.Contrast(img).enhance(3.0)
-        img = ImageEnhance.Sharpness(img).enhance(2.0)
-        img = img.filter(ImageFilter.SHARPEN)
-        text = pytesseract.image_to_string(
-            img,
-            config="--psm 6 --oem 3 -c tessedit_char_whitelist=0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz.+=-/() "
-        ).strip()
-        return text if len(text) > 10 else None
-    except Exception:
-        return None
+def image_to_base64(pil_image):
+    buffered = BytesIO()
+    pil_image.save(buffered, format="PNG")
+    return base64.b64encode(buffered.getvalue()).decode("utf-8")
 
 
 def solve_math(image_input, text_input, grade_level, user_key):
@@ -75,35 +65,31 @@ def solve_math(image_input, text_input, grade_level, user_key):
         return
 
     grade_note = f"Grade level: {grade_level}."
-    ocr_text = ""
-
-    if has_image:
-        yield "🔍 Reading text from your worksheet image..."
-        ocr_text = ocr_image(image_input)
-        if ocr_text:
-            yield f"📄 **Extracted from image:**\n```\n{ocr_text}\n```\n\n⚡ Solving step by step..."
-        else:
-            if not has_text:
-                yield (
-                    "⚠️ **Could not read text from the image clearly.**\n\n"
-                    "Please type the problem in the text box below and try again."
-                )
-                return
-            yield "⚡ Solving step by step..."
-    else:
-        yield "⚡ Solving step by step..."
 
-    parts = []
-    if ocr_text:
-        parts.append(f"Text extracted from worksheet image:\n{ocr_text}")
-    if has_text:
-        parts.append(f"Student's note: {text_input.strip()}")
-    if not parts:
-        parts.append(f"Problem: {text_input.strip()}")
-
-    user_msg = f"{grade_note}\n\n" + "\n\n".join(parts) + "\n\nPlease teach me how to solve this step by step."
+    yield "⚡ Solving step by step..."
 
     try:
+        if has_image:
+            # Use vision model to read the image directly
+            img_b64 = image_to_base64(image_input)
+            text_note = f"\nStudent's additional note: {text_input.strip()}" if has_text else ""
+            user_content = [
+                {
+                    "type": "image_url",
+                    "image_url": {
+                        "url": f"data:image/png;base64,{img_b64}"
+                    }
+                },
+                {
+                    "type": "text",
+                    "text": f"{grade_note}\n\nPlease read this math problem from the image carefully, identify all numbers and labels, then teach me how to solve it step by step.{text_note}"
+                }
+            ]
+            model = VISION_MODEL
+        else:
+            user_content = f"{grade_note}\n\nProblem: {text_input.strip()}\n\nPlease teach me how to solve this step by step."
+            model = TEXT_MODEL
+
         response = requests.post(
             GROQ_API_URL,
             headers={
@@ -111,10 +97,10 @@ def solve_math(image_input, text_input, grade_level, user_key):
                 "Content-Type": "application/json",
             },
             json={
-                "model": MODEL,
+                "model": model,
                 "messages": [
                     {"role": "system", "content": SYSTEM_PROMPT},
-                    {"role": "user", "content": user_msg},
+                    {"role": "user", "content": user_content},
                 ],
                 "max_tokens": 1200,
                 "temperature": 0,
@@ -134,11 +120,7 @@ def solve_math(image_input, text_input, grade_level, user_key):
             yield f"❌ API error {response.status_code}: {response.text[:200]}"
             return
 
-        prefix = ""
-        if ocr_text:
-            prefix = f"📄 **Extracted from image:**\n```\n{ocr_text}\n```\n\n"
-
-        output = prefix
+        output = ""
         for raw in response.iter_lines():
             if not raw:
                 continue
@@ -155,8 +137,8 @@ def solve_math(image_input, text_input, grade_level, user_key):
             except Exception:
                 continue
 
-        if output == prefix:
-            yield prefix + "\n\n⚠️ No response received. Please try again."
+        if not output:
+            yield "⚠️ No response received. Please try again."
 
     except requests.exceptions.Timeout:
         yield "❌ Request timed out. Please try again."
@@ -249,7 +231,7 @@ with gr.Blocks(css=CSS, title="Interactive Math Tutor") as demo:
 
             gr.HTML('<div class="panel-label" style="margin-top:16px">📎 Upload Worksheet</div>')
             image_input = gr.Image(type="pil", label="", height=180)
-            gr.HTML('<div style="font-size:.78rem;color:var(--muted);margin-top:4px">⚠️ For diagrams with numbers inside (triangles, graphs), also type the values in the text box below for best results.</div>')
+            gr.HTML('<div style="font-size:.78rem;color:var(--muted);margin-top:4px">📸 Images are read directly by AI — diagrams, triangles, and graphs all work!</div>')
 
             gr.HTML('<div class="panel-label" style="margin-top:14px">✏️ Or Type Your Problem</div>')
             text_input = gr.Textbox(label="", placeholder="e.g.  Solve for x: 5x − 3 = 22", lines=3)