Create app.py

app.py

@@ -0,0 +1,84 @@
+from flask import Flask, request, jsonify
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+from PIL import Image
+import torch
+import io
+import requests
+
+# ==========================================
+# Initialize Flask App
+# ==========================================
+app = Flask(__name__)
+
+# ==========================================
+# Load Model and Processor
+# ==========================================
+MODEL_NAME = "anuashok/ocr-captcha-v3"
+print("🚀 Loading model...")
+
+processor = TrOCRProcessor.from_pretrained(MODEL_NAME)
+model = VisionEncoderDecoderModel.from_pretrained(MODEL_NAME)
+
+# Use GPU if available
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
+model.eval()
+
+print("✅ Model loaded successfully on", device)
+
+# ==========================================
+# API Endpoint (Root)
+# ==========================================
+@app.route("/", methods=["GET", "POST"])
+def root():
+    """
+    POST / -> Send image file
+    GET /?url=image_url -> Send image URL
+    Returns: JSON with recognized text
+    """
+    try:
+        # ---------------------------
+        # Handle GET with ?url=
+        # ---------------------------
+        if request.method == "GET":
+            image_url = request.args.get("url")
+            if not image_url:
+                return jsonify({
+                    "message": "OCR Captcha API running. Use POST (file) or GET ?url=image_url."
+                })
+            try:
+                img_data = requests.get(image_url, timeout=10).content
+                image = Image.open(io.BytesIO(img_data)).convert("RGB")
+            except Exception as e:
+                return jsonify({"error": f"Failed to fetch image from URL: {str(e)}"}), 400
+
+        # ---------------------------
+        # Handle POST with file upload
+        # ---------------------------
+        elif request.method == "POST":
+            if "file" not in request.files:
+                return jsonify({"error": "No file provided"}), 400
+
+            file = request.files["file"]
+            image = Image.open(io.BytesIO(file.read())).convert("RGB")
+
+        # ---------------------------
+        # Process Image with Model
+        # ---------------------------
+        pixel_values = processor(images=image, return_tensors="pt").pixel_values.to(device)
+        with torch.no_grad():
+            generated_ids = model.generate(pixel_values)
+            text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+        return jsonify({"text": text})
+
+    except Exception as e:
+        return jsonify({"error": str(e)}), 500
+
+
+# ==========================================
+# Run Server
+# ==========================================
+if __name__ == "__main__":
+    # Host on 0.0.0.0 so Hugging Face Space can reach it
+    app.run(host="0.0.0.0", port=7860)