from fastapi import FastAPI, Form, UploadFile, File from fastapi.responses import JSONResponse from PIL import Image import torch, clip, io, re, numpy as np, cv2 from pix2text import Pix2Text import pytesseract import time, os import logging from paddleocr import PaddleOCR from datetime import datetime import json # ------------------------------- # Logging setup # ------------------------------- # Ensure logs directory exists # os.makedirs("logs", exist_ok=True) # #---------------------------------------------------------------------- # Logger Function # ---------------------------------------------------------------------- logs_dir = 'logs' if not os.path.exists(logs_dir): os.makedirs(logs_dir) log_formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s') log_file_path = os.path.join(logs_dir, datetime.now().strftime('%Y-%m-%d') + '.log') file_handler = logging.FileHandler(log_file_path) file_handler.setFormatter(log_formatter) logger = logging.getLogger() logger.addHandler(file_handler) logger.setLevel(logging.INFO) # ------------------------------- # Load models # ------------------------------- device = "cuda" if torch.cuda.is_available() else "cpu" clip_model, preprocess = clip.load("ViT-B/32", device=device) # ocr = PaddleOCR(use_angle_cls=True, lang="en") # Labels for classification math_labels = [ "a math equation", "a handwritten equation", "scientific symbols", "algebraic formula", "a trigonometric equation", "a linear algebra equation", "a calculus expression", ] physics_labels = [ "a physics formula", "physics symbols", "electromagnetic equation", "mechanics equation", "thermodynamics formula", ] chemistry_labels = [ "a chemistry formula", "chemical equation", "molecular structure", "periodic element symbols", "organic chemistry formula", ] text_labels = [ "a page of text", "a handwritten note", "a printed paragraph", "a sentence of text", ] all_labels = math_labels + physics_labels + chemistry_labels + text_labels text_tokens = clip.tokenize(all_labels).to(device) # OCR Models p2t = Pix2Text() # ------------------------------- # Regex for symbols # ------------------------------- SCI_SYMBOLS = re.compile( r"[=+\-*/^√∫∮∑ΣΠπ∞≤≥≠<>±‰°ℏ]" # math/physics operators r"|[ΔΩθλμνρσΦΨΓαβγδ∂∇→←⇌⇔⇄↔]" # greek & physics r"|[∀∃∈∉⊂⊆⊇⊃∪∩∧∨¬⇒⇔]" # logic & set theory r"|[\[\]\(\)\{\}\|⟨⟩‖]" # brackets of all types ) # def run_paddle_ocr(image: Image.Image) -> str: # """Run PaddleOCR on a PIL image and return text.""" # img_np = np.array(image) # results = ocr.ocr(img_np) # if not results: # return "" # lines = [] # for res in results[0]: # text, conf = res[1] # lines.append(text) # return " ".join(lines) # raw_text = run_paddle_ocr(image) # ------------------------------- # Helpers # ------------------------------- def preprocess_text_image(image: Image.Image) -> Image.Image: """Enhance text images for better OCR.""" img = np.array(image) gray = cv2.cvtColor(img, cv2.COLOR_RGB2GRAY) gray = cv2.convertScaleAbs(gray, alpha=1.5, beta=20) # contrast _, thresh = cv2.threshold(gray, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU) return Image.fromarray(thresh) def clean_text(text: str) -> str: """Remove noise like form feeds and empty lines.""" return "\n".join( [line.strip() for line in text.splitlines() if line.strip()] ).replace("\f", "") raw = r"""A 13.7 kg mass is suspended on a \(1 \times1 0^{5} \, \mathrm{N /}\) m spring. The mass oscillates up and down from the equilibrium position $$ y_{\mathrm{{e q}}}=0 $$ according to $$ $ y ( t )=A \, \sin( \omega t+\phi_{0} ) \, . $$\(At time\)t_{0} \,=\, 0\(the mass is at 29.6 cm and moving upward at velocity\)+ 6 . 3 2 \; \mathrm{m / s .}$""" import re def clean_latex(text: str) -> str: r""" Normalize and wrap LaTeX properly: - Fix OCR spacing and broken $ signs - Inline math -> \( ... \) - Display math -> $$ ... $$ """ if not text: return "" # Step 1. Fix spacing text = re.sub(r"\s+", " ", text).strip() # Step 2. Normalize all math markers (\( \), \[ \]) -> $...$ text = text.replace(r"\(", "$").replace(r"\)", "$") text = text.replace(r"\[", "$$").replace(r"\]", "$$") # Step 3. Remove extra/trailing orphan $ (like "$$ $" or "$ $") text = re.sub(r"\$\s*\$", "$", text) # Step 4. Wrap equations properly def wrap_equation(match): expr = match.group(1).strip() # Keywords that suggest display mode display_keywords = [ "=", "\\int", "\\sum", "\\prod", "\\lim", "\\sin", "\\cos", "\\tan", "\\sqrt", "\\frac", "y(", "A ", "\\begin", "\\end", ] if any(kw in expr for kw in display_keywords): return f"$$ {expr} $$" else: return f"\\({expr}\\)" # Replace $...$ with proper wrapping text = re.sub(r"\$(.+?)\$", wrap_equation, text) return text def is_valid_latex(expr: str, conf: float, min_conf: float = 0.5) -> bool: """Check if OCR output looks like valid LaTeX/scientific expression.""" expr = expr.strip() if not expr or conf < min_conf: return False if len(expr) < 3: return False latex_keywords = [ # ---------------- Numbers / General ---------------- "\\frac", "\\dfrac", "\\tfrac", "\\cfrac", "\\binom", "\\sqrt", "\\root", "\\overline", "\\underline", # ---------------- Operators ---------------- "\\int", "\\iint", "\\iiint", "\\iiiint", "\\oint", "\\oiint", "\\sum", "\\prod", "\\coprod", "\\lim", "\\sup", "\\inf", "\\log", "\\ln", "\\exp", # ---------------- Derivatives / Calculus ---------------- "\\partial", "\\nabla", "\\grad", "\\div", "\\curl", "\\dot", "\\ddot", "\\prime", "\\operatorname{d}", # ---------------- Functions ---------------- "\\sin", "\\cos", "\\tan", "\\cot", "\\sec", "\\csc", "\\arcsin", "\\arccos", "\\arctan", "\\arccot", "\\arcsec", "\\arccsc", "\\sinh", "\\cosh", "\\tanh", "\\coth", "\\sech", "\\csch", "\\max", "\\min", # ---------------- Linear Algebra ---------------- "\\vec", "\\hat", "\\tilde", "\\bar", "\\begin{bmatrix}", "\\end{bmatrix}", "\\begin{pmatrix}", "\\end{pmatrix}", "\\begin{matrix}", "\\end{matrix}", "\\begin{vmatrix}", "\\end{vmatrix}", "\\det", "\\operatorname{rank}", "\\operatorname{trace}", "^T", "^{-1}", # ---------------- Sets / Logic ---------------- "\\in", "\\notin", "\\subset", "\\subseteq", "\\supset", "\\supseteq", "\\cup", "\\cap", "\\setminus", "\\emptyset", "\\varnothing", "\\forall", "\\exists", "\\neg", "\\land", "\\lor", "\\Rightarrow", "\\Leftrightarrow", "\\implies", "\\iff", # ---------------- Relations ---------------- "=", "\\approx", "\\equiv", "\\cong", "\\sim", "\\propto", "<", ">", "\\leq", "\\geq", "\\neq", "\\ll", "\\gg", "\\pm", "\\mp", # ---------------- Greek Letters ---------------- "\\alpha", "\\beta", "\\gamma", "\\Gamma", "\\delta", "\\Delta", "\\epsilon", "\\varepsilon", "\\zeta", "\\eta", "\\theta", "\\Theta", "\\iota", "\\kappa", "\\lambda", "\\Lambda", "\\mu", "\\nu", "\\xi", "\\Xi", "\\pi", "\\Pi", "\\rho", "\\sigma", "\\Sigma", "\\tau", "\\upsilon", "\\Upsilon", "\\phi", "\\varphi", "\\Phi", "\\chi", "\\psi", "\\Psi", "\\omega", "\\Omega", # ---------------- Brackets & Delimiters ---------------- "\\left(", "\\right)", "\\left[", "\\right]", "\\left\\{", "\\right\\}", "\\left|", "\\right|", "\\left\\langle", "\\right\\rangle", "\\left\\lfloor", "\\right\\rfloor", "\\left\\lceil", "\\right\\rceil", "\\Big", "\\big", "\\bigg", "\\Bigg", "\\middle", # ---------------- Special Symbols ---------------- "\\infty", "\\aleph", "\\hbar", "\\imath", "\\jmath", "\\Re", "\\Im", "\\nabla", "\\partial", "\\degree", # ---------------- Probability & Stats ---------------- "\\Pr", "\\mathbb{E}", "\\mathbb{V}", "\\operatorname{Var}", "\\operatorname{Cov}", # ---------------- Blackboard (sets) ---------------- "\\mathbb{R}", "\\mathbb{Q}", "\\mathbb{Z}", "\\mathbb{N}", "\\mathbb{C}", # ---------------- Script / Calligraphic ---------------- "\\mathcal{A}", "\\mathcal{B}", "\\mathcal{C}", "\\mathcal{D}", "\\mathcal{E}", "\\mathcal{F}", "\\mathcal{G}", "\\mathcal{H}", "\\mathcal{I}", "\\mathcal{J}", "\\mathcal{K}", "\\mathcal{L}", "\\mathcal{M}", "\\mathcal{N}", "\\mathcal{O}", "\\mathcal{P}", "\\mathcal{Q}", "\\mathcal{R}", "\\mathcal{S}", "\\mathcal{T}", "\\mathcal{U}", "\\mathcal{V}", "\\mathcal{W}", "\\mathcal{X}", "\\mathcal{Y}", "\\mathcal{Z}", # ---------------- Arrows ---------------- "\\to", "\\mapsto", "\\longrightarrow", "\\implies", "\\iff", "\\leftrightarrow", "\\longleftrightarrow", "\\uparrow", "\\downarrow", "\\updownarrow", "\\Rightarrow", "\\Leftarrow", "\\Leftrightarrow", "\\hookrightarrow", "\\hookleftarrow", "\\xrightarrow{}", "\\xleftarrow{}", # AMS math "\\boxed", # main command for boxing equations # LaTeX core "\\fbox", # text mode framing "\\framebox", # allows setting box width/position "\\mbox", # creates an unbreakable box in text "\\makebox", # adjustable width box "\\parbox", # paragraph box "\\raisebox", # vertical adjustment of content # Colors & Fancy Boxes (if color/xcolor or tcolorbox is used) "\\colorbox", # background colored box "\\fcolorbox", # frame + background color "\\textcolor", # colored text "\\highlight", # (often defined by packages) # Math emphasis variants "\\overbrace", "\\underbrace", "\\overline", "\\underline", "\\cancel", "\\bcancel", "\\xcancel", "\\cancelto", # from cancel package # Answer & Highlight environments (common in exams/books) "\\shabox", "\\Ovalbox", "\\doublebox", "\\dashbox", # from fancybox package "\\tcolorbox", "\\begin{tcolorbox}", "\\end{tcolorbox}", # modern colored boxes ] if any(kw in expr for kw in latex_keywords): return True if re.search(SCI_SYMBOLS, expr): return True return False # ------------------------------- # Classifier # ------------------------------- def classify_image_type(image: Image.Image, threshold: float = 0.6): image_tensor = preprocess(image).unsqueeze(0).to(device) with torch.no_grad(): logits_per_image, _ = clip_model(image_tensor, text_tokens) probs = logits_per_image.softmax(dim=-1).cpu().numpy()[0] n_math, n_phys, n_chem = ( len(math_labels), len(physics_labels), len(chemistry_labels), ) math_score = float(sum(probs[:n_math])) phys_score = float(sum(probs[n_math : n_math + n_phys])) chem_score = float(sum(probs[n_math + n_phys : n_math + n_phys + n_chem])) text_score = float(sum(probs[n_math + n_phys + n_chem :])) raw_text = pytesseract.image_to_string(image) sci_chars = len(SCI_SYMBOLS.findall(raw_text)) ratio = sci_chars / max(1, len(raw_text)) # Decide class scores = { "math": math_score, "physics": phys_score, "chemistry": chem_score, "text": text_score, } best_label = max(scores, key=scores.get) best_score = scores[best_label] if best_score > threshold or ratio > 0.25: return best_label, best_score else: return "uncertain", best_score # ------------------------------- # FastAPI App # ------------------------------- app = FastAPI() # ------------------------------- # FastAPI Endpoint # ------------------------------- @app.post("/image") async def classify_and_convert( file: UploadFile = File(...), category_slug: str = Form("other") ): try: contents = await file.read() image = Image.open(io.BytesIO(contents)).convert("RGB") # Save uploaded image with timestamp folder_name = f"received_images/{time.strftime('%Y-%m-%d')}" os.makedirs(folder_name, exist_ok=True) saved_path = os.path.join(folder_name, file.filename) image.save(saved_path) logger.info(f"[UPLOAD] Saved image to {saved_path}") # return # Step 1: classify label, conf = classify_image_type(image) logger.info(f"[CLASSIFY] Predicted label={label}, confidence={conf:.3f}") # Step 2: decide recognition method if (category_slug in ["math", "physics", "chemistry"]) or ( category_slug == "other" and ((label in ["math", "physics", "chemistry"]) or label == "uncertain") ): logger.info("[OCR] Running Pix2Text (math/science mode)") result = p2t.recognize(image) if isinstance(result, list): question = " ".join( [r["text"] if isinstance(r, dict) else str(r) for r in result] ) else: question = str(result).strip() logger.debug(f"[OCR-RAW] Pix2Text result: {question[:200]}") # Step 3: validate + clean LaTeX if is_valid_latex(question, conf): is_latex_flag = is_valid_latex(question, conf) question = clean_latex(question) logger.info("[VALIDATION] Recognized valid LaTeX") else: is_latex_flag = is_valid_latex(question, conf) logger.warning( "[VALIDATION] Invalid LaTeX, falling back to Tesseract OCR" ) processed = preprocess_text_image(image) question = clean_text(pytesseract.image_to_string(processed)) question = clean_latex(question) else: logger.info("[OCR] Running plain OCR (text mode)") processed = preprocess_text_image(image) raw_text = pytesseract.image_to_string(processed) question = clean_text(raw_text) question = clean_latex(question) is_latex_flag = is_valid_latex(question, conf) # Heuristic confidence for text (ratio of valid characters) valid_chars = sum(c.isalnum() or c.isspace() for c in raw_text) text_conf = valid_chars / max(1, len(raw_text)) conf = round((conf + text_conf) / 2, 3) # blend CLIP + OCR heuristic logger.info(f"[TEXT-CONF] Adjusted text confidence={conf:.3f}") # Step 4: build response response = { "status": True, "message": "Image processed successfully", "data": { "text": question, "type": label, "confidence": float(conf), "is_latex": is_latex_flag, # 👈 added }, } logger.info("[RESPONSE] " + json.dumps(response, ensure_ascii=False, indent=2)) return JSONResponse(content=response) except Exception as e: logger.exception("Error in classify_and_convert") return JSONResponse(status_code=500, content={"error": str(e)})