Biifruu commited on
Commit
9bf9d52
·
verified ·
1 Parent(s): 88f4380

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +167 -0
app.py ADDED
@@ -0,0 +1,167 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import gradio as gr
2
+ import fitz # PyMuPDF
3
+ from PIL import Image
4
+ import numpy as np
5
+ import cv2
6
+ import pytesseract
7
+ import base64
8
+ import os
9
+
10
+ # ---------- OCR y limpieza de texto ----------
11
+
12
+ def clean_ocr_text(text):
13
+ lines = text.splitlines()
14
+ cleaned_lines = [line.strip() for line in lines if line.strip()]
15
+ return "\n".join(cleaned_lines)
16
+
17
+ # ---------- Funciones de imagen ----------
18
+
19
+ def text_area_ratio(image):
20
+ np_img = np.array(image.convert("L"))
21
+ _, thresh = cv2.threshold(np_img, 150, 255, cv2.THRESH_BINARY_INV)
22
+ contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
23
+ text_area = sum(w * h for x, y, w, h in [cv2.boundingRect(c) for c in contours if 8 < cv2.boundingRect(c)[3] < 40 and 5 < cv2.boundingRect(c)[2] < 100])
24
+ total_area = np_img.shape[0] * np_img.shape[1]
25
+ return text_area / total_area if total_area > 0 else 0
26
+
27
+ def has_significant_text(image):
28
+ return text_area_ratio(image) > 0.25
29
+
30
+ def is_primarily_text(image, ocr_threshold=30):
31
+ if has_significant_text(image):
32
+ ocr_result = pytesseract.image_to_string(image, lang="eng+spa")
33
+ return len(ocr_result.strip()) > ocr_threshold
34
+ return False
35
+
36
+ def is_likely_photo(crop):
37
+ np_crop = np.array(crop)
38
+ gray = cv2.cvtColor(np_crop, cv2.COLOR_RGB2GRAY)
39
+ return np.std(gray) > 25 and len(np.unique(gray)) > 50
40
+
41
+ def extract_visual_regions(image):
42
+ np_img = np.array(image.convert("RGB"))
43
+ gray = cv2.cvtColor(np_img, cv2.COLOR_RGB2GRAY)
44
+ _, binary = cv2.threshold(gray, 220, 255, cv2.THRESH_BINARY_INV)
45
+ closed = cv2.morphologyEx(binary, cv2.MORPH_CLOSE, cv2.getStructuringElement(cv2.MORPH_RECT, (15, 15)))
46
+
47
+ num_labels, labels, stats, _ = cv2.connectedComponentsWithStats(closed, connectivity=8)
48
+ results = []
49
+ for i in range(1, num_labels):
50
+ x, y, w, h, area = stats[i]
51
+ if area > 2000 and 0.3 < (w / float(h)) < 3.5:
52
+ bbox = (x, y, x + w, y + h)
53
+ crop = image.crop(bbox)
54
+ if is_likely_photo(crop) and text_area_ratio(crop) < 0.25 and not is_primarily_text(crop):
55
+ results.append(crop)
56
+ return results
57
+
58
+ # ---------- Extracción de texto + imágenes ----------
59
+
60
+ def extract_text_markdown(doc, image_paths, page_index, seen_xrefs):
61
+ markdown_output = f"\n## Página {page_index + 1}\n\n"
62
+ image_counter = 1
63
+ elements = []
64
+ page = doc[0]
65
+ blocks = page.get_text("dict")["blocks"]
66
+
67
+ for b in blocks:
68
+ y = b["bbox"][1]
69
+ if b["type"] == 0:
70
+ for line in b["lines"]:
71
+ line_y = line["bbox"][1]
72
+ line_text = " ".join([span["text"] for span in line["spans"]]).strip()
73
+ max_font_size = max([span.get("size", 10) for span in line["spans"]])
74
+ if line_text:
75
+ elements.append((line_y, line_text, max_font_size))
76
+
77
+ images_on_page = page.get_images(full=True)
78
+ for img_index, img in enumerate(images_on_page):
79
+ xref = img[0]
80
+ if xref in seen_xrefs:
81
+ continue
82
+ seen_xrefs.add(xref)
83
+ try:
84
+ base_image = page.parent.extract_image(xref)
85
+ image_bytes = base_image["image"]
86
+ ext = base_image["ext"]
87
+ image_path = f"/tmp/imagen_p{page_index + 1}_{img_index + 1}.{ext}"
88
+ with open(image_path, "wb") as f:
89
+ f.write(image_bytes)
90
+ image_paths.append(image_path)
91
+ elements.append((float("inf") - img_index, f"\n\n![imagen_{image_counter}]({image_path})\n", 10))
92
+ image_counter += 1
93
+ except Exception as e:
94
+ elements.append((float("inf"), f"[Error imagen: {e}]", 10))
95
+
96
+ elements.sort(key=lambda x: x[0])
97
+ previous_y = None
98
+
99
+ for y, text, font_size in elements:
100
+ is_header = font_size >= 14
101
+ if previous_y is not None and abs(y - previous_y) > 10:
102
+ markdown_output += "\n"
103
+ markdown_output += f"\n### {text.strip()}\n" if is_header else text.strip() + "\n"
104
+ previous_y = y
105
+
106
+ markdown_output += "\n---\n\n"
107
+ return markdown_output.strip()
108
+
109
+ # ---------- Función principal ----------
110
+
111
+ def convert(pdf_file):
112
+ temp_pdf_path = pdf_file.name
113
+ doc = fitz.open(temp_pdf_path)
114
+ markdown_output = ""
115
+ image_paths = []
116
+ seen_xrefs = set()
117
+
118
+ for page_num in range(len(doc)):
119
+ page = doc[page_num]
120
+ text = page.get_text("text").strip()
121
+
122
+ if len(text) > 30:
123
+ markdown_output += extract_text_markdown([page], image_paths, page_num, seen_xrefs) + "\n"
124
+ else:
125
+ markdown_output += f"\n## Página {page_num + 1}\n\n"
126
+ pix = page.get_pixmap(dpi=300)
127
+ img = Image.frombytes("RGB", [pix.width, pix.height], pix.samples)
128
+ image_path = f"/tmp/ocr_page_{page_num + 1}.jpg"
129
+ img.save(image_path)
130
+ image_paths.append(image_path)
131
+ markdown_output += f"![imagen_pagina_{page_num + 1}]({image_path})\n"
132
+
133
+ try:
134
+ ocr_text = pytesseract.image_to_string(img)
135
+ except pytesseract.TesseractError:
136
+ ocr_text = ""
137
+ markdown_output += clean_ocr_text(ocr_text) + "\n"
138
+
139
+ crops = extract_visual_regions(img)
140
+ for i, crop in enumerate(crops):
141
+ crop_path = f"/tmp/recorte_p{page_num + 1}_{i + 1}.jpg"
142
+ crop.save(crop_path)
143
+ image_paths.append(crop_path)
144
+ markdown_output += f"\n\n![imagen_detectada]({crop_path})\n"
145
+
146
+ markdown_output += "\n---\n\n"
147
+
148
+ markdown_path = "/tmp/resultado.md"
149
+ with open(markdown_path, "w", encoding="utf-8") as f:
150
+ f.write(markdown_output)
151
+
152
+ return markdown_output.strip(), image_paths, markdown_path
153
+
154
+ # ---------- Interfaz Gradio ----------
155
+
156
+ with gr.Blocks() as demo:
157
+ with gr.Row():
158
+ pdf_input = gr.File(label="Sube tu PDF", type="filepath", file_types=[".pdf"])
159
+ submit_btn = gr.Button("Procesar PDF")
160
+
161
+ markdown_output = gr.Textbox(label="Markdown generado", lines=25, interactive=True)
162
+ gallery_output = gr.Gallery(label="Imágenes extraídas y detectadas", type="file")
163
+ download_md = gr.File(label="Descargar Markdown")
164
+
165
+ submit_btn.click(fn=convert, inputs=[pdf_input], outputs=[markdown_output, gallery_output, download_md])
166
+
167
+ demo.launch()