Update app.py

app.py

@@ -1,23 +1,38 @@
-from flask import Flask, request, jsonify
-from transformers import AutoTokenizer, AutoModel
+import gradio as gr
+from PIL import Image
+# Assuming 'model' and 'tokenizer' are defined elsewhere in your code
+# from your_model_file import model, tokenizer
 
-app = Flask(__name__)
+def load_image(image_file):
+    """Load and preprocess the image."""
+    if isinstance(image_file, Image.Image):  # Check if the input is an Image object
+        return image_file.convert("RGB")  # Convert to RGB if necessary
+    elif isinstance(image_file, str) and (image_file.startswith('http') or image_file.startswith('https')):
+        # Handle URL case (you can use an external library to fetch the image if needed)
+        return Image.open(requests.get(image_file, stream=True).raw).convert("RGB")
+    else:
+        # Handle file path case
+        return Image.open(image_file).convert("RGB")
 
-# Load model and tokenizer
-try:
-    tokenizer = AutoTokenizer.from_pretrained('stepfun-ai/GOT-OCR2_0', revision='cf6b7386bc89a54f09785612ba74cb12de6fa17c', trust_remote_code=True)
-    model = AutoModel.from_pretrained('stepfun-ai/GOT-OCR2_0', revision='cf6b7386bc89a54f09785612ba74cb12de6fa17c', trust_remote_code=True)
-except Exception as e:
-    print(f"Error loading model and tokenizer: {e}")
+def perform_ocr(image):
+    """Perform OCR on the uploaded image."""
+    try:
+        # Load and preprocess the image
+        processed_image = load_image(image)
+        # Use the model for OCR
+        res = model.chat(tokenizer, processed_image, ocr_type='ocr')
+        return res
+    except Exception as e:
+        return str(e)  # Return the error message
 
-@app.route('/predict', methods=['POST'])
-def predict():
-    # Assuming you send image data in the request
-    data = request.json
-    # Add your model inference logic here
-    # e.g., model.forward(data)
+# Gradio interface setup
+iface = gr.Interface(
+    fn=perform_ocr,
+    inputs=gr.Image(type="pil"),  # Ensure Gradio accepts images as PIL images
+    outputs="text",
+    title="OCR Application",
+    description="Upload an image to perform Optical Character Recognition (OCR)."
+)
 
-    return jsonify({"message": "Prediction made successfully!"})
-
-if __name__ == "__main__":
-    app.run(host='0.0.0.0', port=5000)  # Adjust port if necessary
+if _name_ == "_main_":
+    iface.launch()