Update app.py

app.py

@@ -1,7 +1,6 @@
 import os
 os.environ["CUDA_VISIBLE_DEVICES"] = ""
 
-
 import gradio as gr
 import torch
 import numpy as np
@@ -9,10 +8,9 @@ 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")
@@ -21,115 +19,27 @@ clip_processor = CLIPProcessor.from_pretrained("openai/clip-vit-large-patch14")
 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
@@ -160,9 +70,11 @@ def process_image(img_path):
 
     # Perform language detection for each cropped image and then OCR
     predicted_texts = []
+    detected_languages_list = []  # store detected languages
+
     for i, cropped_img in enumerate(cropped_images):
         # Get probabilities
-        inputs = processor(text=candidates, images=cropped_img, return_tensors="pt", padding=True)
+        inputs = clip_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)
@@ -171,6 +83,8 @@ def process_image(img_path):
         best = probs.argmax().item()
         detected_lang_phrase = candidates[best]
         detected_lang = detected_lang_phrase.split()[-2].lower()
+        detected_languages_list.append(detected_lang)  # store detected language
+
         lang_code = lang_map.get(detected_lang, "en")
 
         # Perform OCR for the current cropped image with the detected language
@@ -186,25 +100,22 @@ def process_image(img_path):
 
         text_for_this_image = ""
         if result and result[0] and 'rec_texts' in result[0]:
-             text_for_this_image = " ".join(result[0]['rec_texts'])
+            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 ---
+        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
+        if text:
             center_x, center_y = get_box_center(box)
             all_text_blocks.append({
                 "text": text,
@@ -212,44 +123,40 @@ def process_image(img_path):
                 "center_y": center_y
             })
 
-
-    # --- Step 2: Sort by y-coordinate, then by x-coordinate, and group into lines ---
+    # Step 2: Sort 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
+                if abs(block["center_y"] - current_line[-1]["center_y"]) < 40:
                     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
+    # Step 3: Get unique detected languages
+    unique_languages = sorted(set(detected_languages_list))
+
+    return reconstructed_text, ", ".join(unique_languages)
 
 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."
+    outputs=[
+        gr.Textbox(label="Reconstructed Text"),
+        gr.Textbox(label="Detected Languages")
+    ],
+    title="Image OCR with Language Detection",
+    description="Upload an image to detect text, OCR it, reconstruct in reading order, and list unique detected languages."
 )
 
-if __name__== "__main__":
+if __name__ == "__main__":
     iface.launch(debug=True)
-