app.py

import os
os.environ["CUDA_VISIBLE_DEVICES"] = ""


import gradio as gr
import torch
import numpy as np
import cv2
from PIL import Image
from transformers import CLIPProcessor, CLIPModel
from paddleocr import PaddleOCR, TextDetection
from functools import lru_cache


MODEL_HUB_ID = "imperiusrex/printedpaddle"
# Setup
clip_model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14")
clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")

# Set device to CPU
device = "cpu"
clip_model.to(device)

# Language map for OCR models
def process_image(img_path):
    """
    Processes an image to detect, crop, and OCR text, returning it in reading order.
    Args:
        img_path: The path to the image file.
    Returns:
        A string containing the reconstructed text.
    """
    # Load CLIP model and processor
    clip_model = CLIPModel.from_pretrained("openai/clip-vit-large-patch14")
    processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")

    # Candidate language phrases for detection
    candidates = [
        "This is English text",
        # "This is Hindi text",
        # "This is Tamil text",
        "This is Telugu text",
        # "This is Bengali text",
        # "This is Arabic text",
        "This is Chinese text",
        # "This is Japanese text",
        "This is Korean text",
        "This is Russian text",
        # "This is Kannada text",
        # "This is Malayalam text",
        # "This is Marathi text",
        # "This is Urdu text",
        # "This is French text",
        # "This is Spanish text",
        # "This is Italian text",
        # "This is Portuguese text",
        # "This is Romanian text",
        # "This is Hungarian text",
        # "This is Indonesian text",
        # "This is Lithuanian text",
        # "This is Chinese Traditional text",
        # "This is Malay text",
        # "This is Dutch text",
        # "This is Norwegian text",
        # "This is Bosnian text",
        # "This is Polish text",
        # "This is Czech text",
        # "This is Slovak text",
        # "This is Welsh text",
        # "This is Slovenian text",
        # "This is Danish text",
        # "This is Albanian text",
        # "This is Estonian text",
        # "This is Swedish text",
        # "This is Irish text",
        # "This is Swahili text",
        # "This is Croatian text",
        # "This is Uzbek text",
        # "This is Turkish text",
        "This is Latin text",
        # "This is Belarusian text",
        # "This is Ukrainian text"
    ]

    # Map detected languages to PaddleOCR language codes
    lang_map = {
        "english": "en",
        # "hindi": "hi",
        # "tamil": "ta",
        "telugu": "te",
        # "bengali": "bn",
        # "arabic": "ar",
        "chinese": "ch",
        # "japanese": "japan",
        "korean": "korean",
        "russian": "ru",
        # "kannada": "kn",
        # "malayalam": "ml",
        # "marathi": "mr",
        # "urdu": "ur",
        # "french": "fr",
        # "spanish": "es",
        # "italian": "it",
        # "portuguese": "pt",
        # "romanian": "ro",
        # "hungarian": "hu",
        # "indonesian": "id",
        # "lithuanian": "lt",
        # "chinese traditional": "chinese_cht",
        # "malay": "ms",
        # "dutch": "nl",
        # "norwegian": "no",
        # "bosnian": "bs",
        # "polish": "pl",
        # "czech": "cs",
        # "slovak": "sk",
        # "welsh": "cy",
        # "slovenian": "sl",
        # "danish": "da",
        # "albanian": "sq",
        # "estonian": "et",
        # "swedish": "sv",
        # "irish": "ga",
        # "swahili": "sw",
        # "croatian": "hr",
        # "uzbek": "uz",
        # "turkish": "tr",
        "latin": "la",
        # "belarusian": "be",
        # "ukrainian": "uk"
    }

    # Text Detection
    arr = []
    model_det = TextDetection(model_name="PP-OCRv5_server_det")
    output = model_det.predict(img_path, batch_size=1)
    for res in output:
        polys = res['dt_polys']
        if polys is not None:
            arr.extend(polys.tolist())
    arr = sorted(arr, key=lambda box: (box[0][1], box[0][0]))

    # Image Cropping and Warping
    img = cv2.imread(img_path)
    cropped_images = []
    for i, box in enumerate(arr):
        box = np.array(box, dtype=np.float32)
        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)
        M = cv2.getPerspectiveTransform(box, dst_rect)
        warped = cv2.warpPerspective(img, M, (width, height))
        cropped_images.append(warped)

    # Perform language detection for each cropped image and then OCR
    predicted_texts = []
    for i, cropped_img in enumerate(cropped_images):
        # Get probabilities
        inputs = processor(text=candidates, images=cropped_img, return_tensors="pt", padding=True)
        with torch.no_grad():
            logits_per_image = clip_model(**inputs).logits_per_image
            probs = logits_per_image.softmax(dim=1)

        # Get best language match
        best = probs.argmax().item()
        detected_lang_phrase = candidates[best]
        detected_lang = detected_lang_phrase.split()[-2].lower()
        lang_code = lang_map.get(detected_lang, "en")

        # Perform OCR for the current cropped image with the detected language
        ocr = PaddleOCR(
            use_doc_orientation_classify=False,
            use_doc_unwarping=False,
            use_textline_orientation=False,
            lang=lang_code,
            device="cpu"
        )

        result = ocr.predict(cropped_img)

        text_for_this_image = ""
        if result and result[0] and 'rec_texts' in result[0]:
             text_for_this_image = " ".join(result[0]['rec_texts'])

        predicted_texts.append(text_for_this_image)


    def get_box_center(box):
      """Calculates the center of a bounding box."""
      x_coords = [p[0] for p in box]
      y_coords = [p[1] for p in box]
      center_x = sum(x_coords) / len(x_coords)
      center_y = sum(y_coords) / len(y_coords)
      return center_x, center_y

    # --- Step 1: Read all text and their centroid coordinates ---
    all_text_blocks = []
    for i, box in enumerate(arr):
        # Use the predicted text from the list
        text = predicted_texts[i]
        if text: # Only add if text is not empty
            center_x, center_y = get_box_center(box)
            all_text_blocks.append({
                "text": text,
                "center_x": center_x,
                "center_y": center_y
            })


    # --- Step 2: Sort by y-coordinate, then by x-coordinate, and group into lines ---
    reconstructed_text = ""
    if all_text_blocks:
        # Sort by center_y, then by center_x
        sorted_blocks = sorted(all_text_blocks, key=lambda item: (item["center_y"], item["center_x"]))

        lines = []
        if sorted_blocks:
            current_line = [sorted_blocks[0]]
            for block in sorted_blocks[1:]:
                # Check if the vertical centers are close enough to be on the same line
                if abs(block["center_y"] - current_line[-1]["center_y"]) < 40: # Y-threshold
                    current_line.append(block)
                else:
                    # Sort the current line by x-coordinate and add it to the lines list
                    current_line.sort(key=lambda item: item["center_x"])
                    lines.append(" ".join([item["text"] for item in current_line]))
                    current_line = [block]

            # Add the last line
            if current_line:
                current_line.sort(key=lambda item: item["center_x"])
                lines.append(" ".join([item["text"] for item in current_line]))

        # --- Step 3: Join the lines into a single string ---
        reconstructed_text = "\n".join(lines)

    return reconstructed_text

iface = gr.Interface(
    fn=process_image,
    inputs=gr.Image(type="filepath"),
    outputs=gr.Text(),
    title="Image OCR and Text Reconstruction",
    description="Upload an image to perform text detection, cropping, language detection, OCR, and reconstruct the text in reading order."
)

if __name__== "__main__":
    iface.launch(debug=True)