Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import logging | |
| import os | |
| import json | |
| from PIL import Image, ImageDraw | |
| import torch | |
| from surya.ocr import run_ocr | |
| from surya.detection import batch_text_detection | |
| from surya.layout import batch_layout_detection | |
| from surya.ordering import batch_ordering | |
| from surya.model.detection.model import load_model as load_det_model, load_processor as load_det_processor | |
| from surya.model.recognition.model import load_model as load_rec_model | |
| from surya.model.recognition.processor import load_processor as load_rec_processor | |
| from surya.settings import settings | |
| # Configuração de logging | |
| logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s') | |
| logger = logging.getLogger(__name__) | |
| # Carregamento de modelos | |
| det_processor, det_model = load_det_processor(), load_det_model() | |
| rec_model, rec_processor = load_rec_model(), load_rec_processor() | |
| class CustomJSONEncoder(json.JSONEncoder): | |
| def default(self, obj): | |
| if isinstance(obj, Image.Image): | |
| return "Image object (not serializable)" | |
| return str(obj) | |
| def serialize_result(result): | |
| return json.dumps(result, cls=CustomJSONEncoder, indent=2) | |
| def save_metadata(results): | |
| output_file = "/mnt/data/ocr_metadata.json" # Caminho de armazenamento persistente | |
| with open(output_file, "w") as f: | |
| json.dump(results, f, cls=CustomJSONEncoder, indent=2) | |
| return output_file | |
| def ocr_workflow(images, langs): | |
| logger.info(f"Iniciando workflow OCR para {len(images)} imagens com idiomas: {langs}") | |
| results = [] | |
| for image_file in images: | |
| try: | |
| image = Image.open(image_file.name) | |
| predictions = run_ocr([image], [langs.split(',')], det_model, det_processor, rec_model, rec_processor) | |
| formatted_text = "\n".join([line.text for line in predictions[0].text_lines]) | |
| results.append({ | |
| "image": image_file.name, | |
| "text": formatted_text, | |
| "predictions": predictions[0].text_lines # Assuming text_lines is serializable | |
| }) | |
| except Exception as e: | |
| logger.error(f"Erro com {image_file.name}: {e}") | |
| results.append({"image": image_file.name, "error": str(e)}) | |
| metadata_file = save_metadata(results) | |
| return serialize_result(results), metadata_file | |
| with gr.Blocks(theme=gr.themes.Soft()) as demo: | |
| gr.Markdown("# Análise de Documentos com Surya") | |
| with gr.Tab("OCR"): | |
| gr.Markdown("## Reconhecimento Óptico de Caracteres para Múltiplas Imagens") | |
| with gr.Row(): | |
| ocr_input = gr.Files(label="Carregar Imagens ou PDFs (até 100)") | |
| ocr_langs = gr.Textbox(label="Idiomas (separados por vírgula)", value="en") | |
| ocr_button = gr.Button("Executar OCR") | |
| ocr_output = gr.JSON(label="Resultados OCR") | |
| ocr_file = gr.File(label="Baixar Metadata OCR") | |
| # Executa a função OCR e salva o arquivo de metadata para download | |
| ocr_button.click(ocr_workflow, inputs=[ocr_input, ocr_langs], outputs=[ocr_output, ocr_file]) | |
| if __name__ == "__main__": | |
| logger.info("Iniciando aplicativo Gradio...") | |
| demo.launch() | |