Update app.py

app.py

@@ -8,53 +8,48 @@ 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.postprocessing.heatmap import draw_polys_on_image
 
-# Load models and processors
+# Carregar modelos e processadores
 det_model, det_processor = load_det_model(), load_det_processor()
 rec_model, rec_processor = load_rec_model(), load_rec_processor()
 
-# Load languages from JSON
+# Carregar opções de idioma
 with open("languages.json", "r") as file:
     languages = json.load(file)
-language_options = [(code, language) for code, language in languages.items()]
-
-def ocr_function(img, langs):
-    # Debugging: Print the languages input
-    print("Languages input:", langs)
-    lang_list = langs.split(',')
-    print("Language list:", lang_list)
-    
-    # Ensure we're passing a list of images and a list of languages with matching lengths
-    images = [img]  # Wrap the single image in a list
-    assert len(images) == len(lang_list), "Mismatch in the number of images and languages"
-    
-    predictions = run_ocr(images, lang_list, det_model, det_processor, rec_model, rec_processor)[0]
-    img_with_text = draw_polys_on_image(predictions["polys"], img)
-    return img_with_text, predictions
+language_options = [(language, code) for code, language in languages.items()]
 
+def ocr_function(img, lang_code):
+    # Ajuste aqui para garantir que lang_code é uma lista
+    predictions = run_ocr([img], [lang_code], det_model, det_processor, rec_model, rec_processor)[0]
+    img_with_text = draw_polys_on_image(predictions["polys"], img)
+    return img_with_text, predictions["text"]
 
 def text_line_detection_function(img):
-    preds = batch_detection([img], det_model, det_processor)[0]
+    preds = batch_inference([img], det_model, det_processor)[0]
     img_with_lines = draw_polys_on_image(preds["polygons"], img)
     return img_with_lines, preds
 
 with gr.Blocks() as app:
-    gr.Markdown("# Surya OCR and Text Line Detection Demo")
+    gr.Markdown("# Surya OCR e Detecção de Linhas de Texto")
     with gr.Tab("OCR"):
-        with gr.Row():
-            ocr_input_image = gr.Image(label="Input Image for OCR", type="pil")
-            ocr_language_selector = gr.Dropdown(label="Select Language(s) for OCR", choices=language_options, value="en")
-            ocr_output_image = gr.Image(label="OCR Output Image", type="pil", interactive=False)
-            ocr_json_output = gr.JSON(label="OCR JSON Output")
-        ocr_button = gr.Button("Run OCR")
-        ocr_button.click(fn=ocr_function, inputs=[ocr_input_image, ocr_language_selector], outputs=[ocr_output_image, ocr_json_output])
-
-    with gr.Tab("Text Line Detection"):
-        with gr.Row():
-            detection_input_image = gr.Image(label="Input Image for Detection", type="pil")
-            detection_output_image = gr.Image(label="Detection Output Image", type="pil", interactive=False)
-            detection_json_output = gr.JSON(label="Detection JSON Output")
-        detection_button = gr.Button("Run Text Line Detection")
-        detection_button.click(fn=text_line_detection_function, inputs=detection_input_image, outputs=[detection_output_image, detection_json_output])
+        with gr.Column():
+            ocr_input_image = gr.Image(label="Imagem de Entrada para OCR", type="pil")
+            ocr_language_selector = gr.Dropdown(label="Selecione o Idioma para OCR", choices=language_options, value="en")
+            ocr_run_button = gr.Button("Executar OCR")
+        with gr.Column():
+            ocr_output_image = gr.Image(label="Imagem de Saída do OCR", type="pil", interactive=False)
+            ocr_text_output = gr.TextArea(label="Texto Reconhecido")
+
+        ocr_run_button.click(fn=ocr_function, inputs=[ocr_input_image, ocr_language_selector], outputs=[ocr_output_image, ocr_text_output])
+
+    with gr.Tab("Detecção de Linhas de Texto"):
+        with gr.Column():
+            detection_input_image = gr.Image(label="Imagem de Entrada para Detecção", type="pil")
+            detection_run_button = gr.Button("Executar Detecção de Linhas de Texto")
+        with gr.Column():
+            detection_output_image = gr.Image(label="Imagem de Saída da Detecção", type="pil", interactive=False)
+            detection_json_output = gr.JSON(label="Saída JSON da Detecção")
+
+        detection_run_button.click(fn=text_line_detection_function, inputs=detection_input_image, outputs=[detection_output_image, detection_json_output])
 
 if __name__ == "__main__":
     app.launch()