Initial Captcha OCR Space

app.py

@@ -4,129 +4,73 @@ import onnxruntime as rt
 from torchvision import transforms as T
 from pathlib import Path
 from PIL import Image
-from fastapi import FastAPI, UploadFile, File, Body
-from fastapi.responses import JSONResponse
-from pydantic import BaseModel
 from utils.tokenizer_base import Tokenizer
+import gradio as gr
 import io
-import os
 import base64
-import gradio as gr
+import os
 
-# =========================
+# =====================
 # MODEL SETUP
-# =========================
-model_path = "models/model.onnx"
-cwd = Path(__file__).parent.resolve()
-model_file = os.path.join(cwd, model_path)
-
-if not os.path.exists(model_file):
-    raise FileNotFoundError(f"Model not found at {model_file}")
+# =====================
+model_file = Path(__file__).parent / "models/model.onnx"
+if not model_file.exists():
+    raise RuntimeError(f"Model not found at {model_file}")
 
 img_size = (32, 128)
 vocab = r"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~"
 tokenizer = Tokenizer(vocab)
 
+transform = T.Compose([
+    T.Resize(img_size, T.InterpolationMode.BICUBIC),
+    T.ToTensor(),
+    T.Normalize(0.5, 0.5),
+])
 
-def to_numpy(tensor):
-    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()
-
-
-def get_transform(img_size):
-    return T.Compose([
-        T.Resize(img_size, T.InterpolationMode.BICUBIC),
-        T.ToTensor(),
-        T.Normalize(0.5, 0.5),
-    ])
-
+session = rt.InferenceSession(str(model_file))
 
-def load_model(model_file):
-    transform = get_transform(img_size)
-    onnx_model = onnx.load(model_file)
-    onnx.checker.check_model(onnx_model)
-    session = rt.InferenceSession(model_file)
-    return transform, session
 
+def to_numpy(t):
+    return t.detach().cpu().numpy()
 
-transform, session = load_model(model_file)
 
-# =========================
-# SHARED INFERENCE LOGIC
-# =========================
-def predict_from_image(img: Image.Image) -> str:
+def infer(img: Image.Image):
     x = transform(img.convert("RGB")).unsqueeze(0)
-    ort_inputs = {session.get_inputs()[0].name: to_numpy(x)}
-    logits = session.run(None, ort_inputs)[0]
+    logits = session.run(None, {session.get_inputs()[0].name: to_numpy(x)})[0]
     probs = torch.tensor(logits).softmax(-1)
     preds, _ = tokenizer.decode(probs)
     return preds[0]
 
-# =========================
-# FASTAPI SETUP
-# =========================
-app = FastAPI(title="OCR CAPTCHA API")
-
-class Base64ImageRequest(BaseModel):
-    image_base64: str
-
-
-@app.post("/predict/file")
-async def predict_file(file: UploadFile = File(...)):
-    """
-    Accepts raw bytes (multipart/form-data)
-    """
-    try:
-        contents = await file.read()
-        img = Image.open(io.BytesIO(contents))
-        result = predict_from_image(img)
-        return {"predicted_text": result}
-    except Exception as e:
-        return JSONResponse({"error": str(e)}, status_code=500)
-
-
-@app.post("/predict/base64")
-async def predict_base64(payload: Base64ImageRequest):
-    """
-    Accepts base64-encoded image
-    """
-    try:
-        image_bytes = base64.b64decode(payload.image_base64)
-        img = Image.open(io.BytesIO(image_bytes))
-        result = predict_from_image(img)
-        return {"predicted_text": result}
-    except Exception as e:
-        return JSONResponse({"error": str(e)}, status_code=500)
-
-# =========================
-# GRADIO UI
-# =========================
-def gradio_predict(img: Image.Image):
-    if img is None:
-        return ""
-    return predict_from_image(img)
-
-
-gradio_ui = gr.Interface(
-    fn=gradio_predict,
-    inputs=gr.Image(type="pil", label="Input Image"),
-    outputs=gr.Textbox(label="Predicted Text"),
-    title="OCR CAPTCHA Solver",
-    description="OCR model for captcha images (letters + numbers).",
-    examples=[
-        "examples/1.png",
-        "examples/2.jpg",
-    ],
-)
-
-# =========================
-# MOUNT GRADIO INTO FASTAPI
-# =========================
-app = gr.mount_gradio_app(app, gradio_ui, path="/")
-
-
-# =========================
-# LOCAL RUN
-# =========================
-# if __name__ == "__main__":
-#     import uvicorn
-#     uvicorn.run("app:app", host="0.0.0.0", port=8000, reload=True)
+
+# =====================
+# GRADIO FUNCTIONS
+# =====================
+def predict_image(img):
+    return infer(img)
+
+
+def predict_base64(b64: str):
+    img_bytes = base64.b64decode(b64)
+    img = Image.open(io.BytesIO(img_bytes))
+    return infer(img)
+
+
+# =====================
+# GRADIO APP (REQUIRED)
+# =====================
+with gr.Blocks(title="Captcha OCR") as demo:
+    gr.Markdown("# Captcha OCR")
+    gr.Markdown("OCR for captcha images (letters & numbers)")
+
+    with gr.Tab("Image Upload"):
+        img = gr.Image(type="pil")
+        out = gr.Textbox()
+        gr.Button("Predict").click(predict_image, img, out)
+
+    with gr.Tab("Base64 API"):
+        b64 = gr.Textbox(label="Base64 Image")
+        out2 = gr.Textbox()
+        gr.Button("Predict").click(predict_base64, b64, out2)
+
+demo.queue()
+demo.launch()