def main():
enet = ENet(12)
device = torch.device('cuda:0')
enet = enet.to(device)
# enet.load_state_dict(torch.load('../content/ckpt-enet-CamVid-downloaded.pth')['state_dict'])
# enet.load_state_dict(torch.load('./content/ckpt-enet-100-21.788276344537735.pth')['state_dict'])
enet.load_state_dict(
torch.load('./checkpoints/ckpt-enet-50-314.6487820148468.pth')
['state_dict'])
# testDir = '../content/test/'
# testannotDir = '../content/testannot/'
# imageFileName = 'Seq05VD_f05100.png'
testDir = './content/train/'
testannotDir = './content/trainannot/'
imageFileName = '0006R0_f01260.png'
inputImg = plt.imread(testDir + imageFileName) # read in RGB
inputImg = cv2.resize(inputImg, (512, 512), cv2.INTER_NEAREST)
inputTensor = torch.tensor(inputImg).unsqueeze(0).float()
inputTensor = inputTensor.transpose(2, 3).transpose(1, 2).to(device)
with torch.no_grad():
outTensor = enet(inputTensor.float()).squeeze(0)
# Индекс слоя с максимальным значением в соотв. пикселе и является индексом класса для этого пикселя
# Альтернатива: outTensor.data.max(0)[1].cpu().numpy()
computedLabels = outTensor.data.argmax(0).cpu().numpy()
decodedComputedLabels = decoder(computedLabels)
trueLabels = cv2.imread(testannotDir + imageFileName, cv2.IMREAD_GRAYSCALE)
trueLabels = cv2.resize(trueLabels, (512, 512), cv2.INTER_NEAREST)
decodedTrueLabels = decoder(trueLabels)
figure = plt.figure(figsize=(20, 10))
plt.subplot(1, 3, 1) # rows, cols, index
plt.title('Input Image')
plt.axis('off')
plt.imshow(inputImg)
plt.subplot(1, 3, 2)
plt.title('Output Labels')
plt.axis('off')
plt.imshow(decodedComputedLabels)
plt.subplot(1, 3, 3)
plt.title('True Labels')
plt.axis('off')
plt.imshow(decodedTrueLabels)
plt.show()
def test(FLAGS):
# Check if the pretrained model is available
if not FLAGS.m.endswith('.pth'):
raise RuntimeError('Unknown file passed. Must end with .pth')
if FLAGS.image_path is None or not os.path.exists(FLAGS.image_path):
raise RuntimeError('An image file path must be passed')
h = FLAGS.resize_height
w = FLAGS.resize_width
checkpoint = torch.load(FLAGS.m, map_location=FLAGS.cuda)
# Assuming the dataset is camvid
enet = ENet(12)
enet.load_state_dict(checkpoint['state_dict'])
tmg_ = plt.imread(FLAGS.image_path)
tmg_ = cv2.resize(tmg_, (h, w), cv2.INTER_NEAREST)
tmg = torch.tensor(tmg_).unsqueeze(0).float()
tmg = tmg.transpose(2, 3).transpose(1, 2)
###### MINE ######
device = FLAGS.cuda
tmg = tmg.to(device)
enet = enet.to(device)
###### END MINE #####
with torch.no_grad():
out1 = enet(tmg.float()).squeeze(0)
#smg_ = Image.open('/content/training/semantic/' + fname)
#smg_ = cv2.resize(np.array(smg_), (512, 512), cv2.INTER_NEAREST)
b_ = out1.data.max(0)[1].cpu().numpy()
decoded_segmap = decode_segmap(b_)
images = {
0: ['Input Image', tmg_],
1: ['Predicted Segmentation', b_],
}
show_images(images)
def test(FLAGS):
# Check if the pretrained model is available
if not FLAGS.m.endswith('.pth'):
raise RuntimeError('Unknown file passed. Must end with .pth')
if FLAGS.image_path is None or not os.path.exists(FLAGS.image_path):
raise RuntimeError('An image file path must be passed')
h = FLAGS.resize_height
w = FLAGS.resize_width
nc = FLAGS.num_classes
test_mode = FLAGS.test_mode
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")##dawson add gpu mode
checkpoint = torch.load(FLAGS.m, map_location=lambda storage, loc: storage.cuda(0))#"cpu")
# Assuming the dataset is camvid
enet = ENet(nc)
enet.to(device) ##dawson add gpu mode
enet.load_state_dict(checkpoint['state_dict'])
tmg_ = plt.imread(FLAGS.image_path)
tmg_ = cv2.resize(tmg_, (h, w), cv2.INTER_NEAREST)
tmg = torch.tensor(tmg_).unsqueeze(0).float()
tmg = tmg.transpose(2, 3).transpose(1, 2)
with torch.no_grad():
out1 = enet(tmg.cuda().float()).squeeze(0)##dawson add gpu mode
#smg_ = Image.open('/content/training/semantic/' + fname)
#smg_ = cv2.resize(np.array(smg_), (512, 512), cv2.INTER_NEAREST)
b_ = out1.data.max(0)[1].cpu().numpy()
if test_mode.lower() == 'cityscapes':
decoded_segmap = decode_segmap(b_, True)
else:
decoded_segmap = decode_segmap(b_, False)
cv2.imshow("test", decoded_segmap)
cv2.waitKey()
images = {
0 : ['Input Image', tmg_],
1 : ['Predicted Segmentation', b_],
}
show_images(images)
def train(FLAGS):
# Defining the hyperparameters
device = FLAGS.cuda
batch_size = FLAGS.batch_size
epochs = FLAGS.epochs
lr = FLAGS.learning_rate
print_every = FLAGS.print_every
eval_every = FLAGS.eval_every
save_every = FLAGS.save_every
nc = FLAGS.num_classes
wd = FLAGS.weight_decay
ip = FLAGS.input_path_train
lp = FLAGS.label_path_train
ipv = FLAGS.input_path_val
lpv = FLAGS.label_path_val
print('[INFO]Defined all the hyperparameters successfully!')
# Get the class weights
print('[INFO]Starting to define the class weights...')
pipe = loader(ip, lp, batch_size='all')
class_weights = get_class_weights(pipe, nc)
print('[INFO]Fetched all class weights successfully!')
# Get an instance of the model
enet = ENet(nc)
print('[INFO]Model Instantiated!')
# Define the criterion and the optimizer
criterion = nn.CrossEntropyLoss(
weight=torch.FloatTensor(class_weights).to(device))
optimizer = torch.optim.Adam(enet.parameters(), lr=lr, weight_decay=wd)
print('[INFO]Defined the loss function and the optimizer')
# Training Loop starts
print('[INFO]Staring Training...')
print()
train_losses = []
eval_losses = []
# Assuming we are using the CamVid Dataset
bc_train = 367 // batch_size
bc_eval = 101 // batch_size
pipe = loader(ip, lp, batch_size)
eval_pipe = loader(ipv, lpv, batch_size)
epochs = epochs
for e in range(1, epochs + 1):
train_loss = 0
print('-' * 15, 'Epoch %d' % e, '-' * 15)
enet.train()
for _ in tqdm(range(bc_train)):
X_batch, mask_batch = next(pipe)
#assert (X_batch >= 0. and X_batch <= 1.0).all()
X_batch, mask_batch = X_batch.to(device), mask_batch.to(device)
optimizer.zero_grad()
out = enet(X_batch.float())
loss = criterion(out, mask_batch.long())
loss.backward()
optimizer.step()
train_loss += loss.item()
print()
train_losses.append(train_loss)
if (e + 1) % print_every == 0:
print('Epoch {}/{}...'.format(e, epochs),
'Loss {:6f}'.format(train_loss))
if e % eval_every == 0:
with torch.no_grad():
enet.eval()
eval_loss = 0
for _ in tqdm(range(bc_eval)):
inputs, labels = next(eval_pipe)
inputs, labels = inputs.to(device), labels.to(device)
out = enet(inputs)
out = out.data.max(1)[1]
loss = criterion(out, labels.long())
eval_loss += loss.item()
print()
print('Loss {:6f}'.format(eval_loss))
eval_losses.append(eval_loss)
if e % save_every == 0:
checkpoint = {'epochs': e, 'state_dict': enet.state_dict()}
torch.save(checkpoint,
'./ckpt-enet-{}-{}.pth'.format(e, train_loss))
print('Model saved!')
print('Epoch {}/{}...'.format(e + 1, epochs),
'Total Mean Loss: {:6f}'.format(sum(train_losses) / epochs))
def train(FLAGS):
# Defining the hyperparameters
device = FLAGS.cuda
batch_size = FLAGS.batch_size
epochs = FLAGS.epochs
lr = FLAGS.learning_rate
print_every = FLAGS.print_every
eval_every = FLAGS.eval_every
save_every = FLAGS.save_every
nc = FLAGS.num_classes
wd = FLAGS.weight_decay
ip = FLAGS.input_path_train
lp = FLAGS.label_path_train
ipv = FLAGS.input_path_val
lpv = FLAGS.label_path_val
train_mode = FLAGS.train_mode
pretrain_model = FLAGS.pretrain_model
cityscapes_path = FLAGS.cityscapes_path
resume_model_path = FLAGS.resume_model_path
print('[INFO]Defined all the hyperparameters successfully!')
# Get the class weights
print('[INFO]Starting to define the class weights...')
if len(cityscapes_path):
pipe = loader_cityscapes(ip, cityscapes_path, batch_size='all')
class_weights = get_class_weights(pipe, nc, isCityscapes=True)
#class_weights = np.array([3.03507951, 13.09507946, 4.54913664, 37.64795738, 35.78537802, 31.50943831, 45.88744201, 39.936759,
# 6.05101481, 31.85754823, 16.92219283, 32.07766734, 47.35907214, 11.34163794, 44.31105748, 45.81085476,
# 45.67260936, 48.3493813, 42.02189188])
else:
pipe = loader(ip, lp, batch_size='all')
class_weights = get_class_weights(pipe, nc)
print('[INFO]Fetched all class weights successfully!')
# Get an instance of the model
if train_mode.lower() == 'encoder-decoder':
enet = ENet(nc)
if len(pretrain_model):
checkpoint0 = torch.load(pretrain_model)
pretrain_dict = checkpoint0['state_dict']
enet_dict = enet.state_dict()
pretrain_dict = {
k: v
for k, v in pretrain_dict.items() if k in enet_dict
}
enet_dict.update(pretrain_dict)
enet.load_state_dict(enet_dict)
print('[INFO]Previous model Instantiated!')
else:
enet = ENet_encoder(nc)
print('[INFO]Model Instantiated!')
# Move the model to cuda if available
enet = enet.to(device)
# Define the criterion and the optimizer
if len(cityscapes_path):
criterion = nn.CrossEntropyLoss(
weight=torch.FloatTensor(class_weights).to(device),
ignore_index=255)
else:
criterion = nn.CrossEntropyLoss(
weight=torch.FloatTensor(class_weights).to(device))
optimizer = torch.optim.Adam(enet.parameters(), lr=lr, weight_decay=wd)
scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer,
mode="min",
factor=0.5,
patience=2,
verbose=True,
threshold=0.01)
# scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode="min", factor=0.5, patience=2, verbose=True,
# threshold=0.005)
print('[INFO]Defined the loss function and the optimizer')
# Training Loop starts
print('[INFO]Staring Training...')
print()
train_losses = []
eval_losses = []
if len(cityscapes_path):
# Assuming we are using the Cityscapes Dataset
bc_train = 2975 // batch_size
bc_eval = 500 // batch_size
pipe = loader_cityscapes(ip, cityscapes_path, batch_size)
eval_pipe = loader_cityscapes(ipv, cityscapes_path, batch_size)
else:
# Assuming we are using the CamVid Dataset
bc_train = 367 // batch_size
bc_eval = 101 // batch_size
pipe = loader(ip, lp, batch_size)
eval_pipe = loader(ipv, lpv, batch_size)
epoch = 1
if len(resume_model_path):
checkpoint1 = torch.load(resume_model_path)
epoch = checkpoint1['epochs'] + 1
enet.load_state_dict(checkpoint1['state_dict'])
epochs = epochs
for e in range(epoch, epochs + 1):
train_loss = 0
print('-' * 15, 'Epoch %d' % e, '-' * 15)
enet.train()
for _ in tqdm(range(bc_train)):
X_batch, mask_batch = next(pipe)
#assert (X_batch >= 0. and X_batch <= 1.0).all()
X_batch, mask_batch = X_batch.to(device), mask_batch.to(device)
optimizer.zero_grad()
out = enet(X_batch.float())
loss = criterion(out, mask_batch.long())
loss.backward()
optimizer.step()
train_loss += loss.item()
print()
train_losses.append(train_loss)
if (e + 1) % print_every == 0:
print('Epoch {}/{}...'.format(e, epochs),
'Loss {:6f}'.format(train_loss))
scheduler.step(train_loss)
if e % eval_every == 0:
with torch.no_grad():
enet.eval()
eval_loss = 0
for _ in tqdm(range(bc_eval)):
inputs, labels = next(eval_pipe)
inputs, labels = inputs.to(device), labels.to(device)
out = enet(inputs)
loss = criterion(out, labels.long())
eval_loss += loss.item()
print()
print('Loss {:6f}'.format(eval_loss))
eval_losses.append(eval_loss)
if e % save_every == 0:
checkpoint = {'epochs': e, 'state_dict': enet.state_dict()}
if train_mode.lower() == 'encoder-decoder':
torch.save(
checkpoint, './logs/ckpt-enet-{}-{}-{}.pth'.format(
e,
optimizer.state_dict()['param_groups'][0]['lr'],
train_loss))
else:
torch.save(
checkpoint, './logs/ckpt-enet_encoder-{}-{}-{}.pth'.format(
e,
optimizer.state_dict()['param_groups'][0]['lr'],
train_loss))
print('Model saved!')
print('Epoch {}/{}...'.format(e + 1, epochs),
'Total Mean Loss: {:6f}'.format(sum(train_losses) / epochs))
print('[INFO]Training Process complete!')
def main():
device = torch.device('cuda:0')
enet = ENet(12).to(device)
lr = 5e-4
batch_size = 20
class_weights = get_class_weights(12)
criterion = nn.CrossEntropyLoss(
weight=torch.FloatTensor(class_weights).to(device))
optimizer = torch.optim.Adam(enet.parameters(), lr=lr, weight_decay=2e-4)
print_every = 5
eval_every = 5
# ## Training loop
# In[17]:
train_losses = []
eval_losses = []
bc_train = 367 // batch_size # mini_batch train
bc_eval = 101 // batch_size # mini_batch validation
iterationsPerEpochs = 1000
# Define pipeline objects
pipe = loader('./content/train/', './content/trainannot/', batch_size)
eval_pipe = loader('./content/val/', './content/valannot/', batch_size)
epochs = 100
print()
for e in range(1, epochs + 1):
train_loss = 0
print('-' * 15, 'Epoch %d' % e, '-' * 15)
enet.train()
for _ in tqdm(range(iterationsPerEpochs)):
X_batch, mask_batch = next(pipe)
# assign data to cpu/gpu
X_batch, mask_batch = X_batch.to(device), mask_batch.to(device)
optimizer.zero_grad()
out = enet(X_batch.float())
# loss calculation
loss = criterion(out, mask_batch.long())
# update weights
loss.backward()
optimizer.step()
train_loss += loss.item()
print()
train_losses.append(train_loss)
if (e + 1) % print_every == 0:
print(f'Epoch {e}/{epochs}...', f'Loss {train_loss:6f}')
if e % eval_every == 0:
with torch.no_grad():
enet.eval()
eval_loss = 0
# Validation loop
for _ in tqdm(range(bc_eval)):
inputs, labels = next(eval_pipe)
inputs, labels = inputs.to(device), labels.to(device)
out = enet(inputs)
out = out.data.max(1)[1]
eval_loss += (labels.long() - out.long()).sum()
print()
print(f'Loss {eval_loss:6f}')
eval_losses.append(eval_loss)
if e % print_every == 0:
checkpoint = {'epochs': e, 'state_dict': enet.state_dict()}
torch.save(checkpoint,
f'./checkpoints/ckpt-enet-{e}-{train_loss}.pth')
print('Model saved!')
print(f'Epoch {e}/{epochs}...',
f'Total Mean Loss: {sum(train_losses) / epochs:6f}')
import cv2
import torch
import torch.nn as nn
from utils import decode_segmap
from models.ENet import ENet
if __name__ == '__main__':
enet = ENet(12)
device=torch.device('cuda:0' if torch.cuda.is_available() \
else 'cpu')
checkpoint = torch.load('./ckpt-enet-100-35.71956396102905.pth')
enet.load_state_dict(checkpoint['state_dict'])
enet = enet.to(device)
cam1 = cv2.VideoCapture(0)
#cv2.namedWindow("test")
img_counter = 0
while True:
ret, frame = cam1.read()
if not ret:
print("failed to grab frame")
break
cv2.imshow("test", frame)
def train(FLAGS,frozen_layers=27):
# Defining the hyperparameters
device = FLAGS.cuda
m = FLAGS.m
batch_size = FLAGS.batch_size
epochs = FLAGS.epochs
lr = FLAGS.learning_rate
print_every = FLAGS.print_every
eval_every = FLAGS.eval_every
save_every = FLAGS.save_every
nc = FLAGS.num_classes
wd = FLAGS.weight_decay
ip = FLAGS.input_path_train
lp = FLAGS.label_path_train
ipv = FLAGS.input_path_val
lpv = FLAGS.label_path_val
print ('[INFO]Defined all the hyperparameters successfully!')
"""
# Get the class weights
print ('[INFO]Starting to define the class weights...')
pipe = loader(ip, lp, batch_size='all')
class_weights = get_class_weights(pipe, nc)
print ('[INFO]Fetched all class weights successfully!')
"""
# Get an instance of the model
enet = ENet(nc)
print ('[INFO]Model Instantiated!')
# Transfer learnt weights
pretrained_dict = torch.load(FLAGS.m, map_location=FLAGS.cuda)['state_dict']
model_dict = enet.state_dict()
# 1. filter out unnecessary keys
pretrained_dict = {k: v for k, v in pretrained_dict.items() if k in model_dict}
# 2. overwrite entries in the existing state dict
model_dict.update(pretrained_dict)
# 3. load the new state dict
enet.load_state_dict(model_dict)
#enet.load_state_dict(pretrained_dict)
# Choose frozen layers
count=0
for child in enet.children():
if count<frozen_layers:
for param in child.parameters():
param.requires_grad=False
count+=1
else:
for param in child.parameters():
print(param)
# Move the model to cuda if available
enet = enet.to(device)
# Define the criterion and the optimizer
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(enet.parameters(),lr=lr,weight_decay=wd)
print ('[INFO]Defined the loss function and the optimizer')
# Training Loop starts
print ('[INFO]Staring Training...')
print ()
train_losses = []
eval_losses = []
# Assuming we are using the CamVid Dataset
bc_train = 367 // batch_size
bc_eval = 101 // batch_size
pipe = loader(ip, lp, batch_size)
eval_pipe = loader(ipv, lpv, batch_size)
epochs = epochs
for e in range(1, epochs+1):
train_loss = 0
print ('-'*15,'Epoch %d' % e, '-'*15)
enet.train()
for _ in tqdm(range(bc_train)):
X_batch, mask_batch = next(pipe)
#assert (X_batch >= 0. and X_batch <= 1.0).all()
X_batch, mask_batch = X_batch.to(device), mask_batch.to(device)
optimizer.zero_grad()
out = enet(X_batch.float())
loss = criterion(out, mask_batch.long())
loss.backward()
optimizer.step()
train_loss += loss.item()
print ()
train_losses.append(train_loss)
if (e+1) % print_every == 0:
print ('Epoch {}/{}...'.format(e, epochs),
'Loss {:6f}'.format(train_loss))
if e % eval_every == 0:
with torch.no_grad():
enet.eval()
eval_loss = 0
for _ in tqdm(range(bc_eval)):
inputs, labels = next(eval_pipe)
inputs, labels = inputs.to(device), labels.to(device)
out = enet(inputs)
loss = criterion(out, labels.long())
eval_loss += loss.item()
print ()
print ('Loss {:6f}'.format(eval_loss))
eval_losses.append(eval_loss)
if e % save_every == 0:
checkpoint = {
'epochs' : e,
'state_dict' : enet.state_dict()
}
torch.save(checkpoint, './ckpt-new-transfer-enet-{}.pth'.format(e))
print ('Model saved!')
print ('Epoch {}/{}...'.format(e+1, epochs),
'Total Mean Loss: {:6f}'.format(sum(train_losses) / epochs))
print ('[INFO]Training Process complete!')