artificialguybr commited on
Commit
7d42bca
·
verified ·
1 Parent(s): 94a609f

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +116 -57
app.py CHANGED
@@ -15,90 +15,149 @@ from surya.settings import settings
15
  from surya.model.ordering.processor import load_processor as load_order_processor
16
  from surya.model.ordering.model import load_model as load_order_model
17
 
18
- # Set up logging
19
- logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
20
  logger = logging.getLogger(__name__)
21
 
22
- # Set environment variables for performance
 
23
  os.environ["RECOGNITION_BATCH_SIZE"] = "512"
24
  os.environ["DETECTOR_BATCH_SIZE"] = "36"
25
  os.environ["ORDER_BATCH_SIZE"] = "32"
26
  os.environ["RECOGNITION_STATIC_CACHE"] = "true"
27
 
28
- # Load models
29
- logger.info("Loading models...")
30
- det_processor, det_model = load_det_processor(), load_det_model()
31
- rec_model, rec_processor = load_rec_model(), load_rec_processor()
32
- layout_model = load_det_model(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
33
- layout_processor = load_det_processor(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
34
- order_model = load_order_model()
35
- order_processor = load_order_processor()
36
 
37
- # Compile recognition model
38
- logger.info("Compiling recognition model...")
39
- rec_model.decoder.model = torch.compile(rec_model.decoder.model)
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
40
 
41
  def ocr_workflow(image, langs):
42
- logger.info(f"Starting OCR workflow with languages: {langs}")
43
- image = Image.open(image.name)
44
- predictions = run_ocr([image], [langs.split(',')], det_model, det_processor, rec_model, rec_processor)
45
- logger.info("OCR workflow completed")
46
- return json.dumps(predictions, indent=2)
 
 
 
 
 
47
 
48
  def text_detection_workflow(image):
49
- logger.info("Starting text detection workflow")
50
- image = Image.open(image.name)
51
- predictions = batch_text_detection([image], det_model, det_processor)
52
- logger.info("Text detection workflow completed")
53
- return json.dumps(predictions, indent=2)
 
 
 
 
 
54
 
55
  def layout_analysis_workflow(image):
56
- logger.info("Starting layout analysis workflow")
57
- image = Image.open(image.name)
58
- line_predictions = batch_text_detection([image], det_model, det_processor)
59
- layout_predictions = batch_layout_detection([image], layout_model, layout_processor, line_predictions)
60
- logger.info("Layout analysis workflow completed")
61
- return json.dumps(layout_predictions, indent=2)
 
 
 
 
 
 
62
 
63
  def reading_order_workflow(image):
64
- logger.info("Starting reading order workflow")
65
- image = Image.open(image.name)
66
- line_predictions = batch_text_detection([image], det_model, det_processor)
67
- layout_predictions = batch_layout_detection([image], layout_model, layout_processor, line_predictions)
68
- bboxes = [pred['bbox'] for pred in layout_predictions[0]['bboxes']]
69
- order_predictions = batch_ordering([image], [bboxes], order_model, order_processor)
70
- logger.info("Reading order workflow completed")
71
- return json.dumps(order_predictions, indent=2)
 
 
 
 
 
 
 
72
 
73
  with gr.Blocks(theme=gr.themes.Soft()) as demo:
74
- gr.Markdown("# Surya Document Analysis")
75
 
76
  with gr.Tab("OCR"):
77
- gr.Markdown("## Optical Character Recognition")
78
  with gr.Row():
79
- ocr_input = gr.File(label="Upload Image or PDF")
80
- ocr_langs = gr.Textbox(label="Languages (comma-separated)", value="en")
81
- ocr_button = gr.Button("Run OCR")
82
- ocr_output = gr.JSON(label="OCR Results")
83
  ocr_button.click(ocr_workflow, inputs=[ocr_input, ocr_langs], outputs=ocr_output)
84
 
85
- with gr.Tab("Text Detection"):
86
- gr.Markdown("## Text Line Detection")
87
- det_input = gr.File(label="Upload Image or PDF")
88
- det_button = gr.Button("Run Text Detection")
89
- det_output = gr.JSON(label="Text Detection Results")
90
  det_button.click(text_detection_workflow, inputs=det_input, outputs=det_output)
91
 
92
- with gr.Tab("Layout Analysis"):
93
- gr.Markdown("## Layout Analysis and Reading Order")
94
- layout_input = gr.File(label="Upload Image or PDF")
95
- layout_button = gr.Button("Run Layout Analysis")
96
- order_button = gr.Button("Determine Reading Order")
97
- layout_output = gr.JSON(label="Layout Analysis Results")
98
- order_output = gr.JSON(label="Reading Order Results")
99
  layout_button.click(layout_analysis_workflow, inputs=layout_input, outputs=layout_output)
100
  order_button.click(reading_order_workflow, inputs=layout_input, outputs=order_output)
101
 
102
  if __name__ == "__main__":
103
- logger.info("Starting Gradio app...")
104
  demo.launch()
 
15
  from surya.model.ordering.processor import load_processor as load_order_processor
16
  from surya.model.ordering.model import load_model as load_order_model
17
 
18
+ # Configuração de logging mais detalhada
19
+ logging.basicConfig(level=logging.DEBUG, format='%(asctime)s - %(levelname)s - %(message)s')
20
  logger = logging.getLogger(__name__)
21
 
22
+ # Configuração de variáveis de ambiente
23
+ logger.info("Configurando variáveis de ambiente para otimização de performance")
24
  os.environ["RECOGNITION_BATCH_SIZE"] = "512"
25
  os.environ["DETECTOR_BATCH_SIZE"] = "36"
26
  os.environ["ORDER_BATCH_SIZE"] = "32"
27
  os.environ["RECOGNITION_STATIC_CACHE"] = "true"
28
 
29
+ # Carregamento de modelos
30
+ logger.info("Iniciando carregamento dos modelos...")
 
 
 
 
 
 
31
 
32
+ try:
33
+ logger.debug("Carregando modelo e processador de detecção...")
34
+ det_processor, det_model = load_det_processor(), load_det_model()
35
+ logger.debug("Modelo e processador de detecção carregados com sucesso")
36
+ except Exception as e:
37
+ logger.error(f"Erro ao carregar modelo de detecção: {e}")
38
+ raise
39
+
40
+ try:
41
+ logger.debug("Carregando modelo e processador de reconhecimento...")
42
+ rec_model, rec_processor = load_rec_model(), load_rec_processor()
43
+ logger.debug("Modelo e processador de reconhecimento carregados com sucesso")
44
+ except Exception as e:
45
+ logger.error(f"Erro ao carregar modelo de reconhecimento: {e}")
46
+ raise
47
+
48
+ try:
49
+ logger.debug("Carregando modelo e processador de layout...")
50
+ layout_model = load_det_model(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
51
+ layout_processor = load_det_processor(checkpoint=settings.LAYOUT_MODEL_CHECKPOINT)
52
+ logger.debug("Modelo e processador de layout carregados com sucesso")
53
+ except Exception as e:
54
+ logger.error(f"Erro ao carregar modelo de layout: {e}")
55
+ raise
56
+
57
+ try:
58
+ logger.debug("Carregando modelo e processador de ordenação...")
59
+ order_model = load_order_model()
60
+ order_processor = load_order_processor()
61
+ logger.debug("Modelo e processador de ordenação carregados com sucesso")
62
+ except Exception as e:
63
+ logger.error(f"Erro ao carregar modelo de ordenação: {e}")
64
+ raise
65
+
66
+ logger.info("Todos os modelos foram carregados com sucesso")
67
+
68
+ # Compilação do modelo de reconhecimento
69
+ logger.info("Iniciando compilação do modelo de reconhecimento...")
70
+ try:
71
+ rec_model.decoder.model = torch.compile(rec_model.decoder.model)
72
+ logger.info("Compilação do modelo de reconhecimento concluída com sucesso")
73
+ except Exception as e:
74
+ logger.error(f"Erro durante a compilação do modelo de reconhecimento: {e}")
75
+ logger.warning("Continuando sem compilação do modelo")
76
 
77
  def ocr_workflow(image, langs):
78
+ logger.info(f"Iniciando workflow OCR com idiomas: {langs}")
79
+ try:
80
+ image = Image.open(image.name)
81
+ logger.debug(f"Imagem carregada: {image.size}")
82
+ predictions = run_ocr([image], [langs.split(',')], det_model, det_processor, rec_model, rec_processor)
83
+ logger.info("Workflow OCR concluído com sucesso")
84
+ return json.dumps(predictions, indent=2)
85
+ except Exception as e:
86
+ logger.error(f"Erro durante o workflow OCR: {e}")
87
+ return json.dumps({"error": str(e)})
88
 
89
  def text_detection_workflow(image):
90
+ logger.info("Iniciando workflow de detecção de texto")
91
+ try:
92
+ image = Image.open(image.name)
93
+ logger.debug(f"Imagem carregada: {image.size}")
94
+ predictions = batch_text_detection([image], det_model, det_processor)
95
+ logger.info("Workflow de detecção de texto concluído com sucesso")
96
+ return json.dumps(predictions, indent=2)
97
+ except Exception as e:
98
+ logger.error(f"Erro durante o workflow de detecção de texto: {e}")
99
+ return json.dumps({"error": str(e)})
100
 
101
  def layout_analysis_workflow(image):
102
+ logger.info("Iniciando workflow de análise de layout")
103
+ try:
104
+ image = Image.open(image.name)
105
+ logger.debug(f"Imagem carregada: {image.size}")
106
+ line_predictions = batch_text_detection([image], det_model, det_processor)
107
+ logger.debug(f"Detecção de linhas concluída. Número de linhas detectadas: {len(line_predictions[0]['bboxes'])}")
108
+ layout_predictions = batch_layout_detection([image], layout_model, layout_processor, line_predictions)
109
+ logger.info("Workflow de análise de layout concluído com sucesso")
110
+ return json.dumps(layout_predictions, indent=2)
111
+ except Exception as e:
112
+ logger.error(f"Erro durante o workflow de análise de layout: {e}")
113
+ return json.dumps({"error": str(e)})
114
 
115
  def reading_order_workflow(image):
116
+ logger.info("Iniciando workflow de ordem de leitura")
117
+ try:
118
+ image = Image.open(image.name)
119
+ logger.debug(f"Imagem carregada: {image.size}")
120
+ line_predictions = batch_text_detection([image], det_model, det_processor)
121
+ logger.debug(f"Detecção de linhas concluída. Número de linhas detectadas: {len(line_predictions[0]['bboxes'])}")
122
+ layout_predictions = batch_layout_detection([image], layout_model, layout_processor, line_predictions)
123
+ logger.debug(f"Análise de layout concluída. Número de elementos de layout: {len(layout_predictions[0]['bboxes'])}")
124
+ bboxes = [pred['bbox'] for pred in layout_predictions[0]['bboxes']]
125
+ order_predictions = batch_ordering([image], [bboxes], order_model, order_processor)
126
+ logger.info("Workflow de ordem de leitura concluído com sucesso")
127
+ return json.dumps(order_predictions, indent=2)
128
+ except Exception as e:
129
+ logger.error(f"Erro durante o workflow de ordem de leitura: {e}")
130
+ return json.dumps({"error": str(e)})
131
 
132
  with gr.Blocks(theme=gr.themes.Soft()) as demo:
133
+ gr.Markdown("# Análise de Documentos com Surya")
134
 
135
  with gr.Tab("OCR"):
136
+ gr.Markdown("## Reconhecimento Óptico de Caracteres")
137
  with gr.Row():
138
+ ocr_input = gr.File(label="Carregar Imagem ou PDF")
139
+ ocr_langs = gr.Textbox(label="Idiomas (separados por vírgula)", value="en")
140
+ ocr_button = gr.Button("Executar OCR")
141
+ ocr_output = gr.JSON(label="Resultados OCR")
142
  ocr_button.click(ocr_workflow, inputs=[ocr_input, ocr_langs], outputs=ocr_output)
143
 
144
+ with gr.Tab("Detecção de Texto"):
145
+ gr.Markdown("## Detecção de Linhas de Texto")
146
+ det_input = gr.File(label="Carregar Imagem ou PDF")
147
+ det_button = gr.Button("Executar Detecção de Texto")
148
+ det_output = gr.JSON(label="Resultados da Detecção de Texto")
149
  det_button.click(text_detection_workflow, inputs=det_input, outputs=det_output)
150
 
151
+ with gr.Tab("Análise de Layout"):
152
+ gr.Markdown("## Análise de Layout e Ordem de Leitura")
153
+ layout_input = gr.File(label="Carregar Imagem ou PDF")
154
+ layout_button = gr.Button("Executar Análise de Layout")
155
+ order_button = gr.Button("Determinar Ordem de Leitura")
156
+ layout_output = gr.JSON(label="Resultados da Análise de Layout")
157
+ order_output = gr.JSON(label="Resultados da Ordem de Leitura")
158
  layout_button.click(layout_analysis_workflow, inputs=layout_input, outputs=layout_output)
159
  order_button.click(reading_order_workflow, inputs=layout_input, outputs=order_output)
160
 
161
  if __name__ == "__main__":
162
+ logger.info("Iniciando aplicativo Gradio...")
163
  demo.launch()