214 lines
7.0 KiB
Python
214 lines
7.0 KiB
Python
import torch
|
|
import torch.optim as optim
|
|
import network.modeling as models
|
|
import numpy as np
|
|
from torch.utils.data import Dataset, DataLoader
|
|
from torchvision import transforms as T
|
|
from PIL import Image
|
|
import os
|
|
import argparse
|
|
import matplotlib.pyplot as plt
|
|
|
|
# Configuração do parsing de argumentos
|
|
parser = argparse.ArgumentParser(description='Treinamento do modelo DeepLabV3+')
|
|
parser.add_argument('--checkpoint', type=str, help='Caminho para o checkpoint de onde continuar o treinamento', default=None)
|
|
parser.add_argument('--n_epoch', type=str, help='Número de épocas para o treinamento', default='10')
|
|
args = parser.parse_args()
|
|
|
|
def load_class_map(labelmap_path):
|
|
class_map = {}
|
|
with open(labelmap_path, 'r') as file:
|
|
for line in file:
|
|
# Ignorar comentários e linhas vazias
|
|
if line.startswith('#') or not line.strip():
|
|
continue
|
|
# Extrair informações
|
|
parts = line.strip().split(':')
|
|
if len(parts) >= 2:
|
|
label, color_str = parts[0], parts[1]
|
|
# Converter a string de cor em uma tupla de inteiros
|
|
rgb = tuple(map(int, color_str.split(',')))
|
|
# Atualizar o class_map
|
|
class_map[rgb] = len(class_map)
|
|
return class_map
|
|
|
|
# Transformação personalizada para máscaras
|
|
def mask_transform(mask):
|
|
mask = np.array(mask)
|
|
class_map = load_class_map('dataset/labelmap.txt')
|
|
class_mask = np.zeros((mask.shape[0], mask.shape[1]), dtype=np.int32)
|
|
|
|
for rgb, idx in class_map.items():
|
|
match = (mask == rgb).all(axis=-1)
|
|
class_mask[match] = idx
|
|
|
|
return torch.tensor(class_mask, dtype=torch.long)
|
|
|
|
# Transformações para as imagens
|
|
image_transforms = T.Compose([
|
|
T.Resize((512, 512)),
|
|
T.ToTensor(),
|
|
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
|
|
])
|
|
|
|
# Dataset personalizado
|
|
class CustomVOC(Dataset):
|
|
def __init__(self, image_dirs, mask_dirs, transform=None, mask_transform=None):
|
|
self.transform = transform
|
|
self.mask_transform = mask_transform
|
|
self.images = []
|
|
self.masks = []
|
|
|
|
# Combine todos os arquivos de imagem dos diretórios fornecidos
|
|
for image_dir in image_dirs:
|
|
for img in os.listdir(image_dir):
|
|
self.images.append(os.path.join(image_dir, img))
|
|
|
|
# Combine todos os arquivos de máscara dos diretórios fornecidos
|
|
for mask_dir in mask_dirs:
|
|
for mask in os.listdir(mask_dir):
|
|
self.masks.append(os.path.join(mask_dir, mask.replace('jpg', 'jpeg'))) # Asumindo substituição de extensão
|
|
|
|
def __len__(self):
|
|
return len(self.images)
|
|
|
|
def __getitem__(self, idx):
|
|
img_path = self.images[idx]
|
|
mask_path = self.masks[idx] # Já ajustado na inicialização
|
|
image = Image.open(img_path).convert("RGB")
|
|
mask = Image.open(mask_path).convert("RGB")
|
|
|
|
if self.transform:
|
|
image = self.transform(image)
|
|
if self.mask_transform:
|
|
mask = self.mask_transform(mask)
|
|
|
|
return image, mask
|
|
|
|
# Instanciando o Dataset
|
|
# Altere os caminhos conforme necessário
|
|
image_dirs = ["dataset/images", "dataset/augmented_images"]
|
|
mask_dirs = ["dataset/masks", "dataset/augmented_masks"]
|
|
save_dir = "backup/"
|
|
model_name = "ruasModel"
|
|
train_dataset = CustomVOC(image_dirs=image_dirs, mask_dirs=mask_dirs, transform=image_transforms, mask_transform=mask_transform)
|
|
|
|
# DataLoader
|
|
train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True, num_workers=0)
|
|
|
|
# Definindo o modelo DeepLabV3+
|
|
def create_deeplabv3plus(_num_classes, _output_stride):
|
|
model = models.deeplabv3plus_resnet50(num_classes=_num_classes, output_stride=_output_stride)
|
|
return model
|
|
|
|
|
|
# Configuração do modelo, perda e otimizador
|
|
num_classes = 4
|
|
output_stride = 16
|
|
model = create_deeplabv3plus(num_classes, output_stride)
|
|
criterion = torch.nn.CrossEntropyLoss()
|
|
optimizer = optim.Adam(model.parameters(), lr=0.001)
|
|
|
|
# Preparando o modelo para treinamento
|
|
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
|
|
|
# Carregar checkpoint se fornecido
|
|
start_epoch = 0 # Iniciar do começo se nenhum checkpoint for fornecido
|
|
if args.checkpoint:
|
|
checkpoint = torch.load(args.checkpoint)
|
|
model.load_state_dict(checkpoint['model_state_dict'])
|
|
|
|
# Carrega o estado do otimizador
|
|
optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
|
|
|
|
# Garante que todos os tensores no otimizador estejam no dispositivo correto
|
|
for state in optimizer.state.values():
|
|
for k, v in state.items():
|
|
if isinstance(v, torch.Tensor):
|
|
state[k] = v.to(device)
|
|
|
|
start_epoch = checkpoint['epoch']
|
|
print(f"Continuando o treinamento do checkpoint {args.checkpoint}, a partir da época {start_epoch+1}")
|
|
|
|
model.to(device)
|
|
|
|
|
|
epoch_losses = []
|
|
epoch_accuracies = []
|
|
|
|
plt.ion() # Ativa o modo interativo
|
|
fig, ax1 = plt.subplots()
|
|
|
|
color_loss = 'tab:red'
|
|
ax1.set_xlabel('Época')
|
|
ax1.set_ylabel('Loss', color=color_loss)
|
|
ax1.tick_params(axis='y', labelcolor=color_loss)
|
|
|
|
ax2 = ax1.twinx() # Instancia um segundo eixo que compartilha o mesmo eixo x
|
|
color_accuracy = 'tab:blue'
|
|
ax2.set_ylabel('Accuracy', color=color_accuracy) # Definimos o label do eixo y
|
|
ax2.tick_params(axis='y', labelcolor=color_accuracy)
|
|
|
|
fig.tight_layout() # Ajusta o layout para evitar sobreposições
|
|
|
|
# Treinamento do modelo
|
|
num_epochs = int(args.n_epoch)
|
|
# Loop de treinamento, ajustado para continuar de onde parou
|
|
for epoch in range(start_epoch, num_epochs):
|
|
model.train()
|
|
running_loss = 0.0
|
|
for images, masks in train_loader:
|
|
images = images.to(device)
|
|
masks = masks.to(device)
|
|
|
|
optimizer.zero_grad()
|
|
|
|
out = model(images)
|
|
|
|
outputs = out
|
|
loss = criterion(outputs, masks)
|
|
loss.backward()
|
|
optimizer.step()
|
|
|
|
running_loss += loss.item()
|
|
|
|
correct = 0
|
|
total = 0
|
|
# Convertendo as previsões em rótulos preditos
|
|
_, predicted = torch.max(outputs.data, 1)
|
|
correct += (predicted == masks).sum().item()
|
|
total += masks.numel()
|
|
|
|
accuracy = 100 * correct / total
|
|
|
|
# Salvando checkpoints
|
|
checkpoint_path = f'{save_dir}{model_name}_checkpoint_epoch_{epoch+1}.pth'
|
|
torch.save({
|
|
'epoch': epoch+1,
|
|
'model_state_dict': model.state_dict(),
|
|
'optimizer_state_dict': optimizer.state_dict(),
|
|
'loss': running_loss,
|
|
}, checkpoint_path)
|
|
print(f"Checkpoint salvo em {checkpoint_path}")
|
|
|
|
epoch_losses.append(running_loss / len(train_loader))
|
|
epoch_accuracies.append(accuracy)
|
|
|
|
# Atualiza os dados do gráfico
|
|
ax1.plot(epoch_losses, color=color_loss)
|
|
ax2.plot(epoch_accuracies, color=color_accuracy)
|
|
|
|
# Redesenha o gráfico
|
|
fig.canvas.draw()
|
|
fig.canvas.flush_events()
|
|
|
|
print(f"Epoch {epoch+1}, Loss: {running_loss/len(train_loader)}, Accuracy: {accuracy}")
|
|
|
|
# Salvar o modelo final
|
|
save_path = f'{save_dir}{model_name}_final.pth'
|
|
torch.save(model.state_dict(), save_path)
|
|
print(f"Modelo salvo em {save_path}")
|
|
|
|
plt.ioff() # Desativa o modo interativo
|
|
plt.show()
|