def eval(config: dict, model_path='checkpoints/model_epoch30.pth'):
ssd = build_model(config, is_test=True)
ssd.load_state_dict(torch.load(model_path))
ssd.train(False)
net = Predictor(ssd)
data_transform = transforms.Compose([
transforms.ToRelativeBoxes(),
transforms.Resize(config['image_size']),
transforms.Scale(),
transforms.ToTensor()
])
target_transform = MatchPrior(priors, config)
val_set = CityscapesDataset(config, 'dataset/val', None, data_transform,
target_transform, True)
test_image = val_set.get_image(0)
boxes, labels, conf, _, _ = net.predict(test_image)
drawer = Draw(test_image)
for i in range(boxes.shape[0]):
top_left = tuple(boxes[i][:2].numpy().flatten())
bottom_right = tuple(boxes[i][2:].numpy().flatten())
drawer.rectangle((top_left, bottom_right))
test_image.save("predict.jpg")
def eval_disparity(config: dict, model_path='checkpoints/model_epoch30.pth'):
ssd = build_model(config, is_test=True)
ssd.load_state_dict(torch.load(model_path))
ssd.train(False)
net = Predictor(ssd)
data_transform = transforms.Compose([
transforms.ToRelativeBoxes(),
transforms.Resize(config['image_size']),
transforms.Scale(),
transforms.ToTensor()
])
target_transform = MatchPrior(priors, config)
val_set = CityscapesDataset(config, 'dataset/val', None, data_transform,
target_transform, True)
errors = []
for i, _ in enumerate(val_set):
image = val_set.get_image(i)
gt_disparity = val_set.get_disparity(i)
_, _, _, _, disparity = net.predict(image)
error = ((gt_disparity - disparity)**2).flatten()
errors.append(error)
errors = torch.cat(errors)
print("RMSE: {}".format(math.sqrt(errors.mean().item())))
def eval(config: dict, model_path='checkpoints/model_epoch40.pth'):
ssd = build_model(config, is_test=True)
ssd.load_state_dict(torch.load(model_path))
ssd.train(False)
net = Predictor(ssd)
data_transform = transforms.Compose([
transforms.ToRelativeBoxes(),
transforms.Resize(config['image_size']),
transforms.Scale(),
transforms.ToTensor()
])
target_transform = MatchPrior(priors, config)
val_set = CityscapesDataset(config, 'dataset/val', None, data_transform,
target_transform, True)
for i, _ in enumerate(val_set):
if i % 10 == 0:
print("Image {}".format(i))
image = val_set.get_image(i)
probs, boxes, disparity = net.predict(image)
labels = torch.argmax(probs, dim=probs.dim() - 1)
chosen_indices = []
for class_index in range(1, config['num_classes'] + 1):
class_mask = labels == class_index
# If there's no prediction in this class, skip the class
if class_mask.long().sum() <= 0:
continue
class_probabilities = probs[class_mask, class_index]
class_boxes = boxes[class_mask]
class_indices = nms(class_boxes, class_probabilities, 0.5)
chosen_indices.append(class_indices)
chosen_indices = torch.cat(chosen_indices)
probs = probs[chosen_indices]
boxes = boxes[chosen_indices]
box_drawer = Draw(image)
for box in boxes:
top_left = tuple(box[:2].numpy().flatten())
bottom_right = tuple(box[2:].numpy().flatten())
box_drawer.rectangle((top_left, bottom_right))
image.save('result.jpg')
# TODO change to all image evaluation
break
import numpy as np
import torch
from PIL import Image
from network import transforms
from torchvision.transforms import ToPILImage
transformer = transforms.Compose([
transforms.CustomJitter(),
transforms.RandomExpand(),
transforms.RandomCrop(),
transforms.RandomMirror(),
transforms.ToRelativeBoxes(),
transforms.Resize(1000),
transforms.Scale(),
transforms.ToTensor()
])
test_image = Image.open('../dataset/train/images/aachen_000000_000019_leftImg8bit.png')
test_boxes = np.random.randint(30, 60, (2, 4))
test_labels = np.random.randint(0, 10, (2, 1))
res, _, _ = transformer(test_image, torch.tensor(test_boxes), torch.tensor(test_labels))
tensor_to_img = ToPILImage()
res_img = tensor_to_img(res)
test_image.show()
res_img.show()
def train_ssd(start_epoch: int, end_epoch: int, config: dict, use_gpu: bool = True, model_name='model',
checkpoint_folder='checkpoints',
log_folder='log', redirect_output=True):
if not os.path.isdir(log_folder):
os.makedirs(log_folder)
priors = generate_priors(config)
target_transform = MatchPrior(priors, config)
train_transform = transforms.Compose([
transforms.CustomJitter(),
transforms.ToOpenCV(),
transforms.RandomMirror()
])
data_transform = transforms.Compose([
transforms.ToRelativeBoxes(),
transforms.Resize(config['image_size']),
transforms.Scale(),
transforms.ToTensor()
])
train_set = CityscapesDataset(config, 'dataset/train', train_transform,
data_transform, target_transform)
train_loader = DataLoader(train_set, batch_size=32,
shuffle=True, num_workers=4)
ssd = build_model(config)
if use_gpu:
ssd = ssd.cuda()
ssd.train(True)
if os.path.isfile(os.path.join(checkpoint_folder, "{}_epoch{}.pth".format(model_name, start_epoch - 1))):
ssd.load_state_dict(
torch.load(os.path.join(checkpoint_folder, "{}_epoch{}.pth".format(model_name, start_epoch - 1))))
criterion = MultiBoxLoss(0.5, 0, 3, config)
disparity_criterion = BerHuLoss()
ssd_params = [
{'params': ssd.extractor.parameters()},
{'params': ssd.extras.parameters()},
{'params': itertools.chain(ssd.class_headers.parameters(),
ssd.location_headers.parameters(),
ssd.upsampling.parameters())}
]
optimizer = SGD(ssd_params, lr=0.001, momentum=0.9, weight_decay=0.0005)
lr_scheduler = CosineAnnealingLR(optimizer, 120, last_epoch= -1)
if os.path.isfile(os.path.join(checkpoint_folder, "optimizer_epoch{}.pth".format(start_epoch - 1))):
print("Loading previous optimizer")
optimizer.load_state_dict(
torch.load(os.path.join(checkpoint_folder, "optimizer_epoch{}.pth".format(start_epoch - 1))))
for epoch in range(start_epoch, end_epoch):
lr_scheduler.step()
running_loss = 0.0
running_regression_loss = 0.0
running_classification_loss = 0.0
running_disparity_loss = 0.0
num_steps = len(train_loader)
aps = torch.zeros((config['num_classes'],))
running_map = 0
if redirect_output:
sys.stdout = open(os.path.join(log_folder, 'train_epoch_{}.txt'.format(epoch)), 'w')
for i, batch in enumerate(train_loader):
images, gt_locations, labels, gt_disparity = batch
if use_gpu:
images = images.cuda()
gt_locations = gt_locations.cuda()
labels = labels.cuda()
gt_disparity = gt_disparity.cuda()
optimizer.zero_grad()
confidences, locations, disparity = ssd(images)
regression_loss, classification_loss = criterion.forward(confidences, locations, labels, gt_locations)
disparity_loss = disparity_criterion.forward(disparity, gt_disparity)
loss = regression_loss + classification_loss + 10 * disparity_loss
loss.backward()
optimizer.step()
running_loss += loss.item()
running_regression_loss += regression_loss.item()
running_classification_loss += classification_loss.item()
running_disparity_loss += disparity_loss
with torch.no_grad():
boxes = convert_locations_to_boxes(locations, priors.cuda(), config['variance'][0],
config['variance'][1])
gt_boxes = convert_locations_to_boxes(gt_locations, priors.cuda(), config['variance'][0],
config['variance'][1])
batch_map, batch_ap = calculate_map(confidences, labels, boxes, gt_boxes)
running_map += batch_map
aps += batch_ap
avg_loss = running_loss / num_steps
avg_reg_loss = running_regression_loss / num_steps
avg_class_loss = running_classification_loss / num_steps
avg_disp_loss = running_disparity_loss / num_steps
mean_ap = running_map / num_steps
epoch_ap = aps / num_steps
print("Epoch {}".format(epoch))
print("Average Loss: {:.2f}".format(avg_loss))
print("Average Regression Loss: {:.2f}".format(avg_reg_loss))
print("Average Classification Loss: {:.2f}".format(avg_class_loss))
print("Average Disparity Loss: {:.2f}".format(avg_disp_loss))
print("Average mAP: {:.2f}".format(mean_ap))
print("Average AP per class: {}".format(epoch_ap))
torch.save(ssd.state_dict(), os.path.join(checkpoint_folder, "{}_epoch{}.pth".format(model_name, epoch)))
torch.save(optimizer.state_dict(), os.path.join(checkpoint_folder, "optimizer_epoch{}.pth".format(epoch)))
if sys.stdout != sys.__stdout__:
sys.stdout.close()
sys.stdout = sys.__stdout__