|
import os |
|
import logging |
|
import torch |
|
import cv2 |
|
import numpy as np |
|
from preprocess import unsharp_masking |
|
import time |
|
from sklearn.cluster import KMeans |
|
import os |
|
from preprocessar_imagem import preprocessar_imagem |
|
|
|
|
|
device = "cuda" if torch.cuda.is_available() else "cpu" |
|
|
|
|
|
def ordenar_e_preprocessar_imagem(img, modelo): |
|
ori = img.copy() |
|
img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) |
|
h, w = img.shape |
|
img_out = preprocessar_imagem(img, modelo) |
|
return img_out, h, w, img, ori |
|
|
|
|
|
def processar_imagem_de_entrada(img, modelo, pipe, device, salvar_resultado=False): |
|
try: |
|
|
|
img = img.copy() |
|
|
|
|
|
pipe = pipe.to(device).eval() |
|
|
|
|
|
start = time.time() |
|
|
|
|
|
img, h, w, ori_gray, ori = ordenar_e_preprocessar_imagem(img, modelo) |
|
|
|
|
|
img = torch.FloatTensor(img).unsqueeze(0).to(device) |
|
|
|
|
|
with torch.no_grad(): |
|
if modelo == 'AngioNet': |
|
img = torch.cat([img, img], dim=0) |
|
logit = np.round(torch.softmax(pipe.forward(img), dim=1).detach().cpu().numpy()[0, 0]).astype(np.uint8) |
|
|
|
|
|
spent = time.time() - start |
|
spent = f"{spent:.3f} segundos" |
|
|
|
|
|
if h != 512 or w != 512: |
|
logit = cv2.resize(logit, (h, w)) |
|
|
|
|
|
logit = logit.astype(bool) |
|
|
|
|
|
if salvar_resultado: |
|
nome_arquivo = os.path.join(caminho_salvar_resultado, f'resultado_{int(time.time())}.png') |
|
cv2.imwrite(nome_arquivo, logit) |
|
|
|
return spent, logit |
|
|
|
except Exception as e: |
|
logging.error(f"Error: {str(e)}") |
|
return str(e), None |
|
|
|
|
|
|