Spaces:
Running
Running
| import gradio as gr | |
| import pytesseract | |
| from pdf2image import convert_from_path | |
| from PIL import Image | |
| import os | |
| # Function to perform OCR and search for a keyword | |
| def ocr_and_search(input_file, keyword, lang='fas'): # 'fas': Persian language (Farsi) | |
| extracted_text = "" | |
| keyword_found = False | |
| # Check if the input file is a PDF or an image | |
| if isinstance(input_file, str) and input_file.endswith('.pdf'): # Check if the file is a PDF | |
| # Convert PDF to images | |
| images = convert_from_path(input_file) | |
| # Loop through each image and perform OCR | |
| for page_number, image in enumerate(images): | |
| text = pytesseract.image_to_string(image, lang=lang) | |
| extracted_text += text | |
| # Check if the keyword is in the extracted text | |
| if keyword.lower() in text.lower(): | |
| keyword_found = True | |
| elif isinstance(input_file, Image.Image): # If the input is an image | |
| text = pytesseract.image_to_string(input_file, lang=lang) | |
| extracted_text = text | |
| # Check if the keyword is in the extracted text | |
| if keyword.lower() in text.lower(): | |
| keyword_found = True | |
| if not keyword_found: | |
| result_message = f"Keyword '{keyword}' not found in the document." | |
| else: | |
| result_message = f"Keyword '{keyword}' found in the document." | |
| return extracted_text, result_message | |
| # Create Gradio interface | |
| def gradio_interface(): | |
| # Define Gradio inputs and outputs | |
| input_type = gr.Radio(["PDF", "Image"], label="Choose Input Type", value="PDF") # Option to choose file type | |
| file_input = gr.File(label="Upload PDF/Image") | |
| keyword_input = gr.Textbox(label="Enter Keyword", value="ููุณูู") # Default keyword is 'ููุณูู' | |
| output_text = gr.Textbox(label="Extracted Text", interactive=False) | |
| output_message = gr.Textbox(label="Keyword Search Result", interactive=False) | |
| # Function to process the inputs and return the outputs | |
| def process(input_type, file, keyword): | |
| # Handle PDF and image accordingly | |
| if input_type == "PDF": | |
| extracted_text, result_message = ocr_and_search(file.name, keyword) | |
| else: # Handle image input | |
| image = Image.open(file.name) # Open image file | |
| extracted_text, result_message = ocr_and_search(image, keyword) | |
| return extracted_text, result_message | |
| # Create and launch Gradio interface | |
| gr.Interface(fn=process, | |
| inputs=[input_type, file_input, keyword_input], | |
| outputs=[output_text, output_message], | |
| title="OCR Keyword Search (PDF/Image)", | |
| description="Upload a PDF or Image, enter a keyword, and see the OCR results along with a search for the keyword." | |
| ).launch() | |
| # Call the function to create the interface | |
| gradio_interface() | |