Spaces:
Sleeping
Sleeping
| 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) |