remove loadimg

app.py

@@ -6,6 +6,7 @@ import base64
 import spaces
 from loadimg import load_img
 from PIL import Image
+import numpy as np
 
 title = """# 🙋🏻‍♂️Welcome to Tonic's🫴🏻📸GOT-OCR"""
 description = """"
@@ -46,38 +47,45 @@ 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')
+
 @spaces.GPU
 def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None, render=False):
-    if isinstance(image, str):
-        img_path = image
-        img = Image.open(img_path)
-    else:
-        img = Image.open(image)
+    try:
+        img = load_image(image)
         img_path = "/tmp/temp_image.png"
         img.save(img_path)
 
-    if task == "Plain Text OCR":
-        res = model.chat(tokenizer, img, ocr_type='ocr')
-    elif task == "Format Text OCR":
-        res = model.chat(tokenizer, img, ocr_type='format')
-    elif task == "Fine-grained OCR (Box)":
-        res = model.chat(tokenizer, img, ocr_type=ocr_type, ocr_box=ocr_box)
-    elif task == "Fine-grained OCR (Color)":
-        res = model.chat(tokenizer, img, ocr_type=ocr_type, ocr_color=ocr_color)
-    elif task == "Multi-crop OCR":
-        res = model.chat_crop(tokenizer, image_file=img_path)
-    elif task == "Render Formatted OCR":
-        res = model.chat(tokenizer, img, 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
-    
-    # Clean up 
-    if img_path.startswith("/tmp/"):
+        if task == "Plain Text OCR":
+            res = model.chat(tokenizer, img_path, ocr_type='ocr')
+        elif task == "Format Text OCR":
+            res = model.chat(tokenizer, img_path, ocr_type='format')
+        elif task == "Fine-grained OCR (Box)":
+            res = model.chat(tokenizer, img_path, 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)
+        elif task == "Multi-crop OCR":
+            res = model.chat_crop(tokenizer, image_file=img_path)
+        elif task == "Render Formatted OCR":
+            res = model.chat(tokenizer, img_path, 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
+        
+        # Clean up 
         os.remove(img_path)
+        
+        return res, None
+    except Exception as e:
+        return str(e), None
     
-    return res, None
-
 def update_inputs(task):
     if task == "Plain Text OCR" or task == "Format Text OCR" or task == "Multi-crop OCR":
         return [gr.update(visible=False)] * 4

requirements.txt

@@ -8,4 +8,5 @@
  cairosvg
  accelerate
  numpy 
- loadimg
+ loadimg
+ pillow