Create app.py

app.py

@@ -0,0 +1,244 @@
+# Import the GPU decorator for ZeroGPU Spaces
+# This will be a no-op if the space is not configured for ZeroGPU
+# but it is required for the specified hardware to work correctly.
+from spaces import GPU
+
+import os
+import cv2
+import numpy as np
+import torch
+import tempfile
+import gradio as gr
+from PIL import Image
+from pdf2image import convert_from_path
+from reportlab.lib.pagesizes import letter
+from reportlab.pdfgen import canvas
+from transformers import TrOCRProcessor, VisionEncoderDecoderModel
+from paddleocr import PaddleOCR, TextDetection
+
+# Set the GPU device if available
+# The `spaces.GPU` decorator handles the dynamic GPU allocation, but we still need to
+# specify the device for PyTorch and other GPU-enabled libraries.
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+print(f"Using device: {device}")
+
+# --- MODEL LOADING ---
+# Load models globally so they are only initialized once when the app starts.
+
+# Initialize the PaddleOCR detection model
+# `use_angle_cls=False` is set for efficiency, as we are already using
+# perspective warping to straighten the text.
+print("Initializing PaddleOCR text detection model...")
+try:
+    # Use the PaddleOCR class with a specific model for detection only
+    det_model = PaddleOCR(use_angle_cls=False, lang='en', use_gpu=torch.cuda.is_available(), show_log=False)
+except Exception as e:
+    print(f"Error initializing PaddleOCR: {e}")
+    det_model = None
+
+# Initialize the TrOCR recognition model and processor
+print("Initializing TrOCR text recognition model...")
+try:
+    trocr_processor = TrOCRProcessor.from_pretrained("microsoft/trocr-large-handwritten")
+    trocr_model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-large-handwritten")
+    trocr_model.eval()
+    trocr_model.to(device)
+except Exception as e:
+    print(f"Error initializing TrOCR: {e}")
+    trocr_model = None
+    trocr_processor = None
+
+# Helper function to save a temp image
+def save_temp_image(img):
+    """Save an image array to a temporary file and return the path."""
+    temp_fd, temp_path = tempfile.mkstemp(suffix='.png')
+    cv2.imwrite(temp_path, img)
+    os.close(temp_fd)
+    return temp_path
+
+def process_image_page(img):
+    """
+    Process a single image to detect polygons, crop regions, and recognize text.
+    Returns a list of [box, text] for each cropped region and the original PIL image.
+    """
+    if det_model is None or trocr_model is None:
+        raise RuntimeError("OCR models are not loaded. Please check logs for errors.")
+
+    # Convert OpenCV image (BGR numpy array) to PIL Image (RGB)
+    original_pil_image = Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
+    
+    # PaddleOCR's predict method takes a file path, so we'll save the image to a temp file
+    temp_image_path = save_temp_image(img)
+    
+    # Detect polygons using PaddleOCR
+    # The `ocr` method in PaddleOCR returns both detection and recognition results.
+    # We will use it just for the detection polygons.
+    ocr_result = det_model.ocr(temp_image_path)
+    os.remove(temp_image_path)
+    
+    arr = []
+    # The OCR result is a list of lists, where each inner list represents a text line.
+    # The first element is the bounding box coordinates.
+    for line in ocr_result[0]:
+        arr.append(line[0])
+
+    print(f"Detected {len(arr)} lines in this page.")
+
+    cropped_images = []
+    for box in arr:
+        box = np.array(box, dtype=np.float32)
+
+        # Compute width and height of the straightened image
+        width_a = np.linalg.norm(box[0] - box[1])
+        width_b = np.linalg.norm(box[2] - box[3])
+        height_a = np.linalg.norm(box[0] - box[3])
+        height_b = np.linalg.norm(box[1] - box[2])
+
+        width = int(max(width_a, width_b))
+        height = int(max(height_a, height_b))
+
+        dst_rect = np.array([
+            [0, 0],
+            [width - 1, 0],
+            [width - 1, height - 1],
+            [0, height - 1]
+        ], dtype=np.float32)
+
+        # Perspective transform
+        M = cv2.getPerspectiveTransform(box, dst_rect)
+        warped = cv2.warpPerspective(img, M, (width, height))
+        cropped_images.append(warped)
+
+    # Reverse cropped images and corresponding boxes
+    cropped_images.reverse()
+    arr.reverse()
+
+    # Text recognition with TrOCR
+    results = []
+    for i, crop in enumerate(cropped_images):
+        image_pil = Image.fromarray(cv2.cvtColor(crop, cv2.COLOR_BGR2RGB))
+        pixel_values = trocr_processor(images=image_pil, return_tensors="pt").pixel_values.to(device)
+
+        with torch.no_grad():
+            generated_ids = trocr_model.generate(pixel_values, max_new_tokens=64)
+            generated_text = trocr_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
+
+        results.append([arr[i], generated_text])
+        print(f"Recognized: {generated_text}")
+
+    return results, original_pil_image
+
+def process_file_and_create_pdf(file):
+    """
+    Main function to process a file (image or PDF) and return a path to a new PDF.
+    The @GPU decorator ensures this function is run on the GPU.
+    """
+    if file is None:
+        return None, "Please upload a file."
+
+    temp_output_dir = tempfile.mkdtemp()
+    output_pdf_path = os.path.join(temp_output_dir, "ocr_results.pdf")
+
+    try:
+        if file.name.lower().endswith('.pdf'):
+            # Convert PDF to images
+            print(f"Converting PDF {file.name} to images...")
+            # Use `poppler_path` if poppler is installed on the system, otherwise
+            # it might be necessary to install it via a `packages.txt` file.
+            # Here we assume it's available.
+            images = convert_from_path(file.name, dpi=300)
+
+            c = canvas.Canvas(output_pdf_path, pagesize=letter)
+            width, height = letter
+            
+            for page_num, page in enumerate(images):
+                print(f"\nProcessing page {page_num + 1}")
+                img_cv = cv2.cvtColor(np.array(page), cv2.COLOR_RGB2BGR)
+                results, original_image = process_image_page(img_cv)
+
+                c.setFont("Helvetica-Bold", 14)
+                c.drawString(50, height - 40, f"Page {page_num + 1} - OCR Results")
+                
+                y = height - 60
+                c.setFont("Helvetica", 12)
+                for _, text in results:
+                    c.drawString(50, y, text)
+                    y -= 15
+                    if y < 50:
+                        c.showPage()
+                        c.setFont("Helvetica-Bold", 14)
+                        c.drawString(50, height - 40, f"Page {page_num + 1} (cont.) - OCR Results")
+                        y = height - 60
+                c.showPage()
+            c.save()
+
+        else: # Handle single image file
+            img_cv = cv2.imread(file.name)
+            if img_cv is None:
+                raise ValueError("Failed to load image.")
+            
+            results, original_image = process_image_page(img_cv)
+
+            c = canvas.Canvas(output_pdf_path, pagesize=letter)
+            width, height = letter
+            c.setFont("Helvetica-Bold", 14)
+            c.drawString(50, height - 40, "Image OCR Results")
+            
+            # The input file from Gradio is a temp file that will be cleaned up.
+            # We can't display it directly in the PDF from its path.
+            # To draw it in the PDF, we save it to a new temporary path.
+            temp_img_path = os.path.join(temp_output_dir, "original_image.png")
+            original_image.save(temp_img_path)
+            
+            # Draw the image on the PDF
+            c.drawImage(temp_img_path, 50, height - 300, width=200, preserveAspectRatio=True)
+            
+            y = height - 350
+            c.setFont("Helvetica", 12)
+            for _, text in results:
+                c.drawString(50, y, text)
+                y -= 15
+                if y < 50:
+                    c.showPage()
+                    c.setFont("Helvetica", 12)
+                    y = height - 50
+            c.save()
+            os.remove(temp_img_path)
+
+        return output_pdf_path
+    
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        # Clean up temporary directory on error
+        # shutil.rmtree(temp_output_dir)
+        return None
+
+# Gradio Interface
+# The `@GPU` decorator is used here to ensure this function runs on a GPU.
+@GPU
+def process_file_for_gradio(file):
+    # This wrapper function is needed because Gradio's `File` component passes a temp file.
+    # We call our main processing function and return the path to the output PDF.
+    output_path = process_file_and_create_pdf(file)
+    if output_path is None:
+        return None
+    return output_path
+
+demo = gr.Interface(
+    fn=process_file_for_gradio,
+    inputs=gr.File(label="Upload an Image (PNG, JPG) or a PDF", file_types=['.png', '.jpg', '.jpeg', '.pdf']),
+    outputs=gr.File(label="Download OCR Results PDF"),
+    title="OCR App with PaddleOCR and TrOCR",
+    description="Upload an image or a multi-page PDF to get an output PDF with the recognized text from each page.",
+    examples=[
+        # Here you can provide paths to example files in your repo
+        # "example.png",
+        # "example.pdf"
+    ]
+)
+
+if __name__ == "__main__":
+    # You will need to set the hardware configuration in the `README.md` file
+    # of your Hugging Face Space for the GPU to be available.
+    demo.launch()