Update app.py

app.py

@@ -10,6 +10,48 @@ model = LatexOCR()
 def preprocess_handwritten_image(pil_img):
     return pil_img.convert('RGB')
 
+# --- Tab 1 cleaner (UNCHANGED) ---
+def clean_latex(latex):
+    latex = latex.replace('\\ ', '')
+    latex = latex.replace('\\\\', '\\')
+    latex = re.sub(r'\\[ \t\n\r\f\v]*', '', latex)
+    latex = re.sub(r'\\([+\-=])', r'\1', latex)
+    replacements = {
+        r'\\chi': 'x', r'chi': 'x',
+        r'\\xi': 'x', r'xi': 'x',
+        r'\\alpha': 'x', r'alpha': 'x',
+        r'\\beta': 'b', r'beta': 'b',
+        r'\\gamma': 'y', r'gamma': 'y',
+        r'\\vartheta': '3', r'vartheta': '3',
+        r'\\mathcalW': 'x', r'mathcalW': 'x',
+        r'\\pi': 'pi', r'pi': 'pi',
+        r'\\mathrm': '', r'mathrm': '',
+    }
+    for wrong, correct in replacements.items():
+        latex = re.sub(wrong, correct, latex)
+    latex = re.sub(r'\\(cal|mathcal)\s*\{?\s*[Xx]\s*\}?', 'x', latex)
+    latex = re.sub(r'\\(cal|mathcal)\s*\{?\s*[Yy]\s*\}?', 'y', latex)
+    latex = re.sub(r'\\(cal|mathcal)\s*\{?\s*[Zz]\s*\}?', 'z', latex)
+    latex = latex.replace('cal x', 'x').replace('cal X', 'x')
+    latex = latex.replace('mathcal x', 'x').replace('mathcal X', 'x')
+    latex = latex.replace('{', '').replace('}', '')
+    latex = latex.strip().rstrip(',.')
+    latex = re.sub(r'(?<![a-zA-Z0-9])e(?![a-zA-Z0-9])', 'E', latex)
+    latex = re.sub(r'(\d)([a-zA-Z])', r'\1*\2', latex)
+    latex = re.sub(r'(\d+)\s*i', r'\1*I', latex)
+    latex = re.sub(r'(?<![a-zA-Z0-9])i(?![a-zA-Z0-9])', 'I', latex)
+    latex = re.sub(r'\(([^()]+?)\)\s*([a-zA-Z](\^\d+)?)', r'(\1)*\2', latex)
+    latex = latex.replace(r'\cdot', '*')
+    latex = latex.replace('−', '-')
+    latex = re.sub(r'[^\w\s^=+*\-().]', '', latex)
+    if '=' not in latex:
+        latex += '=0'
+    latex = latex.replace('pi', '3.1416')
+    latex = latex.replace('e', '2.7183')
+    latex = latex.replace('E', '2.7183')
+    return latex
+
+# --- Tab 2 cleaner ---
 def clean_latex2(latex):
     latex = latex.replace('\\ ', '')
     latex = latex.replace('\\\\', '\\')
@@ -34,59 +76,96 @@ def clean_latex2(latex):
         latex = re.sub(wrong, right, latex)
     return latex.strip()
 
-def solve_from_multiple_images(img1, img2, img3):
+# --- Polynomial Solver (unchanged) ---
+def solve_polynomial(image):
     try:
-        images = [img for img in [img1, img2, img3] if img is not None]
-        if len(images) < 2:
-            return "❌ Please upload at least 2 images (each with one equation)."
+        img = preprocess_handwritten_image(image)
+        latex_result = model(img)
+        if not latex_result or len(latex_result.strip()) < 2:
+            return "❌ Could not extract valid LaTeX from image.", "", ""
+        cleaned_latex = clean_latex(latex_result)
+        expr = parse_latex(cleaned_latex)
+        expr = expr.subs(sp.pi, sp.Float(3.1416)).subs(sp.E, sp.Float(2.7183))
+        if not isinstance(expr, sp.Equality):
+            raise ValueError("Expression is not an equation.")
+        lhs = expr.lhs - expr.rhs
+        if not lhs.is_polynomial():
+            raise ValueError("Not a polynomial")
+        output = f"## 📄 Extracted LaTeX\n```latex\n{latex_result}\n```\n"
+        output += f"---\n## 🧹 Cleaned LaTeX\n```latex\n{cleaned_latex}\n```\n"
+        output += f"---\n## 🧠 Parsed Expression\n$$ {sp.latex(expr)} $$\n"
+        factor = sp.factor(lhs)
+        output += f"---\n## ✏️ Standard Form\n$$ {sp.latex(lhs)} = 0 $$\n"
+        output += f"---\n## 🧩 Factorized\n$$ {sp.latex(factor)} = 0 $$\n"
+        x = next(iter(lhs.free_symbols))
+        roots = sp.solve(lhs, x)
+        output += "## ✅ Roots\n"
+        output += "$$\n\\begin{aligned}\n"
+        for i, r in enumerate(roots, 1):
+            root_val = sp.N(r, 6)
+            output += f"\\text{{Root {i}}}:\\quad x &\\approx {sp.latex(root_val)} \\\\\n"
+        output += "\\end{aligned}\n$$\n"
+        return output, cleaned_latex, ""
+    except Exception as e:
+        return f"❌ Error: {str(e)}", "", ""
 
-        raw_latexes = []
-        cleaned_eqs = []
-        for img in images:
+# --- NEW Multi-image System Solver ---
+def solve_system_multi(img1, img2, img3):
+    images = [img1, img2, img3]
+    cleaned_equations = []
+    raw_outputs = []
+    for img in images:
+        if img is None:
+            continue
+        try:
             img = preprocess_handwritten_image(img)
             latex = model(img)
-            raw_latexes.append(latex)
+            raw_outputs.append(latex)
             cleaned = clean_latex2(latex)
-            if '=' in cleaned:
-                try:
-                    parsed = parse_latex(cleaned)
-                    if isinstance(parsed, sp.Equality):
-                        cleaned_eqs.append(parsed)
-                except:
-                    continue
+            parsed = parse_latex(cleaned)
+            if isinstance(parsed, sp.Equality):
+                cleaned_equations.append(parsed)
+        except:
+            continue
 
-        if len(cleaned_eqs) < 2:
-            return "❌ Could not parse enough valid equations."
+    if len(cleaned_equations) < 2:
+        return f"❌ Could not parse enough valid equations."
 
-        symbols = set()
-        for eq in cleaned_eqs:
-            symbols.update(eq.free_symbols)
+    symbols = set()
+    for eq in cleaned_equations:
+        symbols.update(eq.free_symbols)
 
-        solution = sp.solve(cleaned_eqs, list(symbols), dict=True)
-        output = "## 🧾 Extracted Equations\n"
-        for latex in raw_latexes:
-            output += f"```latex\n{latex}\n```\n"
-        output += "---\n## ✏️ Parsed & Cleaned:\n"
-        for eq in cleaned_eqs:
-            output += f"$$ {sp.latex(eq)} $$\n"
-        output += "---\n## ✅ Solution:\n"
-        if solution:
-            output += "$$" + sp.latex(solution[0]) + "$$"
-        else:
-            output += "❌ No solution found or system is inconsistent."
-        return output
-    except Exception as e:
-        return f"❌ Error: {str(e)}"
+    sol = sp.solve(cleaned_equations, list(symbols), dict=True)
+    output = "## 📄 Raw LaTeX:\n"
+    for latex in raw_outputs:
+        output += f"- ```latex\n{latex}\n```\n"
+    output += "---\n## ✏️ Parsed Equations:\n"
+    for eq in cleaned_equations:
+        output += f"$$ {sp.latex(eq)} $$\n"
+    output += "---\n## ✅ Solution:\n"
+    if sol:
+        output += "$$" + sp.latex(sol[0]) + "$$"
+    else:
+        output += "❌ No solution found or system may be inconsistent."
+    return output
 
+# === UI ===
 with gr.Blocks() as demo:
-    gr.Markdown("## 📐 Solve System from Multiple Equation Images")
-    with gr.Row():
+    with gr.Tab("🖼️ Parse from Image"):
+        image_input = gr.Image(type="pil", label="📷 Upload Image of Polynomial")
+        hidden_latex = gr.Textbox(visible=False)
+        explanation_prompt = gr.Textbox(visible=False)
+        output_box = gr.Markdown(label="📋 Step-by-step Solution")
+        submit_btn = gr.Button("🔍 Solve")
+        submit_btn.click(fn=solve_polynomial, inputs=[image_input], outputs=[output_box, hidden_latex, explanation_prompt])
+
+    with gr.Tab("📐 Solve System (Multi-Image)"):
         img1 = gr.Image(type="pil", label="📷 Equation 1")
         img2 = gr.Image(type="pil", label="📷 Equation 2")
-        img3 = gr.Image(type="pil", label="📷 (Optional) Equation 3")
-    solve_btn = gr.Button("📌 Solve System")
-    output = gr.Markdown()
-    solve_btn.click(fn=solve_from_multiple_images, inputs=[img1, img2, img3], outputs=output)
+        img3 = gr.Image(type="pil", label="📷 Equation 3 (Optional)")
+        sys_output = gr.Markdown(label="📋 Solved System Output")
+        sys_btn = gr.Button("📌 Solve System")
+        sys_btn.click(fn=solve_system_multi, inputs=[img1, img2, img3], outputs=[sys_output])
 
 if __name__ == "__main__":
     demo.launch()