add Image Editor GOT OCR Demo

README.md

@@ -1,6 +1,6 @@
 ---
-title: Tonic's Pixtral
-emoji: 📸🌬️
+title: Tonic's ImageEditor GOT OCR
+emoji: 📸🫴🏻👁
 colorFrom: blue
 colorTo: red
 sdk: gradio
@@ -10,7 +10,7 @@ pinned: true
 license: mit
 thumbnail: >-
   https://cdn-uploads.huggingface.co/production/uploads/62a3bb1cd0d8c2c2169f0b88/DlATYnzPl5cLHA_ua48Wl.png
-short_description: 'Pixtral by Mistral : Small-ish Vision Model With Prompting'
+short_description: 'Using Gradio Image Editor for GOT-OCR color ocr'
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -1,207 +1,172 @@
-import os
 import gradio as gr
-from vllm import LLM, SamplingParams
-from PIL import Image
-from io import BytesIO
-import base64
-import requests
-from huggingface_hub import login
 import torch
-import torch.nn.functional as F
-# import spaces
-import json
-from huggingface_hub import snapshot_download
-# import traceback
-
-login(os.environ.get("HUGGINGFACE_TOKEN"))
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:80"
-
-repo_id = "mistralai/Pixtral-12B-2409"
-max_tokens_per_img = 2048
-max_img_per_msg = 2
-
-title = "# **WIP / DEMO** 🙋🏻‍♂️Welcome to Tonic's Pixtral Model Demo"
-description = """
-### Join us : 
-🌟TeamTonic🌟 is always making cool demos! Join our active builder's 🛠️community 👻 [![Join us on Discord](https://img.shields.io/discord/1109943800132010065?label=Discord&logo=discord&style=flat-square)](https://discord.gg/qdfnvSPcqP) On 🤗Huggingface:[MultiTransformer](https://huggingface.co/MultiTransformer) On 🌐Github: [Tonic-AI](https://github.com/tonic-ai) & contribute to🌟 [Build Tonic](https://git.tonic-ai.com/contribute)🤗Big thanks to Yuvi Sharma and all the folks at huggingface for the community grant 🤗
-"""
-HUGGINGFACE_TOKEN = os.environ.get("HUGGINGFACE_TOKEN")
-model_path = snapshot_download(repo_id="mistralai/Pixtral-12B-2409", token=HUGGINGFACE_TOKEN)
-
-with open(f'{model_path}/params.json', 'r') as f:
-    params = json.load(f)
-
-with open(f'{model_path}/tekken.json', 'r') as f:
-    tokenizer_config = json.load(f)
-
-model_name = "mistralai/Pixtral-12B-2409"
-
-llm = LLM(
-    model=model_name,
-    tokenizer_mode="mistral",
-    max_num_batched_tokens=max_img_per_msg * max_tokens_per_img,
-    dtype="float16"
-)
-
-def clear_cuda_cache():
-    torch.cuda.empty_cache()
-
-def encode_image(image: Image.Image, image_format="PNG") -> str:
-    im_file = BytesIO()
-    image.save(im_file, format=image_format)
-    im_bytes = im_file.getvalue()
-    im_64 = base64.b64encode(im_bytes).decode("utf-8")
-    return im_64
-
-def infer(image_url, prompt, temperature, max_tokens, progress=gr.Progress(track_tqdm=True)):
-    if llm is None:
-        return "Error: LLM initialization failed. Please try again later."
+from transformers import AutoModel, AutoTokenizer, AutoConfig
+import os
+import base64
+import io
+from PIL import Image
+import numpy as np
+import uuid
+import cv2
+import re
+from globe import title, description, modelinfor, joinus, howto
+
+model_name = 'ucaslcl/GOT-OCR2_0'
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModel.from_pretrained(model_name, trust_remote_code=True, low_cpu_mem_usage=True, device_map='cuda', use_safetensors=True, pad_token_id=tokenizer.eos_token_id)
+model = model.eval().cuda()
+model.config.pad_token_id = tokenizer.eos_token_id
+
+UPLOAD_FOLDER = "./uploads"
+RESULTS_FOLDER = "./results"
+
+for folder in [UPLOAD_FOLDER, RESULTS_FOLDER]:
+    if not os.path.exists(folder):
+        os.makedirs(folder)
+
+def image_to_base64(image):
+    buffered = io.BytesIO()
+    image.save(buffered, format="PNG")
+    return base64.b64encode(buffered.getvalue()).decode()
+
+def process_image(image, ocr_type, ocr_box=None, ocr_color=None):
+    unique_id = str(uuid.uuid4())
+    image_path = os.path.join(UPLOAD_FOLDER, f"{unique_id}.png")
+    result_path = os.path.join(RESULTS_FOLDER, f"{unique_id}.html")
     
     try:
-        sampling_params = SamplingParams(max_tokens=max_tokens, temperature=temperature)
+        if isinstance(image, dict):
+            composite_image = image.get("composite")
+            if composite_image is not None:
+                if isinstance(composite_image, np.ndarray):
+                    cv2.imwrite(image_path, cv2.cvtColor(composite_image, cv2.COLOR_RGB2BGR))
+                elif isinstance(composite_image, Image.Image):
+                    composite_image.save(image_path)
+                else:
+                    return "Error: Unsupported image format from ImageEditor", None
+            else:
+                return "Error: No composite image found in ImageEditor output", None
+        else:
+            return "Error: Unsupported image format", None
+
+        if ocr_color:
+            res = model.chat(tokenizer, image_path, ocr_type=ocr_type, ocr_color=ocr_color, render=True, save_render_file=result_path)
+        else:
+            res = model.chat(tokenizer, image_path, ocr_type=ocr_type, ocr_box=ocr_box, render=True, save_render_file=result_path)
         
-        image = Image.open(BytesIO(requests.get(image_url).content))
-        image = image.resize((3844, 2408))
-        new_image_url = f"data:image/png;base64,{encode_image(image, image_format='PNG')}"
-
-        messages = [
-            {
-                "role": "user",
-                "content": [{"type": "text", "text": prompt}, {"type": "image_url", "image_url": {"url": new_image_url}}]
-            },
-        ]
-
-        outputs = llm.chat(messages, sampling_params=sampling_params)
-        clear_cuda_cache()
-        return outputs[0].outputs[0].text
+        if os.path.exists(result_path):
+            with open(result_path, 'r') as f:
+                html_content = f.read()
+            return res, html_content
+        else:
+            return res, None
     except Exception as e:
-        clear_cuda_cache()
-        return f"Error during inference: {e}"
-
-def compare_images(image1_url, image2_url, prompt, temperature, max_tokens, progress=gr.Progress(track_tqdm=True)):
-    if llm is None:
-        return "Error: LLM initialization failed. Please try again later."
-    
-    try:
-        sampling_params = SamplingParams(max_tokens=max_tokens, temperature=temperature)
+        return f"Error: {str(e)}", None
+    finally:
+        if os.path.exists(image_path):
+            os.remove(image_path)
+
+def parse_latex_output(res):
+    lines = re.split(r'(\$\$.*?\$\$)', res, flags=re.DOTALL)
+    parsed_lines = []
+    in_latex = False
+    latex_buffer = []
+
+    for line in lines:
+        if line == '\n':
+            if in_latex:
+                latex_buffer.append(line)
+            else:
+                parsed_lines.append(line)
+            continue
+
+        line = line.strip()
         
-        image1 = Image.open(BytesIO(requests.get(image1_url).content))
-        image2 = Image.open(BytesIO(requests.get(image2_url).content))
-        image1 = image1.resize((3844, 2408))
-        image2 = image2.resize((3844, 2408))
-        new_image1_url = f"data:image/png;base64,{encode_image(image1, image_format='PNG')}"
-        new_image2_url = f"data:image/png;base64,{encode_image(image2, image_format='PNG')}"
-
-        messages = [
-            {
-                "role": "user",
-                "content": [
-                    {"type": "text", "text": prompt},
-                    {"type": "image_url", "image_url": {"url": new_image1_url}},
-                    {"type": "image_url", "image_url": {"url": new_image2_url}}
-                ]
-            },
-        ]
-
-        outputs = llm.chat(messages, sampling_params=sampling_params)
-        clear_cuda_cache()
-        return outputs[0].outputs[0].text
-    except Exception as e:
-        clear_cuda_cache()
-        return f"Error during image comparison: {e}"
-
-def calculate_image_similarity(image1_url, image2_url):
-    if llm is None:
-        return "Error: LLM initialization failed. Please try again later."
+        latex_patterns = [r'\{', r'\}', r'\[', r'\]', r'\\', r'\$', r'_', r'^', r'"']
+        contains_latex = any(re.search(pattern, line) for pattern in latex_patterns)
+        
+        if contains_latex:
+            if not in_latex:
+                in_latex = True
+                latex_buffer = ['$$']
+            latex_buffer.append(line)
+        else:
+            if in_latex:
+                latex_buffer.append('$$')
+                parsed_lines.extend(latex_buffer)
+                in_latex = False
+                latex_buffer = []
+            parsed_lines.append(line)
+
+    if in_latex:
+        latex_buffer.append('$$')
+        parsed_lines.extend(latex_buffer)
+
+    return '$$\\$$\n'.join(parsed_lines)
+
+def ocr_demo(image, ocr_type, ocr_color):
+    res, html_content = process_image(image, ocr_type, ocr_color=ocr_color)
     
-    try:
-        image1 = Image.open(BytesIO(requests.get(image1_url).content)).convert('RGB')
-        image2 = Image.open(BytesIO(requests.get(image2_url).content)).convert('RGB')
-        image1 = image1.resize((224, 224))  # Resize to match model input size
-        image2 = image2.resize((224, 224))
-
-        image1_tensor = torch.tensor(list(image1.getdata())).view(1, 3, 224, 224).float() / 255.0
-        image2_tensor = torch.tensor(list(image2.getdata())).view(1, 3, 224, 224).float() / 255.0
-
-        with torch.no_grad():
-            embedding1 = llm.model.vision_encoder([image1_tensor])
-            embedding2 = llm.model.vision_encoder([image2_tensor])
+    if isinstance(res, str) and res.startswith("Error:"):
+        return res, None
 
-        similarity = F.cosine_similarity(embedding1.mean(dim=0), embedding2.mean(dim=0), dim=0).item()
-        clear_cuda_cache()
-        return similarity
-    except Exception as e:
-        clear_cuda_cache()
-        return f"Error during image similarity calculation: {e}"
-
-with gr.Blocks() as demo:
-    gr.Markdown(title)
-    gr.Markdown("## How it works")
-    gr.Markdown("1. The image is processed by a Vision Encoder using 2D ROPE (Rotary Position Embedding).")
-    gr.Markdown("2. The encoder uses SiLU activation in its feed-forward layers.")
-    gr.Markdown("3. The encoded image is used for text generation or similarity comparison.")
-    gr.Markdown(
-        """
-        ## How to use
-        1. For Image-to-Text Generation:
-           - Enter the URL of an image
-           - Provide a prompt describing what you want to know about the image
-           - Adjust the temperature and max tokens
-           - Click "Generate" to get the model's response
-        2. For Image Comparison:
-           - Enter URLs for two images you want to compare
-           - Provide a prompt asking about the comparison
-           - Adjust the temperature and max tokens
-           - Click "Compare" to get the model's analysis
-        3. For Image Similarity:
-           - Enter URLs for two images you want to compare
-           - Click "Calculate Similarity" to get a similarity score between 0 and 1
-        """
-    )
-    gr.Markdown(description)
-    with gr.Tabs():
-        with gr.TabItem("Image-to-Text Generation"):
-            with gr.Row():
-                image_url = gr.Text(label="Image URL")
-                prompt = gr.Text(label="Prompt")
-            with gr.Row():
-                temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, label="Temperature")
-                max_tokens = gr.Number(value=4096, label="Max Tokens")
-            generate_button = gr.Button("Generate")
-            output = gr.Text(label="Generated Text")
-            
-            generate_button.click(infer, inputs=[image_url, prompt, temperature, max_tokens], outputs=output)
-        
-        with gr.TabItem("Image Comparison"):
-            with gr.Row():
-                image1_url = gr.Text(label="Image 1 URL")
-                image2_url = gr.Text(label="Image 2 URL")
-            comparison_prompt = gr.Text(label="Comparison Prompt")
+    res = res.replace("\\title", "\\title ")
+    formatted_res = parse_latex_output(res)
+    
+    if html_content:
+        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_{uuid.uuid4()}.html">Download Full Result</a>'
+        return formatted_res, f"{iframe}<br>{download_link}"
+    return formatted_res, None
+
+with gr.Blocks(theme=gr.themes.Base()) as demo:
+    with gr.Row():
+        gr.Markdown(title)
+    with gr.Row():
+        with gr.Column(scale=1):
+            with gr.Group():                    
+                gr.Markdown(description)
+        with gr.Column(scale=1):
+            with gr.Group():
+                gr.Markdown(modelinfor)
+                gr.Markdown(joinus)
+    with gr.Row():
+        with gr.Accordion("How to use 🫴🏻👁GOT OCR", open=True):
             with gr.Row():
-                comparison_temperature = gr.Slider(minimum=0.1, maximum=2.0, value=0.7, label="Temperature")
-                comparison_max_tokens = gr.Number(value=4096, label="Max Tokens")
-            compare_button = gr.Button("Compare")
-            comparison_output = gr.Text(label="Comparison Result")
-            
-            compare_button.click(compare_images, inputs=[image1_url, image2_url, comparison_prompt, comparison_temperature, comparison_max_tokens], outputs=comparison_output)
-
-        with gr.TabItem("Image Similarity"):
+                gr.Image("res/image/howto_1.png", label="Select the Following Parameters")
+                gr.Image("res/image/howto_2.png", label="Click on Paintbrush in the Image Editor")
+                gr.Image("res/image/howto_3.png", label="Select your Brush Color (Red)")
+                gr.Image("res/image/howto_4.png", label="Make a Box Around The Text")
             with gr.Row():
-                sim_image1_url = gr.Text(label="Image 1 URL")
-                sim_image2_url = gr.Text(label="Image 2 URL")
-            similarity_button = gr.Button("Calculate Similarity")
-            similarity_output = gr.Number(label="Similarity Score")
-
-            similarity_button.click(calculate_image_similarity, inputs=[sim_image1_url, sim_image2_url], outputs=similarity_output)
-    gr.Markdown("## Model Details")
-    gr.Markdown(f"- Model Dimension: {params['dim']}")
-    gr.Markdown(f"- Number of Layers: {params['n_layers']}")
-    gr.Markdown(f"- Number of Attention Heads: {params['n_heads']}")
-    gr.Markdown(f"- Vision Encoder Hidden Size: {params['vision_encoder']['hidden_size']}")
-    gr.Markdown(f"- Number of Vision Encoder Layers: {params['vision_encoder']['num_hidden_layers']}")
-    gr.Markdown(f"- Number of Vision Encoder Attention Heads: {params['vision_encoder']['num_attention_heads']}")
-    gr.Markdown(f"- Image Size: {params['vision_encoder']['image_size']}x{params['vision_encoder']['image_size']}")
-    gr.Markdown(f"- Patch Size: {params['vision_encoder']['patch_size']}x{params['vision_encoder']['patch_size']}")
+                with gr.Group():
+                    gr.Markdown(howto)
+    with gr.Row():
+        with gr.Column(scale=1):
+            image_editor = gr.ImageEditor(label="Image Editor", type="pil")
+            ocr_type_dropdown = gr.Dropdown(
+                choices=["ocr", "format"],
+                label="OCR Type",
+                value="ocr"
+            )
+            ocr_color_dropdown = gr.Dropdown(
+                choices=["red", "green", "blue"],
+                label="OCR Color",
+                value="red"
+            )
+            submit_button = gr.Button("Process")
+
+        with gr.Column(scale=1):
+            output_markdown = gr.Markdown(label="OCR Result")
+            output_html = gr.HTML(label="Rendered Result")
+
+    submit_button.click(
+        ocr_demo,
+        inputs=[image_editor, ocr_type_dropdown, ocr_color_dropdown],
+        outputs=[output_markdown, output_html]
+    )
 
 if __name__ == "__main__":
     demo.launch()

globe.py

@@ -0,0 +1,68 @@
+
+title = """# 🙋🏻‍♂️Welcome to Tonic's ImageEditor📸🫴🏻👁GOT-OCR Demo
+---
+"""
+
+description = """
+The **🫴🏻👁GOT-OCR model** is a cutting-edge OCR system with **580M parameters**, designed to process a wide range of "characters." Equipped with a **high-compression encoder** and a **long-context decoder**, it excels in both scene and document-style images. The model supports **multi-page** and **dynamic resolution OCR**, enhancing its versatility.
+
+### Key Features
+
+- **Plain Text OCR**: Extracts text from images.
+- **Formatted Text OCR**: Retains the original formatting, including tables and formulas.
+- **Fine-grained OCR**: Offers box-based and color-based OCR for precision in specific regions.
+- **Multi-crop OCR**: Handles multiple cropped sections within an image.
+
+## Supported Content Types
+
+- Plain text
+- Math/molecular formulas
+- Tables and charts
+- Sheet music
+- Geometric shapes
+
+"""
+joinus = """
+## Join us :
+
+🌟TeamTonic🌟 is always making cool demos! Join our active builder's 🛠️community 👻 [![Join us on Discord](https://img.shields.io/discord/1109943800132010065?label=Discord&logo=discord&style=flat-square)](https://discord.gg/qdfnvSPcqP) On 🤗Huggingface:[MultiTransformer](https://huggingface.co/MultiTransformer) On 🌐Github: [Tonic-AI](https://github.com/tonic-ai) & contribute to🌟 [Build Tonic](https://git.tonic-ai.com/contribute)🤗Big thanks to Yuvi Sharma and all the folks at huggingface for the community grant 🤗
+"""
+modelinfor = """
+## How to Use
+
+1. Select a task from the dropdown menu.
+2. Upload an image.
+3. Use Paintbrush to draw a circle around the text you want to OCR.
+4. Click **Process** to view the results.
+
+## Model Information
+
+- **Model Name**: GOT-OCR 2.0
+- **Hugging Face Repository**: [ucaslcl/GOT-OCR2_0](https://huggingface.co/ucaslcl/GOT-OCR2_0)
+- **Environment**: CUDA 11.8 + PyTorch 2.0.1
+
+"""
+
+tasks = [
+    "Plain Text OCR",
+    "Format Text OCR",
+    "Fine-grained OCR (Box)",
+    "Fine-grained OCR (Color)",
+    "Multi-crop OCR",
+    "Render Formatted OCR"
+]
+
+ocr_types = ["ocr", "format"]
+ocr_colors = ["red", "green", "blue"]
+
+howto =         """
+                ## To use Fine-grained OCR (Color):
+                1. Click on 'Fine-grained OCR (Color)' in the task dropdown.
+                2. Set 'OCR Type' to 'ocr'.
+                3. This will display the image editor.
+                4. Upload an image to the editor.
+                5. Use the drawing tools to draw a circle around the text you want to OCR.
+                6. Select the color that matches your circle in the 'OCR Color' dropdown.
+                7. Click 'Process Edited Image' to run the OCR on the selected area.
+                """
+           

requirements.txt

@@ -1,7 +1,14 @@
-torch>=1.9.0
-safetensors>=0.3.1
-Pillow>=9.0.0
-numpy>=1.21.0
-mistral_common
+torch==2.0.1
+torchvision==0.15.2
+transformers==4.37.2
+megfile==3.1.2
+tiktoken
+verovio
+opencv-python
+cairosvg
+accelerate
+numpy==1.26.4
 loadimg
-vllm==0.6.1
+pillow
+markdown
+shutils