Spaces:
Paused
Paused
Update app.py
Browse files
app.py
CHANGED
@@ -1,44 +1,42 @@
|
|
1 |
-
|
2 |
from transformers import VisionEncoderDecoderModel, DonutProcessor
|
3 |
from PIL import Image
|
4 |
from pdf2image import convert_from_bytes
|
5 |
import gradio as gr
|
6 |
|
7 |
# Configuraci贸n del modelo Donut
|
8 |
-
MODEL_ID = "mychen76/invoice-and-receipts_donut_v1"
|
9 |
print("Cargando modelo Donut...")
|
10 |
model = VisionEncoderDecoderModel.from_pretrained(MODEL_ID)
|
11 |
processor = DonutProcessor.from_pretrained(MODEL_ID)
|
12 |
model.eval()
|
13 |
|
14 |
# Funci贸n para procesar documentos
|
15 |
-
def process_document(
|
16 |
-
#
|
17 |
-
if
|
18 |
-
|
19 |
-
|
20 |
else:
|
21 |
-
|
22 |
-
images = [Image.open(file).convert("RGB")]
|
23 |
|
24 |
results = []
|
25 |
for img in images:
|
26 |
# Preprocesar la imagen
|
27 |
inputs = processor(img, return_tensors="pt", max_patches=1024)
|
28 |
-
# Generar
|
29 |
with torch.no_grad():
|
30 |
outputs = model.generate(**inputs)
|
31 |
# Decodificar resultado
|
32 |
result = processor.batch_decode(outputs, skip_special_tokens=True)[0]
|
33 |
results.append(result)
|
34 |
|
35 |
-
# Combina los resultados si son m煤ltiples p谩ginas
|
36 |
return results
|
37 |
|
38 |
# Interfaz Gradio
|
39 |
iface = gr.Interface(
|
40 |
fn=process_document,
|
41 |
-
inputs=gr.File(label="Sube tu factura o recibo (PDF o imagen)", type="
|
42 |
outputs="json",
|
43 |
title="Donut OCR - Extracci贸n de datos de facturas",
|
44 |
description="Sube un PDF o imagen y extrae informaci贸n estructurada (n煤mero de factura, fecha, monto, etc.) utilizando Donut OCR."
|
|
|
1 |
+
iimport torch
|
2 |
from transformers import VisionEncoderDecoderModel, DonutProcessor
|
3 |
from PIL import Image
|
4 |
from pdf2image import convert_from_bytes
|
5 |
import gradio as gr
|
6 |
|
7 |
# Configuraci贸n del modelo Donut
|
8 |
+
MODEL_ID = "mychen76/invoice-and-receipts_donut_v1"
|
9 |
print("Cargando modelo Donut...")
|
10 |
model = VisionEncoderDecoderModel.from_pretrained(MODEL_ID)
|
11 |
processor = DonutProcessor.from_pretrained(MODEL_ID)
|
12 |
model.eval()
|
13 |
|
14 |
# Funci贸n para procesar documentos
|
15 |
+
def process_document(file_path):
|
16 |
+
# Leer y procesar el archivo
|
17 |
+
if file_path.endswith(".pdf"):
|
18 |
+
with open(file_path, "rb") as pdf_file:
|
19 |
+
images = convert_from_bytes(pdf_file.read(), dpi=300)
|
20 |
else:
|
21 |
+
images = [Image.open(file_path).convert("RGB")]
|
|
|
22 |
|
23 |
results = []
|
24 |
for img in images:
|
25 |
# Preprocesar la imagen
|
26 |
inputs = processor(img, return_tensors="pt", max_patches=1024)
|
27 |
+
# Generar predicci贸n
|
28 |
with torch.no_grad():
|
29 |
outputs = model.generate(**inputs)
|
30 |
# Decodificar resultado
|
31 |
result = processor.batch_decode(outputs, skip_special_tokens=True)[0]
|
32 |
results.append(result)
|
33 |
|
|
|
34 |
return results
|
35 |
|
36 |
# Interfaz Gradio
|
37 |
iface = gr.Interface(
|
38 |
fn=process_document,
|
39 |
+
inputs=gr.File(label="Sube tu factura o recibo (PDF o imagen)", type="filepath"),
|
40 |
outputs="json",
|
41 |
title="Donut OCR - Extracci贸n de datos de facturas",
|
42 |
description="Sube un PDF o imagen y extrae informaci贸n estructurada (n煤mero de factura, fecha, monto, etc.) utilizando Donut OCR."
|