Update app.py

app.py

@@ -1,71 +1,84 @@
+import os
 import gradio as gr
-from PIL import Image
-from pix2tex.cli import LatexOCR
-import sympy as sp
-from sympy.parsing.latex import parse_latex
-import re
 
-def preprocess_handwritten_image(pil_img):
-    return pil_img.convert('RGB')
+if os.path.exists("trained_model"):
+    from PIL import Image
+    from pix2tex.cli import LatexOCR
+    import sympy as sp
+    from sympy.parsing.latex import parse_latex
+    import re
 
-model = LatexOCR(weights='trained_model')
+    def preprocess_handwritten_image(pil_img):
+        return pil_img.convert('RGB')
 
-def clean_latex(latex):
-    latex = re.sub(r'\\(cal|mathcal)\s*X', 'x', latex)
-    latex = latex.replace('{', '').replace('}', '')
-    latex = latex.strip().rstrip(',.')
-    latex = re.sub(r'(\d+)\s*\\pi', r'(\1*3.1416)', latex)
-    latex = latex.replace(r'\pi', '3.1416')
-    latex = re.sub(r'(\d+)\s*e', r'(\1*2.7183)', latex)
-    latex = re.sub(r'(?<![a-zA-Z0-9])e(?![a-zA-Z0-9])', '2.7183', 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*([xX](\^\d+)?)', r'(\1)*\2', latex)
-    if '=' not in latex:
-        latex += '=0'
-    return latex
+    model = LatexOCR(weights='trained_model')
 
-def solve_polynomial(image):
-    try:
-        img = preprocess_handwritten_image(image)
-        latex_result = model(img)
-        cleaned_latex = clean_latex(latex_result)
-        expr = parse_latex(cleaned_latex)
+    def clean_latex(latex):
+        latex = re.sub(r'\\(cal|mathcal)\s*X', 'x', latex)
+        latex = latex.replace('{', '').replace('}', '')
+        latex = latex.strip().rstrip(',.')
+        latex = re.sub(r'(\d+)\s*\\pi', r'(\1*3.1416)', latex)
+        latex = latex.replace(r'\pi', '3.1416')
+        latex = re.sub(r'(\d+)\s*e', r'(\1*2.7183)', latex)
+        latex = re.sub(r'(?<![a-zA-Z0-9])e(?![a-zA-Z0-9])', '2.7183', 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*([xX](\^\d+)?)', r'(\1)*\2', latex)
+        if '=' not in latex:
+            latex += '=0'
+        return latex
 
-        output = f"## 📄 Extracted LaTeX\n```\n{latex_result}\n```\n---\n"
-        output += f"## 🧹 Cleaned LaTeX Used\n```\n{cleaned_latex}\n```\n---\n"
-        output += f"## 🧠 Parsed Expression\n\n$$ {sp.latex(expr)} $$\n---\n"
+    def solve_polynomial(image):
+        try:
+            img = preprocess_handwritten_image(image)
+            latex_result = model(img)
+            cleaned_latex = clean_latex(latex_result)
+            expr = parse_latex(cleaned_latex)
 
-        if isinstance(expr, sp.Equality):
-            lhs = expr.lhs - expr.rhs
-            output += "## ✏️ Step 1: Standard Form\n"
-            output += f"$$ {sp.latex(lhs)} = 0 $$\n---\n"
-            output += "## 🧩 Step 2: Factorized\n"
-            output += f"$$ {sp.latex(sp.factor(lhs))} = 0 $$\n---\n"
-            output += "## ✅ Step 3: Roots\n"
-            roots = sp.solve(sp.Eq(lhs, 0), dict=True)
-            if roots:
-                output += "$$\n\\begin{aligned}\n"
-                for i, sol in enumerate(roots, 1):
-                    for var, val in sol.items():
-                        output += f"\\text{{Root {i}}}: {var} &= {sp.latex(val)}\\\\\n"
-                output += "\\end{aligned}\n$$"
-        else:
-            output += "## ➕ Simplified\n"
-            output += f"$$ {sp.latex(sp.simplify(expr))} $$"
+            output = f"## 📄 Extracted LaTeX\n```\n{latex_result}\n```\n---\n"
+            output += f"## 🧹 Cleaned LaTeX Used\n```\n{cleaned_latex}\n```\n---\n"
+            output += f"## 🧠 Parsed Expression\n\n$$ {sp.latex(expr)} $$\n---\n"
 
-        return output
-    except Exception as e:
-        return f"❌ **Error**: {str(e)}"
+            if isinstance(expr, sp.Equality):
+                lhs = expr.lhs - expr.rhs
+                output += "## ✏️ Step 1: Standard Form\n"
+                output += f"$$ {sp.latex(lhs)} = 0 $$\n---\n"
+                output += "## 🧩 Step 2: Factorized\n"
+                output += f"$$ {sp.latex(sp.factor(lhs))} = 0 $$\n---\n"
+                output += "## ✅ Step 3: Roots\n"
+                roots = sp.solve(sp.Eq(lhs, 0), dict=True)
+                if roots:
+                    output += "$$\n\\begin{aligned}\n"
+                    for i, sol in enumerate(roots, 1):
+                        for var, val in sol.items():
+                            output += f"\\text{{Root {i}}}: {var} &= {sp.latex(val)}\\\\\n"
+                    output += "\\end{aligned}\n$$"
+            else:
+                output += "## ➕ Simplified\n"
+                output += f"$$ {sp.latex(sp.simplify(expr))} $$"
 
-demo = gr.Interface(
-    fn=solve_polynomial,
-    inputs=gr.Image(type="pil", label="📷 Upload Image"),
-    outputs=gr.Markdown(label="📋 Solution"),
-    title="🧠 Handwritten Polynomial Solver",
-    description="Extract handwritten polynomials and solve step-by-step.",
-    allow_flagging="never"
-)
+            return output
+        except Exception as e:
+            return f"❌ **Error**: {str(e)}"
+
+    demo = gr.Interface(
+        fn=solve_polynomial,
+        inputs=gr.Image(type="pil", label="📷 Upload Image"),
+        outputs=gr.Markdown(label="📋 Solution"),
+        title="🧠 Handwritten Polynomial Solver",
+        description="Upload a handwritten polynomial to extract, solve, and view steps.",
+        allow_flagging="never"
+    )
+
+else:
+    # Dummy interface before training is done
+    demo = gr.Interface(
+        fn=lambda: "Training not completed yet. Please run `train.py` to generate `trained_model/`.",
+        inputs=[],
+        outputs="text",
+        title="⚠️ Model Not Ready",
+        description="Run training first to enable polynomial solving."
+    )
 
 demo.launch()