remove loadimg

app.py

@@ -4,7 +4,7 @@ from transformers import AutoModel, AutoTokenizer, AutoConfig
 import os
 import base64
 import spaces
-from loadimg import load_img
+import io
 from PIL import Image
 import numpy as np
 
@@ -47,34 +47,32 @@ model = AutoModel.from_pretrained('ucaslcl/GOT-OCR2_0', trust_remote_code=True,
 model = model.eval().cuda()
 model.config.pad_token_id = tokenizer.eos_token_id
 
-def load_image(image_file):
-    if isinstance(image_file, str):
-        if image_file.startswith('http') or image_file.startswith('https'):
-            return Image.open(requests.get(image_file, stream=True).raw).convert('RGB')
-        else:
-            return Image.open(image_file).convert('RGB')
-    else:
-        return image_file.convert('RGB')
+def image_to_base64(image):
+    buffered = io.BytesIO()
+    image.save(buffered, format="PNG")
+    return base64.b64encode(buffered.getvalue()).decode()
 
 @spaces.GPU
 def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None, render=False):
     try:
-        img = load_image(image)
-        img_path = "/tmp/temp_image.png"
-        img.save(img_path)
+        if image is None:
+            return "No image provided", None
+
+        # Convert image to base64
+        image_base64 = image_to_base64(image)
 
         if task == "Plain Text OCR":
-            res = model.chat(tokenizer, img_path, ocr_type='ocr')
+            res = model.chat(tokenizer, image_base64, ocr_type='ocr')
         elif task == "Format Text OCR":
-            res = model.chat(tokenizer, img_path, ocr_type='format')
+            res = model.chat(tokenizer, image_base64, ocr_type='format')
         elif task == "Fine-grained OCR (Box)":
-            res = model.chat(tokenizer, img_path, ocr_type=ocr_type, ocr_box=ocr_box)
+            res = model.chat(tokenizer, image_base64, ocr_type=ocr_type, ocr_box=ocr_box)
         elif task == "Fine-grained OCR (Color)":
-            res = model.chat(tokenizer, img_path, ocr_type=ocr_type, ocr_color=ocr_color)
+            res = model.chat(tokenizer, image_base64, ocr_type=ocr_type, ocr_color=ocr_color)
         elif task == "Multi-crop OCR":
-            res = model.chat_crop(tokenizer, image_file=img_path)
+            res = model.chat_crop(tokenizer, image_file=image_base64)
         elif task == "Render Formatted OCR":
-            res = model.chat(tokenizer, img_path, ocr_type='format', render=True, save_render_file='./results/demo.html')
+            res = model.chat(tokenizer, image_base64, ocr_type='format', render=True, save_render_file='./results/demo.html')
             with open('./results/demo.html', 'r') as f:
                 html_content = f.read()
             return res, html_content
@@ -147,7 +145,7 @@ with gr.Blocks() as demo:
             )
             render_checkbox = gr.Checkbox(
                 label="Render Result",
-                visible=False
+                visible=False   
             )
             submit_button = gr.Button("Process")