Segmento_de_Angio_Coronariana_v6 / processar_imagem_de_entrada.py
DHEIVER's picture
Update processar_imagem_de_entrada.py
3d7a36a
raw
history blame
No virus
2.41 kB
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
# Define the "device" variable at the global scope
device = "cuda" if torch.cuda.is_available() else "cpu"
# Função para ordenar e pré-processar a imagem de entrada
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
# Função para processar a imagem de entrada
def processar_imagem_de_entrada(img, modelo, pipe, device, salvar_resultado=False):
try:
# Faça uma cópia da imagem original
img = img.copy()
# Coloque o modelo na GPU (se disponível) e configure-o para modo de avaliação
pipe = pipe.to(device).eval()
# Registre o tempo de início
start = time.time()
# Pré-processe a imagem e obtenha informações de dimensão
img, h, w, ori_gray, ori = ordenar_e_preprocessar_imagem(img, modelo)
# Converta a imagem para o formato esperado pelo modelo e coloque-a na GPU
img = torch.FloatTensor(img).unsqueeze(0).to(device)
# Realize a inferência do modelo sem gradientes
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)
# Calcule o tempo decorrido
spent = time.time() - start
spent = f"{spent:.3f} segundos"
# Redimensione o resultado, se necessário
if h != 512 or w != 512:
logit = cv2.resize(logit, (h, w))
# Converta o resultado para um formato booleano
logit = logit.astype(bool)
# Salve o resultado em um arquivo se a opção estiver ativada
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
# ... (rest of your code)