function.py 7.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287
  1. # utils.py
  2. import fitz
  3. import cv2
  4. import numpy as np
  5. import re
  6. def redresser_image_auto(img_array):
  7. gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
  8. binary = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 15, 4)
  9. h, w = binary.shape
  10. kh = cv2.getStructuringElement(cv2.MORPH_RECT, (w // 10, 1))
  11. kv = cv2.getStructuringElement(cv2.MORPH_RECT, (1, h // 10))
  12. score_h = cv2.countNonZero(cv2.morphologyEx(binary, cv2.MORPH_OPEN, kh))
  13. score_v = cv2.countNonZero(cv2.morphologyEx(binary, cv2.MORPH_OPEN, kv))
  14. if (score_v / 1.5) > (score_h * 1.3):
  15. return cv2.rotate(img_array, cv2.ROTATE_90_COUNTERCLOCKWISE)
  16. return img_array
  17. def obtenir_zone_tableau_total(img_array):
  18. gray = cv2.cvtColor(img_array, cv2.COLOR_RGB2GRAY)
  19. blurred = cv2.GaussianBlur(gray, (5, 5), 0)
  20. binary = cv2.adaptiveThreshold(blurred, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 15, 4)
  21. h, w = binary.shape
  22. kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (20, 5))
  23. dilated = cv2.dilate(binary, kernel, iterations=3)
  24. contours, _ = cv2.findContours(dilated, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
  25. if not contours: return 0, h
  26. y_points = []
  27. for c in contours:
  28. x, y, w_c, h_c = cv2.boundingRect(c)
  29. if h_c > 10:
  30. y_points.append(y); y_points.append(y + h_c)
  31. if not y_points: return 0, h
  32. return max(0, min(y_points) - 100), min(h, max(y_points) + 100)
  33. def preparer_image_zoom_hd(pdf_path, page_index):
  34. """Gère le double passage pour extraire une image HD cadrée."""
  35. doc = fitz.open(pdf_path)
  36. page = doc.load_page(page_index)
  37. # 1. Localisation basse résolution
  38. pix_low = page.get_pixmap(matrix=fitz.Matrix(1, 1))
  39. img_low = np.frombuffer(pix_low.samples, dtype=np.uint8).reshape(pix_low.h, pix_low.w, 3)
  40. img_low = redresser_image_auto(img_low)
  41. y_min, y_max = obtenir_zone_tableau_total(img_low)
  42. # 2. Calcul du recadrage
  43. h_low = img_low.shape[0]
  44. y_start_pct = y_min / h_low
  45. y_end_pct = y_max / h_low
  46. full_rect = page.rect
  47. crop_rect = fitz.Rect(full_rect.x0, full_rect.y0 + (full_rect.height * y_start_pct),
  48. full_rect.x1, full_rect.y0 + (full_rect.height * y_end_pct))
  49. # 3. Rendu Haute Résolution (x4)
  50. pix_high = page.get_pixmap(matrix=fitz.Matrix(4, 4), clip=crop_rect, colorspace=fitz.csRGB)
  51. img_finale = np.frombuffer(pix_high.samples, dtype=np.uint8).reshape(pix_high.h, pix_high.w, 3)
  52. img_finale = redresser_image_auto(img_finale)
  53. doc.close()
  54. return img_finale
  55. def extraire_donnees_ocr(img, ocr_model):
  56. """Lance l'OCR et structure les résultats par coordonnées."""
  57. h_f, w_f = img.shape[:2]
  58. result = ocr_model.ocr(img, cls=True)
  59. extracted = []
  60. if result and result[0]:
  61. for line in result[0]:
  62. box, (text, conf) = line[0], line[1]
  63. if len(re.sub(r'[^a-zA-Z]', '', text)) > 2:
  64. continue
  65. if conf >= 0.6:
  66. x_c, y_c = sum(p[0] for p in box) / 4, sum(p[1] for p in box) / 4
  67. extracted.append({
  68. "text": text,
  69. "x_pct": round(x_c / w_f * 100, 1),
  70. "y_pct": round(y_c / h_f * 100, 1),
  71. "y_c": y_c,
  72. "x_c": x_c
  73. })
  74. extracted.sort(key=lambda r: (r["y_c"], r["x_c"]))
  75. return extracted
  76. def nettoyer_texte_ocr(text):
  77. if not text:
  78. return ""
  79. text = str(text)
  80. # supprimer artefacts OCR fréquents
  81. text = text.replace("]", "").replace("[", "").replace("/", "")
  82. # corriger O → 0 uniquement si texte numérique
  83. if re.match(r'^[\d\sO]+$', text):
  84. text = text.replace("O", "0")
  85. # corriger erreurs classiques R/RO
  86. text = text.replace("RO", "R0")
  87. # espaces propres
  88. text = re.sub(r"\s+", " ", text)
  89. return text.strip()
  90. def est_code_metier(text):
  91. return bool(re.match(r'^[A-Z]\d+$', text))
  92. def nettoyage_sortie_ocr(data):
  93. cleaned = []
  94. for r in data:
  95. # 🔥 sécurité : dict OU string
  96. if isinstance(r, dict):
  97. text = nettoyer_texte_ocr(r.get("text", ""))
  98. x = r.get("x_pct")
  99. y = r.get("y_pct")
  100. else:
  101. text = nettoyer_texte_ocr(str(r))
  102. x = None
  103. y = None
  104. # ❌ ignorer vide
  105. if not text:
  106. continue
  107. # ❌ bruit pur
  108. if re.match(r"^[\W_]+$", text):
  109. continue
  110. # 🔥 garder codes métier tels quels
  111. if est_code_metier(text):
  112. cleaned.append({
  113. "x_pct": x,
  114. "y_pct": y,
  115. "text": text
  116. })
  117. continue
  118. # ❌ filtrage normal
  119. if not re.search(r"[A-Za-z0-9]", text):
  120. continue
  121. cleaned.append({
  122. "x_pct": x,
  123. "y_pct": y,
  124. "text": text
  125. })
  126. return cleaned
  127. def formater_donnees_section(data_page, page_index):
  128. """
  129. Format lisible + prêt à parser
  130. """
  131. lignes = [f"\n--- DONNÉES PAGE {page_index + 1} ---"]
  132. for r in data_page:
  133. x = float(r["x_pct"])
  134. y = float(r["y_pct"])
  135. text = r["text"]
  136. lignes.append(f"x={x:.1f}% | y={y:.1f}% | {text}")
  137. return lignes
  138. def to_points(data):
  139. """
  140. Convertit directement en format DBSCAN :
  141. (x, y, text)
  142. """
  143. return [(r["x_pct"], r["y_pct"], r["text"]) for r in data]
  144. def sauvegarder_fichier_unique(contenu_total, pdf_path, section_name):
  145. """Sauvegarde toutes les pages accumulées dans un seul fichier."""
  146. # Nettoyage du nom de fichier
  147. nom_propre = pdf_path.replace('.pdf', '').replace(' ', '_')
  148. filename = f"{nom_propre}_{section_name}_complet.txt"
  149. with open(filename, "w", encoding="utf-8") as f:
  150. f.write("\n".join(contenu_total))
  151. return filename
  152. ########## DBSCAN
  153. import numpy as np
  154. def prepare_for_dbscan(points):
  155. """
  156. points = [[x, y, text], ...]
  157. """
  158. coords = np.array([[p[1]] for p in points]) # 🔥 uniquement Y
  159. return coords
  160. from sklearn.cluster import DBSCAN
  161. def cluster_lines(points, eps=0.4, min_samples=2):
  162. coords = prepare_for_dbscan(points)
  163. db = DBSCAN(eps=eps, min_samples=min_samples)
  164. labels = db.fit_predict(coords)
  165. clusters = {}
  166. for label, point in zip(labels, points):
  167. if label == -1:
  168. continue # bruit
  169. clusters.setdefault(label, []).append(point)
  170. return list(clusters.values())
  171. def build_lines(clusters):
  172. lignes = []
  173. for cluster in clusters:
  174. # tri gauche → droite
  175. cluster_sorted = sorted(cluster, key=lambda p: p[0])
  176. texte = " ".join([p[2] for p in cluster_sorted])
  177. lignes.append({
  178. "y": np.mean([p[1] for p in cluster]),
  179. "text": texte,
  180. "points": cluster_sorted
  181. })
  182. # tri haut → bas
  183. lignes = sorted(lignes, key=lambda l: l["y"])
  184. return lignes
  185. def merge_close_lines(lignes, threshold=0.6):
  186. merged = []
  187. prev = None
  188. for line in lignes:
  189. if prev is None:
  190. prev = line
  191. continue
  192. if abs(line["y"] - prev["y"]) < threshold:
  193. prev["text"] += " " + line["text"]
  194. else:
  195. merged.append(prev)
  196. prev = line
  197. if prev:
  198. merged.append(prev)
  199. return merged
  200. def split_tables(points, eps_y=2.0):
  201. """
  202. Sépare les tableaux par distance verticale (Y)
  203. """
  204. import numpy as np
  205. from sklearn.cluster import DBSCAN
  206. y_coords = np.array([[p[1]] for p in points]) # فقط Y
  207. clustering = DBSCAN(eps=eps_y, min_samples=5).fit(y_coords)
  208. tables = {}
  209. for label, point in zip(clustering.labels_, points):
  210. if label == -1:
  211. continue
  212. tables.setdefault(label, []).append(point)
  213. return list(tables.values())