add html file handling

app.py

@@ -13,6 +13,8 @@ from globe import title, description, modelinfor, joinus
 import uuid
 import tempfile
 import time
+import shutil
+
 
 model_name = 'ucaslcl/GOT-OCR2_0'
 
@@ -27,35 +29,51 @@ def image_to_base64(image):
     image.save(buffered, format="PNG")
     return base64.b64encode(buffered.getvalue()).decode()
 
-results_folder = Path('./results')
-results_folder.mkdir(parents=True, exist_ok=True)
+UPLOAD_FOLDER = "./uploads"
+RESULTS_FOLDER = "./results"
+
+for folder in [UPLOAD_FOLDER, RESULTS_FOLDER]:
+    if not os.path.exists(folder):
+        os.makedirs(folder)
 
-@spaces.GPU
+@spaces.GPU()
 def process_image(image, task, ocr_type=None, ocr_box=None, ocr_color=None):
+    if image is None:
+        return "Error: No image provided", None, None
+
     unique_id = str(uuid.uuid4())
-    temp_html_path = results_folder / f"{unique_id}.html"
+    image_path = os.path.join(UPLOAD_FOLDER, f"{unique_id}.png")
+    result_path = os.path.join(RESULTS_FOLDER, f"{unique_id}.html")
     
-    if task == "Plain Text OCR":
-        res = model.chat(tokenizer, image, ocr_type='ocr')
-        return res, None, unique_id
-    else:
-        if task == "Format Text OCR":
-            res = model.chat(tokenizer, image, ocr_type='format', render=True, save_render_file=str(temp_html_path))
-        elif task == "Fine-grained OCR (Box)":
-            res = model.chat(tokenizer, image, ocr_type=ocr_type, ocr_box=ocr_box, render=True, save_render_file=str(temp_html_path))
-        elif task == "Fine-grained OCR (Color)":
-            res = model.chat(tokenizer, image, ocr_type=ocr_type, ocr_color=ocr_color, render=True, save_render_file=str(temp_html_path))
-        elif task == "Multi-crop OCR":
-            res = model.chat_crop(tokenizer, image, ocr_type='format', render=True, save_render_file=str(temp_html_path))
-        elif task == "Render Formatted OCR":
-            res = model.chat(tokenizer, image, ocr_type='format', render=True, save_render_file=str(temp_html_path))
-        
-        if temp_html_path.exists():
-            with open(temp_html_path, 'r') as f:
-                html_content = f.read()
-            return res, html_content, unique_id
-        else:
+    shutil.copy(image, image_path)
+    
+    try:
+        if task == "Plain Text OCR":
+            res = model.chat(tokenizer, image_path, ocr_type='ocr')
             return res, None, unique_id
+        else:
+            if task == "Format Text OCR":
+                res = model.chat(tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
+            elif task == "Fine-grained OCR (Box)":
+                res = model.chat(tokenizer, image_path, ocr_type=ocr_type, ocr_box=ocr_box, render=True, save_render_file=result_path)
+            elif task == "Fine-grained OCR (Color)":
+                res = model.chat(tokenizer, image_path, ocr_type=ocr_type, ocr_color=ocr_color, render=True, save_render_file=result_path)
+            elif task == "Multi-crop OCR":
+                res = model.chat_crop(tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
+            elif task == "Render Formatted OCR":
+                res = model.chat(tokenizer, image_path, ocr_type='format', render=True, save_render_file=result_path)
+            
+            if os.path.exists(result_path):
+                with open(result_path, 'r') as f:
+                    html_content = f.read()
+                return res, html_content, unique_id
+            else:
+                return res, None, unique_id
+    except Exception as e:
+        return f"Error: {str(e)}", None, None
+    finally:
+        if os.path.exists(image_path):
+            os.remove(image_path)
     
 def update_inputs(task):
     if task in ["Plain Text OCR", "Format Text OCR", "Multi-crop OCR", "Render Formatted OCR"]:
@@ -72,22 +90,29 @@ def update_inputs(task):
             gr.update(visible=False),
             gr.update(visible=True, choices=["red", "green", "blue"]),
         ]
+
 def ocr_demo(image, task, ocr_type, ocr_box, ocr_color):
     res, html_content, unique_id = process_image(image, task, ocr_type, ocr_box, ocr_color)
     
+    if res.startswith("Error:"):
+        return res, None
+    
     res = f"$$ {res} $$"
     
     if html_content:
-        iframe = f'<iframe srcdoc="{html_content}" width="100%" height="600px"></iframe>'
-        link = f'<a href="file={results_folder / f"{unique_id}.html"}" target="_blank">View Full Result</a>'
-        return res, f"{link}<br>{iframe}"
+        encoded_html = base64.b64encode(html_content.encode('utf-8')).decode('utf-8')
+        iframe_src = f"data:text/html;base64,{encoded_html}"
+        iframe = f'<iframe src="{iframe_src}" width="100%" height="600px"></iframe>'
+        download_link = f'<a href="data:text/html;base64,{encoded_html}" download="result_{unique_id}.html">Download Full Result</a>'
+        return res, f"{download_link}<br>{iframe}"
     return res, None
 
 def cleanup_old_files():
     current_time = time.time()
-    for file_path in results_folder.glob('*.html'):
-        if current_time - file_path.stat().st_mtime > 3600:  # 1 hour
-            file_path.unlink()
+    for folder in [UPLOAD_FOLDER, RESULTS_FOLDER]:
+        for file_path in Path(folder).glob('*'):
+            if current_time - file_path.stat().st_mtime > 3600:  # 1 hour
+                file_path.unlink()
 
 with gr.Blocks() as demo:
     gr.Markdown(title)

requirements.txt

@@ -10,4 +10,5 @@
  numpy==1.26.4
  loadimg
  pillow
- markdown
+ markdown
+ py-shutils