101 lines
3.1 KiB
Python
101 lines
3.1 KiB
Python
import cv2
|
|
import numpy as np
|
|
import sys
|
|
import json
|
|
from pathlib import Path
|
|
import os
|
|
from PIL import Image
|
|
|
|
def order_points(pts):
|
|
rect = np.zeros((4, 2), dtype="float32")
|
|
s = pts.sum(axis=1)
|
|
rect[0] = pts[np.argmin(s)]
|
|
rect[2] = pts[np.argmax(s)]
|
|
diff = np.diff(pts, axis=1)
|
|
rect[1] = pts[np.argmin(diff)]
|
|
rect[3] = pts[np.argmax(diff)]
|
|
return rect
|
|
|
|
def process_document(image_path, doc_type, output_dir):
|
|
try:
|
|
# Carrega e verifica a imagem
|
|
image = cv2.imread(str(image_path))
|
|
if image is None:
|
|
raise Exception("Não foi possível carregar a imagem")
|
|
|
|
original = image.copy()
|
|
|
|
# Pré-processamento
|
|
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
|
blurred = cv2.GaussianBlur(gray, (5, 5), 0)
|
|
edges = cv2.Canny(blurred, 50, 150)
|
|
|
|
# Encontra contornos
|
|
contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
|
|
if not contours:
|
|
raise Exception("Nenhum contorno encontrado")
|
|
|
|
# Ordena contornos por área
|
|
contours = sorted(contours, key=cv2.contourArea, reverse=True)
|
|
doc_corners = None
|
|
|
|
# Procura o documento
|
|
for contour in contours:
|
|
peri = cv2.arcLength(contour, True)
|
|
approx = cv2.approxPolyDP(contour, 0.02 * peri, True)
|
|
if len(approx) == 4:
|
|
doc_corners = approx
|
|
break
|
|
|
|
if doc_corners is None:
|
|
raise Exception("Documento não encontrado")
|
|
|
|
# Processa cantos e aplica transformação
|
|
doc_corners = doc_corners.reshape(4, 2)
|
|
ordered_corners = order_points(doc_corners)
|
|
|
|
(tl, tr, br, bl) = ordered_corners
|
|
widthA = np.linalg.norm(br - bl)
|
|
widthB = np.linalg.norm(tr - tl)
|
|
maxWidth = max(int(widthA), int(widthB))
|
|
|
|
heightA = np.linalg.norm(tr - br)
|
|
heightB = np.linalg.norm(tl - bl)
|
|
maxHeight = max(int(heightA), int(heightB))
|
|
|
|
dst = np.array([
|
|
[0, 0],
|
|
[maxWidth - 1, 0],
|
|
[maxWidth - 1, maxHeight - 1],
|
|
[0, maxHeight - 1]
|
|
], dtype="float32")
|
|
|
|
# Aplica transformação perspectiva
|
|
matrix = cv2.getPerspectiveTransform(ordered_corners, dst)
|
|
warped = cv2.warpPerspective(original, matrix, (maxWidth, maxHeight))
|
|
|
|
# Salva resultado
|
|
output_path = os.path.join(output_dir, f"{doc_type}_{Path(image_path).name}")
|
|
cv2.imwrite(str(output_path), warped)
|
|
|
|
return {
|
|
"success": True,
|
|
"message": "Processamento concluído com sucesso",
|
|
"width": maxWidth,
|
|
"height": maxHeight
|
|
}
|
|
|
|
except Exception as e:
|
|
return {
|
|
"success": False,
|
|
"message": str(e)
|
|
}
|
|
|
|
if __name__ == "__main__":
|
|
if len(sys.argv) != 4:
|
|
print(json.dumps({"success": False, "message": "Uso: python document_detector.py <input_path> <doc_type> <output_dir>"}))
|
|
sys.exit(1)
|
|
|
|
result = process_document(Path(sys.argv[1]), sys.argv[2], Path(sys.argv[3]))
|
|
print(json.dumps(result))
|