Update app.py

app.py

@@ -6,6 +6,9 @@ from sympy.parsing.latex import parse_latex
 import re
 import os
 
+# Optional: import training function
+from train import train_model
+
 # Trigger data download only once
 if not os.path.exists("dataset/train"):
     print("🚀 Running data preparation scripts...")
@@ -13,58 +16,41 @@ if not os.path.exists("dataset/train"):
     os.system("python generate_csv.py")
 
 
-# Preprocessing for handwritten image
+# Preprocessing
 def preprocess_handwritten_image(pil_img):
-    return pil_img.convert('RGB')  # Minimal processing, just convert to RGB
+    return pil_img.convert('RGB')
 
 
-# Load Pix2Tex model (once)
+# Load Pix2Tex model
 model = LatexOCR()
 
+
 # Clean LaTeX output
 def clean_latex(latex):
-    # Replace \mathcal{X} or \cal X with 'x'
     latex = re.sub(r'\\(cal|mathcal)\s*X', 'x', latex)
-
-    # Remove curly braces
     latex = latex.replace('{', '').replace('}', '')
     latex = latex.strip().rstrip(',.')
-    
-    # Replace coefficients like 5\pi with (5*3.1416)
     latex = re.sub(r'(\d+)\s*\\pi', r'(\1*3.1416)', latex)
     latex = latex.replace(r'\pi', '3.1416')
-
-    # Replace coefficients like 5e with (5*2.7183)
     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)
-
-    # Insert * between number and variable (e.g., 45x → 45*x)
     latex = re.sub(r'(\d)([a-zA-Z])', r'\1*\2', latex)
-
-    # Replace number followed by i with number*I
     latex = re.sub(r'(\d+)\s*i', r'\1*I', latex)
-
-    # Replace standalone i with I
     latex = re.sub(r'(?<![a-zA-Z0-9])i(?![a-zA-Z0-9])', 'I', latex)
-
-    # Wrap complex coefficients with variables: (a+bI)x^n → (a+b*I)*x^n
     latex = re.sub(r'\(([^()]+?)\)\s*([xX](\^\d+)?)', r'(\1)*\2', latex)
-
-    # Append '=0' if not already present
     if '=' not in latex:
         latex += '=0'
-
     return latex
 
 
-# Main function
+# Solver logic
 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)
+
         output = f"## 📄 Extracted LaTeX\n```\n{latex_result}\n```\n"
         output += "---\n"
         output += f"## 🧹 Cleaned LaTeX Used\n```\n{cleaned_latex}\n```\n"
@@ -74,7 +60,6 @@ def solve_polynomial(image):
 
         if isinstance(expr, sp.Equality):
             lhs = expr.lhs - expr.rhs
-
             output += "## ✏️ Step 1: Standard Form of the Polynomial\n"
             output += f"$$ {sp.latex(lhs)} = 0 $$\n"
             output += "---\n"
@@ -86,34 +71,49 @@ def solve_polynomial(image):
 
             output += "## ✅ Step 3: Solve for 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}}}:\\quad {var} &= {sp.latex(val)} \\\\\n"
                 output += "\\end{aligned}\n$$\n"
-
         else:
             simplified = sp.simplify(expr)
             output += "## ➕ Simplified Expression\n"
             output += f"$$ {sp.latex(simplified)} $$"
 
         return output
-
     except Exception as e:
         return f"❌ **Error**: {str(e)}"
 
 
-# Gradio UI
-demo = gr.Interface(
-    fn=solve_polynomial,
-    inputs=gr.Image(type="pil", label="📷 Upload Image of Polynomial"),
-    outputs=gr.Markdown(label="📋 Step-by-step Solution"),
-    title="🧠 Polynomial Solver from Image",
-    description="Upload an image of a polynomial (typed or handwritten). The app will extract, solve, and explain it step-by-step.",
-    allow_flagging="never"
-)
+# Trigger training
+def run_training():
+    try:
+        train_model("train.yaml")  # path to your training config
+        return "✅ Training completed successfully."
+    except Exception as e:
+        return f"❌ Training failed: {str(e)}"
+
+
+# Gradio UI with training button
+with gr.Blocks() as demo:
+    gr.Markdown("## 🧠 Polynomial Solver from Handwritten Image")
+    gr.Markdown("Upload a handwritten polynomial image. The app will extract and solve it step-by-step.")
+
+    with gr.Row():
+        image_input = gr.Image(type="pil", label="📷 Upload Polynomial Image")
+        solution_output = gr.Markdown(label="📋 Step-by-step Solution")
+
+    image_input.change(fn=solve_polynomial, inputs=image_input, outputs=solution_output)
+
+    gr.Markdown("----")
+
+    gr.Markdown("## 🛠 Optional: Fine-tune Pix2Tex (CPU)")
+    train_btn = gr.Button("🎓 Start CPU Training")
+    train_output = gr.Textbox(label="Training Status")
+
+    train_btn.click(fn=run_training, outputs=train_output)
+
+demo.launch()
 
-if __name__ == "__main__":
-    demo.launch()