update app [Bounding Boxes]

app.py

@@ -15,7 +15,7 @@ import tempfile
 import gradio as gr
 import requests
 import torch
-from PIL import Image
+from PIL import Image, ImageDraw
 import fitz
 import numpy as np
 
@@ -130,7 +130,7 @@ def generate_and_preview_pdf(image: Image.Image, text_content: str, font_size: i
 def process_document_stream(
     image: Image.Image, 
     prompt_input: str,
-    image_scale_factor: float, # New parameter for image scaling
+    image_scale_factor: float,
     max_new_tokens: int,
     temperature: float,
     top_p: float,
@@ -138,7 +138,7 @@ def process_document_stream(
     repetition_penalty: float
 ):
     """
-    Main function that handles model inference using tencent/POINTS-Reader.
+    Main function that handles model inference for general OCR.
     """
     if image is None:
         yield "Please upload an image.", ""
@@ -147,42 +147,29 @@ def process_document_stream(
         yield "Please enter a prompt.", ""
         return
 
-    # --- IMPLEMENTATION: Image Scaling based on user input ---
     if image_scale_factor > 1.0:
         try:
             original_width, original_height = image.size
             new_width = int(original_width * image_scale_factor)
             new_height = int(original_height * image_scale_factor)
             print(f"Scaling image from {image.size} to ({new_width}, {new_height}) with factor {image_scale_factor}.")
-            # Use a high-quality resampling filter for better results
             image = image.resize((new_width, new_height), Image.Resampling.LANCZOS)
         except Exception as e:
             print(f"Error during image scaling: {e}")
-            # Continue with the original image if scaling fails
             pass
-    # --- END IMPLEMENTATION ---
 
     temp_image_path = None
     try:
-        # --- FIX: Save the PIL Image to a temporary file ---
-        # The model expects a file path, not a PIL object.
         temp_dir = tempfile.gettempdir()
         temp_image_path = os.path.join(temp_dir, f"temp_image_{uuid.uuid4()}.png")
         image.save(temp_image_path)
         
-        # Prepare content for the model using the temporary file path
         content = [
             dict(type='image', image=temp_image_path),
             dict(type='text', text=prompt_input)
         ]
-        messages = [
-            {
-                'role': 'user',
-                'content': content
-            }
-        ]
+        messages = [{'role': 'user', 'content': content}]
         
-        # Prepare generation configuration from UI inputs
         generation_config = {
             'max_new_tokens': max_new_tokens,
             'repetition_penalty': repetition_penalty,
@@ -192,21 +179,78 @@ def process_document_stream(
             'do_sample': True if temperature > 0 else False
         }
 
-        # Run inference
-        response = model.chat(
-            messages,
-            tokenizer,
-            image_processor,
-            generation_config
-        )
-        # Yield the full response at once
+        response = model.chat(messages, tokenizer, image_processor, generation_config)
         yield response, response
 
     except Exception as e:
         traceback.print_exc()
         yield f"An error occurred during processing: {str(e)}", ""
     finally:
-        # --- Clean up the temporary image file ---
+        if temp_image_path and os.path.exists(temp_image_path):
+            os.remove(temp_image_path)
+
+# --- Bounding Box Extraction Logic ---
+@spaces.GPU
+def extract_text_with_coordinates(image: Image.Image):
+    """
+    Runs the model with a specific prompt to get OCR and bounding boxes,
+    then processes the output to create a visualization.
+    """
+    if image is None:
+        raise gr.Error("Please upload an image first in the main tab.")
+
+    prompt = "Please perform OCR on the image and provide the bounding box for each recognized text line. The format should be 'text<box>x1, y1, x2, y2</box>'."
+    temp_image_path = None
+    try:
+        temp_dir = tempfile.gettempdir()
+        temp_image_path = os.path.join(temp_dir, f"temp_image_{uuid.uuid4()}.png")
+        image.save(temp_image_path)
+        
+        content = [dict(type='image', image=temp_image_path), dict(type='text', text=prompt)]
+        messages = [{'role': 'user', 'content': content}]
+        generation_config = {'max_new_tokens': 4096} 
+
+        response = model.chat(messages, tokenizer, image_processor, generation_config)
+
+        original_width, original_height = image.size
+        
+        # Regex to find coordinates inside <box> tags
+        pattern_coords = r"<box>(\d+,\s*\d+,\s*\d+,\s*\d+)</box>"
+        # Regex to split the string by the full box tag to isolate text
+        pattern_splitter = r"<box>\d+,\s*\d+,\s*\d+,\s*\d+</box>"
+        
+        bboxs_raw = re.findall(pattern_coords, response)
+        lines = [line.strip() for line in re.split(pattern_splitter, response) if line.strip()]
+        
+        num_items = min(len(lines), len(bboxs_raw))
+        vis_image = image.copy()
+        draw = ImageDraw.Draw(vis_image)
+        output_text = ""
+
+        for i in range(num_items):
+            line_text = lines[i]
+            box_coords = [int(c.strip()) for c in bboxs_raw[i].split(',')]
+            
+            if len(box_coords) == 4:
+                x0, y0, x1, y1 = box_coords
+                
+                # Scale coordinates from the model's 1000px basis to the original image size
+                x0_s = int(x0 * original_width / 1000)
+                y0_s = int(y0 * original_height / 1000)
+                x1_s = int(x1 * original_width / 1000)
+                y1_s = int(y1 * original_height / 1000)
+                
+                draw.rectangle([x0_s, y0_s, x1_s, y1_s], outline="red", width=2)
+                
+                # Format output as a polygon (quadrilateral) and the extracted text
+                output_text += f"{x0_s},{y0_s},{x1_s},{y0_s},{x1_s},{y1_s},{x0_s},{y1_s},{line_text}\n"
+
+        return output_text.strip(), vis_image
+
+    except Exception as e:
+        traceback.print_exc()
+        return f"An error occurred: {str(e)}", None
+    finally:
         if temp_image_path and os.path.exists(temp_image_path):
             os.remove(temp_image_path)
 
@@ -233,29 +277,12 @@ def create_gradio_interface():
         with gr.Row():
             # Left Column (Inputs)
             with gr.Column(scale=1):
-                gr.Textbox(
-                    label="Model in Use ⚡",
-                    value="tencent/POINTS-Reader",
-                    interactive=False
-                )
-                prompt_input = gr.Textbox(
-                    label="Query Input",
-                    placeholder="✦︎ Enter the prompt",
-                    value="Perform OCR on the image precisely.",
-                )
+                gr.Textbox(label="Model in Use ⚡", value="tencent/POINTS-Reader", interactive=False)
+                prompt_input = gr.Textbox(label="Query Input", placeholder="✦︎ Enter the prompt", value="Perform OCR on the image precisely.")
                 image_input = gr.Image(label="Upload Image", type="pil", sources=['upload'])
                 
                 with gr.Accordion("Advanced Settings", open=False):
-                    # --- NEW UI ELEMENT: Image Scaling Slider ---
-                    image_scale_factor = gr.Slider(
-                        minimum=1.0, 
-                        maximum=3.0, 
-                        value=1.0, 
-                        step=0.1, 
-                        label="Image Upscale Factor",
-                        info="Increases image size before processing. Can improve OCR on small text. Default: 1.0 (no change)."
-                    )
-                    # --- END NEW UI ELEMENT ---
+                    image_scale_factor = gr.Slider(minimum=1.0, maximum=3.0, value=1.0, step=0.1, label="Image Upscale Factor", info="Increases image size before processing. Can improve OCR on small text.")
                     max_new_tokens = gr.Slider(minimum=512, maximum=8192, value=2048, step=256, label="Max New Tokens")
                     temperature = gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, step=0.05, value=0.7)
                     top_p = gr.Slider(label="Top-p (nucleus sampling)", minimum=0.05, maximum=1.0, step=0.05, value=0.8)
@@ -277,19 +304,21 @@ def create_gradio_interface():
                     with gr.Tab("📝 Extracted Content"):
                         raw_output_stream = gr.Textbox(label="Raw Model Output (max T ≤ 120s)", interactive=False, lines=15, show_copy_button=True)
                         with gr.Row():
-                            examples = gr.Examples(
-                                examples=["examples/1.jpeg", 
-                                          "examples/2.jpeg", 
-                                          "examples/3.jpeg",
-                                          "examples/4.jpeg", 
-                                          "examples/5.jpeg"],
-                                inputs=image_input, label="Examples"
-                            )
+                            examples = gr.Examples(examples=["examples/1.jpeg", "examples/2.jpeg", "examples/3.jpeg", "examples/4.jpeg", "examples/5.jpeg"], inputs=image_input, label="Examples")
                         gr.Markdown("[Report-Bug💻](https://huggingface.co/spaces/prithivMLmods/POINTS-Reader-OCR/discussions) | [prithivMLmods🤗](https://huggingface.co/prithivMLmods)")
                     
                     with gr.Tab("📰 README.md"):
                         with gr.Accordion("(Result.md)", open=True): 
                             markdown_output = gr.Markdown()
+                    
+                    with gr.Tab("Bounding Boxes"):
+                        gr.Markdown("Click the button to extract text and visualize its location on the image. This uses a specialized prompt to get coordinates from the model.")
+                        with gr.Row():
+                            with gr.Column(scale=1):
+                                ocr_button = gr.Button("🔍 Extract Text with Coordinates", variant="primary")
+                                ocr_text = gr.Textbox(label="Extracted Text with Coordinates", info="Format: x1,y1,x2,y2,x3,y3,x4,y4,text", lines=15, show_copy_button=True)
+                            with gr.Column(scale=1):
+                                ocr_vis = gr.Image(label="Visualization (Red boxes show detected text)")
 
                     with gr.Tab("📋 PDF Preview"):
                         generate_pdf_btn = gr.Button("📄 Generate PDF & Render", variant="primary")
@@ -298,15 +327,21 @@ def create_gradio_interface():
 
         # Event Handlers
         def clear_all_outputs():
-            return None, "", "Raw output will appear here.", "", None, None
+            # Clear all input and output fields across all tabs
+            return None, "", "Raw output will appear here.", "", None, None, "", None
 
         process_btn.click(
             fn=process_document_stream,
-            # --- UPDATE: Add the new slider to the inputs list ---
             inputs=[image_input, prompt_input, image_scale_factor, max_new_tokens, temperature, top_p, top_k, repetition_penalty],
             outputs=[raw_output_stream, markdown_output]
         )
         
+        ocr_button.click(
+            fn=extract_text_with_coordinates,
+            inputs=[image_input],
+            outputs=[ocr_text, ocr_vis]
+        )
+
         generate_pdf_btn.click(
             fn=generate_and_preview_pdf,
             inputs=[image_input, raw_output_stream, font_size, line_spacing, alignment, image_size],
@@ -315,7 +350,7 @@ def create_gradio_interface():
 
         clear_btn.click(
             clear_all_outputs,
-            outputs=[image_input, prompt_input, raw_output_stream, markdown_output, pdf_output_file, pdf_preview_gallery]
+            outputs=[image_input, prompt_input, raw_output_stream, markdown_output, pdf_output_file, pdf_preview_gallery, ocr_text, ocr_vis]
         )
     return demo