add image editor

app.py

@@ -14,6 +14,7 @@ import uuid
 import tempfile
 import time
 import shutil
+import cv2
 
 
 model_name = 'ucaslcl/GOT-OCR2_0'
@@ -38,6 +39,9 @@ def image_to_base64(image):
 
 
 
+@spaces.GPU()
+import cv2
+
 @spaces.GPU()
 def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None):
     if image is None:
@@ -52,15 +56,13 @@ def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None):
             composite_image = image.get("composite")
             if composite_image is not None:
                 if isinstance(composite_image, np.ndarray):
-                    Image.fromarray(composite_image).save(image_path)
-                elif isinstance(composite_image, str):
-                    shutil.copy(composite_image, image_path)
+                    cv2.imwrite(image_path, cv2.cvtColor(composite_image, cv2.COLOR_RGB2BGR))
                 else:
                     return "Error: Unsupported image format from ImageEditor", None, None
             else:
                 return "Error: No composite image found in ImageEditor output", None, None
         elif isinstance(image, np.ndarray):
-            Image.fromarray(image).save(image_path)
+            cv2.imwrite(image_path, cv2.cvtColor(image, cv2.COLOR_RGB2BGR))
         elif isinstance(image, str):
             shutil.copy(image, image_path)
         else:
@@ -92,6 +94,7 @@ def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None):
     finally:
         if os.path.exists(image_path):
             os.remove(image_path)
+            
 def update_image_input(task):
     if task == "Fine-grained OCR (Color)":
         return gr.update(visible=False), gr.update(visible=True), gr.update(visible=True)