Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| import cv2 | |
| from matplotlib import pyplot as plt | |
| import numpy as np | |
| from typing import List | |
| from ultralytics import YOLO | |
| import cv2 | |
| import numpy as np | |
| import os | |
| import shutil | |
| def recortar_notas(imagen_path: str) -> int: | |
| if os.path.exists('deteccion'): | |
| shutil.rmtree('deteccion') | |
| if os.path.exists('diarios'): | |
| shutil.rmtree('diarios') | |
| if os.path.exists('recorte'): | |
| shutil.rmtree('recorte') | |
| model = YOLO("seg-640.pt") | |
| results = model.predict(source=imagen_path, save=True, save_txt=True, project="diarios", classes=0, conf=0.5) | |
| # cargar imagen tif | |
| image = cv2.imread(imagen_path) | |
| # crear una carpeta para guardar los recortes | |
| filename = os.path.splitext(os.path.basename(results[0].path))[0] | |
| foldername = os.path.splitext(filename)[0] | |
| if not os.path.exists(f"recorte/{foldername}"): | |
| os.makedirs(f"recorte/{foldername}") | |
| datos = results[0].masks.xy | |
| num_notas = 0 | |
| for i, coords in enumerate(datos): | |
| # coordenadas de segmentación | |
| coordenadas = np.array(coords) | |
| # crear una máscara vacía | |
| mask = np.zeros((image.shape[0], image.shape[1]), dtype=np.uint8) | |
| # dibujar un polígono en la máscara | |
| cv2.fillPoly(mask, [coordenadas.astype(np.int32)], 255) | |
| # aplicar la máscara a la imagen | |
| masked = cv2.bitwise_and(image, image, mask=mask) | |
| # recortar la imagen utilizando la caja delimitadora de la máscara | |
| x, y, w, h = cv2.boundingRect(mask) | |
| recorte = masked[y:y+h, x:x+w] | |
| # guardar el recorte como una imagen | |
| recorte_path = os.path.join(f"recorte/{foldername}", f"nota {i}.jpg") | |
| cv2.imwrite(recorte_path, recorte) | |
| num_notas += 1 | |
| noticias(foldername) | |
| segmentacion = f"diarios/predict/{filename}.png" | |
| recorte1 = f"recorte/{foldername}/nota 0.jpg" | |
| prediccion = f"deteccion/{foldername}/predict/nota 0.jpg" | |
| try: | |
| recorte2 = f"recorte/{foldername}/nota 1.jpg" | |
| prediccion2 = f"deteccion/{foldername}/predict/nota 1.jpg" | |
| if not os.path.exists(recorte2): | |
| raise Exception(f"El archivo {recorte2} no existe.") | |
| except: | |
| recorte2 = "sin_nota.jpg" | |
| prediccion2 = "sin_nota.jpg" | |
| return cv2.imread(segmentacion), num_notas, recorte1, recorte2,prediccion, prediccion2 | |
| def noticias(carpeta): | |
| nombre = carpeta | |
| model = YOLO("best-detect.pt") | |
| results = model.predict(source=f"recorte/{nombre}", save=True, save_txt=True,project=f"deteccion/{nombre}", conf=0.75) # save plotted images | |
| return print(f"Imagen {nombre} procesada correctamente") | |
| gr.Interface( | |
| fn = recortar_notas, | |
| inputs = [ gr.components.Image(type="filepath", label="Input") ], | |
| outputs = [ gr.components.Image(type="filepath", label="Segmentacion"), gr.Label(label="Numero de Notas"), | |
| gr.components.Image(type="filepath", label="Recorte Nota"),gr.components.Image(type="filepath", label="Recorte Nota"), | |
| gr.components.Image(type="filepath", label="Deteccion"),gr.components.Image(type="filepath", label="Deteccion"), ], | |
| title="IA por la Identidad", | |
| description="Datathon Fundacion Sadosky", | |
| examples=[ ['clarin.jpg']], | |
| allow_flagging="never" | |
| ).launch(debug=True, enable_queue=True) |