118 lines
4.2 KiB
Python
118 lines
4.2 KiB
Python
import torch
|
|
from PIL import Image
|
|
import torchvision.transforms as T
|
|
import matplotlib.pyplot as plt
|
|
import numpy as np
|
|
|
|
# Supondo que você tenha a estrutura do repositório e o módulo `network` conforme descrito no README
|
|
from network.modeling import deeplabv3_resnet101 as deeplabv3_model
|
|
|
|
# Configurações Iniciais
|
|
NUM_CLASSES = 21 # Pascal VOC possui 20 classes + 1 para o fundo
|
|
OUTPUT_STRIDE = 16 # Valor comum para DeepLab
|
|
MODEL_PATH = 'weights/best_deeplabv3_resnet101_voc_os16.pth' # Caminho para o modelo pré-treinado
|
|
|
|
# Função para carregar o modelo
|
|
def load_model(model_path):
|
|
model = deeplabv3_model(num_classes=NUM_CLASSES, output_stride=OUTPUT_STRIDE)
|
|
model.load_state_dict(torch.load(model_path)['model_state'])
|
|
model.eval() # Modo de avaliação
|
|
return model
|
|
|
|
# Carregar o modelo
|
|
model = load_model(MODEL_PATH)
|
|
|
|
# Função para processar uma imagem e realizar a segmentação
|
|
def segment_image(image_path, model):
|
|
image = Image.open(image_path).convert("RGB")
|
|
transform = T.Compose([
|
|
T.Resize(520),
|
|
T.ToTensor(),
|
|
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
|
])
|
|
|
|
input_tensor = transform(image).unsqueeze(0)
|
|
with torch.no_grad():
|
|
output = model(input_tensor)
|
|
output_predictions = output.max(1)[1].squeeze().detach().cpu().numpy()
|
|
return output_predictions
|
|
|
|
def decode_segmap(image, nc=21):
|
|
label_colors = np.array([
|
|
(0, 0, 0), # Fundo
|
|
(128, 0, 0), (0, 128, 0), (128, 128, 0), (0, 0, 128),
|
|
(128, 0, 128), (0, 128, 128), (128, 128, 128), (64, 0, 0),
|
|
(192, 0, 0), (64, 128, 0), (192, 128, 0), (64, 0, 128),
|
|
(192, 0, 128), (64, 128, 128), (192, 128, 128), (0, 64, 0),
|
|
(128, 64, 0), (0, 192, 0), (128, 192, 0), (0, 64, 128),
|
|
])
|
|
# Nomes das classes conforme Pascal VOC
|
|
class_names = [
|
|
"Fundo", "Avião", "Bicicleta", "Pássaro", "Barco",
|
|
"Garrafa", "Ônibus", "Carro", "Gato", "Cadeira",
|
|
"Vaca", "Mesa", "Cachorro", "Cavalo", "Moto",
|
|
"Pessoa", "Vaso de Planta", "Ovelha", "Sofá", "Trem", "Monitor/TV"
|
|
]
|
|
|
|
detected_classes = set(np.unique(image)) # Conjunto das classes detectadas na imagem
|
|
|
|
r = np.zeros_like(image).astype(np.uint8)
|
|
g = np.zeros_like(image).astype(np.uint8)
|
|
b = np.zeros_like(image).astype(np.uint8)
|
|
|
|
for l in detected_classes: # Usar apenas as classes detectadas
|
|
idx = image == l
|
|
r[idx] = label_colors[l, 0]
|
|
g[idx] = label_colors[l, 1]
|
|
b[idx] = label_colors[l, 2]
|
|
|
|
rgb = np.stack([r, g, b], axis=2)
|
|
return rgb, class_names, label_colors, detected_classes
|
|
|
|
def display_segmentation(input_image_path, output_predictions):
|
|
original_image = Image.open(input_image_path)
|
|
original_image = np.array(original_image)
|
|
|
|
seg_image, class_names, label_colors, detected_classes = decode_segmap(output_predictions)
|
|
|
|
seg_image_resized = Image.fromarray(seg_image).resize((original_image.shape[1], original_image.shape[0]), resample=Image.NEAREST)
|
|
seg_image_resized = np.array(seg_image_resized)
|
|
|
|
overlayed_img = original_image * 0.6 + seg_image_resized * 0.4
|
|
overlayed_img = overlayed_img.astype(np.uint8)
|
|
|
|
plt.figure(figsize=(20, 10))
|
|
|
|
# Imagem Original
|
|
plt.subplot(1, 3, 1)
|
|
plt.imshow(original_image)
|
|
plt.title('Imagem Original')
|
|
plt.axis('off')
|
|
|
|
# Representação da Segmentação
|
|
plt.subplot(1, 3, 2)
|
|
plt.imshow(seg_image_resized)
|
|
plt.title('Representação da Segmentação')
|
|
plt.axis('off')
|
|
|
|
# Imagem com Overlay
|
|
plt.subplot(1, 3, 3)
|
|
plt.imshow(overlayed_img)
|
|
plt.title('Imagem com Overlay')
|
|
plt.axis('off')
|
|
|
|
import matplotlib.patches as mpatches
|
|
patches = [mpatches.Patch(color=np.array(label_colors[i])/255.0, label=class_names[i]) for i in detected_classes]
|
|
plt.figlegend(handles=patches, loc='lower center', ncol=5, labelspacing=0.)
|
|
|
|
plt.subplots_adjust(bottom=0.25) # Ajuste conforme necessário para acomodar a legenda
|
|
plt.show()
|
|
|
|
|
|
image_path = 'images/img2.jpg'
|
|
# Exemplo de como usar a função `segment_image`
|
|
output_predictions = segment_image(image_path, model)
|
|
# Chamada da função
|
|
display_segmentation(image_path, output_predictions)
|
|
|