app.py

import gradio as gr
import numpy as np
#import imutils
import cv2
import torch
import matplotlib.pyplot as plt
from pdf2image import convert_from_path, convert_from_bytes
from collections import namedtuple
import pytesseract
import argparse

def convert_pdf_to_png(file_name):
    pages = convert_from_path(file_name)
    for i, page in enumerate(pages):
        page.save(f"{file_name[:-4]}_{i+1}.png", "PNG")


def alinearImagen(image, template, maxFeatures = 500, keepPercent = 0.2, debug = False):
  # Convertimos ambas imágenes en escala de grises
  imageGray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
  templateGray = cv2.cvtColor(template, cv2.COLOR_BGR2GRAY)
 
 	# Usamos ORB para detectar los keypoints y extraer los descriptores invariantes
  orb = cv2.ORB_create(maxFeatures)
  (kpsA, descsA) = orb.detectAndCompute(imageGray, None)
  (kpsB, descsB) = orb.detectAndCompute(templateGray, None)
  
  # Hacemos que coincidan
  method = cv2.DESCRIPTOR_MATCHER_BRUTEFORCE_HAMMING
  matcher = cv2.DescriptorMatcher_create(method)
  matches = matcher.match(descsA, descsB, None)

  # Ordenamos los matches por si distancia, cuando más pequeña sea, más similares son las regiones de los keypoints
  matches = sorted(matches, key=lambda x:x.distance)

  # Almacenamos únicamente los mejores matches
  keep = int(len(matches) * keepPercent)
  matches = matches[:keep]

  # Visualizamos la imagen con las líneas de relación correspondientes si estamos en modo debug
  if debug:
    matchedVis = cv2.drawMatches(image, kpsA, template, kpsB, matches, None)
    matchedVis = imutils.resize(matchedVis, width=1000)
    cv2_imshow(matchedVis)
    cv2.waitKey(0)

  # Reservamos memorias para las coordenadas de los keypoints con mejores matches
  ptsA = np.zeros((len(matches), 2), dtype="float")
  ptsB = np.zeros((len(matches), 2), dtype="float")

  # Recorremos los matches
  for (i, m) in enumerate(matches):
		# Indicamos que los keypoints A y B son mapas de uno a otro
    ptsA[i] = kpsA[m.queryIdx].pt
    ptsB[i] = kpsB[m.trainIdx].pt

  # Calcula la matriz homogrófica entre los dos conjuntos de puntos de matches
  (H, mask) = cv2.findHomography(ptsA, ptsB, method=cv2.RANSAC)
  
  # Usamos la matriz homográfica para alinear las imágenes
  (h, w) = template.shape[:2]
  aligned = cv2.warpPerspective(image, H, (w, h))

  # Devolvemos la imagen alineada
  return aligned

def cleanup_text(text):
  return "".join([c if ord(c) < 128 else "" for c in text]).strip()

def aplicaFiltro(imagen):
  r,g,b = cv2.split(imagen)
  th = cv2.adaptiveThreshold(b,255,cv2.ADAPTIVE_THRESH_GAUSSIAN_C,cv2.THRESH_BINARY,11,7)
  resultado = cv2.fastNlMeansDenoising(th, [50])
  return resultado

def recorte(imagen):
    _, thresh1 = cv2.threshold(imagen, 0, 255, cv2.THRESH_OTSU | cv2.THRESH_BINARY_INV)
    cnts = cv2.findContours(thresh1, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    cnts = imutils.grab_contours(cnts)

    xs = []
    ys = []
    hs = []
    ws = []
    
    for c in cnts:
        (x, y, w, h) = cv2.boundingRect(c)
        if w >= 5 and h >= 5 and w < 100 and h < 30:
            xs.append(x)
            ys.append(y)
            hs.append(h)
            ws.append(w)

    if len(xs) > 0:
        x1 = min(xs)
        y1 = min(ys)
        h1 = max(ys) + max(hs) - y1
        w1 = max(xs) + max(ws) - x1

        clone = imagen.copy()
        im = imagen[y1:y1 + h1, x1:x1 + w1]
        return (im, True)
    else:
        return (None, False)

from transformers import TrOCRProcessor, VisionEncoderDecoderModel
import requests

processor = TrOCRProcessor.from_pretrained("microsoft/trocr-base-handwritten")
model = VisionEncoderDecoderModel.from_pretrained("microsoft/trocr-base-handwritten")

def OCR3_F(nombreFichero):
  # Primero convertimos el fichero dado a PNG
  file_name = nombreFichero.name
  convert_pdf_to_png(file_name)

  # Leemos las dos imagenes generadas
  fich1 = cv2.imread(nombreFichero[:-4]+"_1.png")
  fich2 = cv2.imread(nombreFichero[:-4]+"_2.png")

  # Leemos las dos plantillas
  plantilla1 = cv2.imread("plantilla_0.png")
  plantilla2 = cv2.imread("plantilla_1.png")

  # Alineamos las dos imágenes, cada una con su correspondiente plantilla
  fich1_alineado = alinearImagen(fich1,plantilla1)
  fich2_alineado = alinearImagen(fich2,plantilla2)

  # Concatenamos las imágenes
  fichero = cv2.vconcat([fich1_alineado,fich2_alineado])
  plantilla = cv2.vconcat([plantilla1,plantilla2])

  # Generamos las localizaciones en las que aplicaremos OCR
  OCRLocation = namedtuple("OCRLocation", ["id", "bbox", "onlyNumber"])
  OCR_LOCATIONS = [
      OCRLocation("hojaCatastral", (1550, 55, 200, 32), True),
      OCRLocation("numeroParcela", (1550, 93, 200, 36), True),
      OCRLocation("idemPoligono", (1550, 134, 200, 35), True),
      OCRLocation("idemFotografia", (1550, 134, 200, 34), False),

      OCRLocation("terminoMunicipal", (1240, 254, 520, 38), False),
      OCRLocation("pago", (1300, 305, 460, 40), False),
      OCRLocation("partidoJudicial", (500, 312, 442, 38), False),

      OCRLocation("nombre", (386, 634, 604, 43), False),
      OCRLocation("pueblo", (1006, 632, 240, 44), False),

      OCRLocation("cultivos1", (212, 1373, 493, 43), False),
      OCRLocation("cultivos2", (212, 1426, 493, 42), False),
      OCRLocation("cultivos3", (212, 1480, 493, 41), False)
    ]

  fichFiltro = aplicaFiltro(fichero)

  parsingResults = []
  # loop over the locations of the document we are going to OCR
  for loc in OCR_LOCATIONS:
    # extract the OCR ROI from the aligned image
     (x, y, w, h) = loc.bbox
     roi = fichFiltro[y:y + h, x:x + w]

     im,b = recorte(roi)
     if (b == True):
        # OCR the ROI using Tesseract
        rgb = cv2.cvtColor(roi, cv2.COLOR_BGR2RGB)
        pixel_values = processor(rgb, return_tensors="pt").pixel_values
        generated_ids = model.generate(pixel_values)
        text = processor.batch_decode(generated_ids, skip_special_tokens=True)[0]

        for line in text.split("\n"):
        # if the line is empty, ignore it
          if len(line) == 0:
            continue
      
        parsingResults.append((loc, line))

  results = {}
  for (loc, line) in parsingResults:
    r = results.get(loc.id, None)
    if r is None:
      results[loc.id] = (line, loc._asdict())
    else:
      (existingText, loc) = r
      text = "{}\n{}".format(existingText, line)
      results[loc["id"]] = (text, loc)


  for (locID, result) in results.items():
    (text, loc) = result
    
    (x, y, w, h) = loc["bbox"]
    clean = cleanup_text(text)
    
    cv2.rectangle(fichero, (x, y), (x + w, y + h), (0, 255, 0), 2)
    
    for (i, line) in enumerate(text.split("\n")):
      startY = y + (i * 70) + 40
      cv2.putText(fichero, line, (x, startY),
			cv2.FONT_HERSHEY_SIMPLEX, 1.8, (0, 0, 255), 5)

  return fichero
    
# Creamos la interfaz y la lanzamos. 
gr.Interface(fn=OCR3_F, inputs="file", outputs="image", examples=["fich3.pdf", "fich4.pdf"]).launch()