Update app.py

app.py

@@ -1,38 +1,36 @@
 import gradio as gr
+from transformers import AutoTokenizer, AutoModelForImageClassification
 from PIL import Image
-# Assuming 'model' and 'tokenizer' are defined elsewhere in your code
-# from your_model_file import model, tokenizer
+import torch
 
-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 the model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained('stepfun-ai/GOT-OCR2_0')
+model = AutoModelForImageClassification.from_pretrained('stepfun-ai/GOT-OCR2_0')
 
 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
+    # Ensure the image is in the right format
+    if isinstance(image, Image.Image):
+        image = image.convert("RGB")
+    else:
+        raise ValueError("Input must be a PIL Image")
+
+    # Use the model to perform OCR
+    inputs = tokenizer(image, return_tensors="pt")
+    with torch.no_grad():
+        outputs = model(**inputs)
+    
+    # Get the predictions
+    predictions = outputs.logits.argmax(dim=1).item()
+    return predictions
 
-# Gradio interface setup
+# Create the Gradio interface
 iface = gr.Interface(
     fn=perform_ocr,
-    inputs=gr.Image(type="pil"),  # Ensure Gradio accepts images as PIL images
+    inputs=gr.inputs.Image(type="pil"),
     outputs="text",
-    title="OCR Application",
-    description="Upload an image to perform Optical Character Recognition (OCR)."
+    title="OCR with GOT-OCR2.0",
+    description="Upload an image for Optical Character Recognition."
 )
 
-if __name__ == "__main__":
-    iface.launch()
+# Launch the interface
+iface.launch()