Update app.py

app.py

@@ -20,119 +20,234 @@ device = "cpu"
 clip_model.to(device)
 
 # Language map for OCR models
-lang_map = {
-    "english": "en",
-    "telugu": "te",
-    "chinese": "ch",
-    "korean": "korean",
-}
-
-# Candidate labels for CLIP-based language detection
-candidates = [
-    "This is English text",
-    "This is Telugu text",
-    "This is Chinese text",
-    "This is Korean text"
-]
-
-# Initialize PaddleOCR for text detection only
-ocr_detector = PaddleOCR(use_textline_orientation=False, lang='en')
-
-
-
-# Cache OCR recognizers by language
-@lru_cache(maxsize=4)
-def get_ocr_recognizer(lang_code):
-    return PaddleOCR(lang=lang_code, use_textline_orientation=False, det=False, rec=True, use_gpu=False)
-
-# Helper function to validate detection result
-def is_valid_ocr_detection(det_result):
-    return det_result and isinstance(det_result[0], list) and len(det_result[0]) > 0
-
-# Main OCR pipeline
-def ocr_pipeline(image_np):
-    image_pil = Image.fromarray(image_np).convert("RGB")
-    img_cv = cv2.cvtColor(np.array(image_pil), cv2.COLOR_RGB2BGR)
-
-    detection_result = ocr_detector.ocr(image_np, det=True, rec=False)
-
-    if not is_valid_ocr_detection(detection_result):
-        return "❌ No text detected."
-
-    arr = [line[0] for line in detection_result[0]]
+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 = []
-    warped_boxes = []
-
-    for box in arr:
+    for i, box in enumerate(arr):
         box = np.array(box, dtype=np.float32)
-        width = int(max(np.linalg.norm(box[0] - box[1]), np.linalg.norm(box[2] - box[3])))
-        height = int(max(np.linalg.norm(box[0] - box[3]), np.linalg.norm(box[1] - box[2])))
+        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_cv, M, (width, height))
+        warped = cv2.warpPerspective(img, M, (width, height))
         cropped_images.append(warped)
-        warped_boxes.append(box)
-
-    final_output_lines = []
-
-    for i, crop in enumerate(cropped_images):
-        if crop.shape[0] < 10 or crop.shape[1] < 10:
-            continue
-
-        clip_inputs = clip_processor(text=candidates, images=Image.fromarray(crop), return_tensors="pt", padding=True)
-        clip_inputs = {k: v.to(device) for k, v in clip_inputs.items()}
 
+    # 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():
-            probs = clip_model(**clip_inputs).logits_per_image.softmax(dim=1)
-
-        lang_index = probs.argmax().item()
-        lang_detected = candidates[lang_index].split()[-2].lower()
-        lang_code = lang_map.get(lang_detected, "en")
-
-        ocr = get_ocr_recognizer(lang_code)
-        result = ocr.ocr(crop)
-
-        if not result or not isinstance(result[0], list) or len(result[0]) == 0:
-            continue
-
-        for line in result[0]:
-            text = line[1][0]
-            box = line[0]
-            center_x = sum([p[0] for p in box]) / 4
-            center_y = sum([p[1] for p in box]) / 4
-            final_output_lines.append({"text": text, "cx": center_x, "cy": center_y})
-
-    if not final_output_lines:
-        return "❌ No text detected."
-
-    sorted_blocks = sorted(final_output_lines, key=lambda b: (b["cy"], b["cx"]))
-    lines = []
-    current_line = [sorted_blocks[0]]
-
-    for block in sorted_blocks[1:]:
-        if abs(block["cy"] - current_line[-1]["cy"]) < 40:
-            current_line.append(block)
-        else:
-            lines.append(" ".join([x["text"] for x in sorted(current_line, key=lambda b: b["cx"])]))
-            current_line = [block]
-
-    if current_line:
-        lines.append(" ".join([x["text"] for x in sorted(current_line, key=lambda b: b["cx"])]))
-
-    return "\n".join(lines)
-
-# Gradio interface setup
-def build_interface():
-    return gr.Interface(
-        fn=ocr_pipeline,
-        inputs=gr.Image(type="numpy", label="Upload Printed Image"),
-        outputs="text",
-        title="\U0001F310 Multilingual Printed OCR with CLIP + PaddleOCR",
-        description="\U0001F4C4 Upload a printed document image. Detects language using CLIP and performs text detection + recognition with PaddleOCR."
-    )
-
-if __name__ == "__main__":
-    iface = build_interface()
-    iface.launch()
+            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)
+