def __init__(self):
self.net = ENet(only_encode=True)
# # encoder_weight = torch.load('./model/encoder_ep_497_mIoU_0.5098.pth', map_location='cpu')
model_path = '/Users/yuxuanliu/Desktop/enet.pytorch/ckpt/kitti_checkpoint_19-02-19_22-04-01_encoder_ENet_city_[320, 640]_lr_0.0005.pth'
# model_path = '/Users/yuxuanliu/Desktop/enet.pytorch/ckpt/kitti_checkpoint_19-02-19_22-04-01_encoder_ENet_city_[320, 640]_lr_0.0005.pth'
encoder_weight = torch.load(model_path, map_location='cpu')
self.net.encoder.load_state_dict(encoder_weight)
mean_std = cfg.DATA.MEAN_STD
self.transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
class LaneNet_ENet_1E1D(nn.Module):
'''
A LaneNet model made up of ENet, please refer to 'Efficient Net'.
1E1D means one encoder and one decoder, that is the embedding and the binary segmentation
share the same encoder and decoder, except the last layer of the decoder.
'''
def __init__(self):
super().__init__()
self.model = ENet(64)
self.conv_logit = nn.Conv2d(64, 2, 1, 1)
self.conv_embedding = nn.Conv2d(64, 3, 1, 1)
def forward(self, input):
x = self.model.forward(input)
logit = self.conv_logit(x)
embedding = self.conv_embedding(x)
return embedding, logit
def get_model(name):
if name == 'hlnet':
model = HLNet(input_shape=(IMG_SIZE, IMG_SIZE, 3), cls_num=CLS_NUM)
elif name == 'fastscnn':
model = Fast_SCNN(num_classes=CLS_NUM,
input_shape=(IMG_SIZE, IMG_SIZE, 3)).model()
elif name == 'lednet':
model = LEDNet(groups=2,
classes=CLS_NUM,
input_shape=(IMG_SIZE, IMG_SIZE, 3)).model()
elif name == 'dfanet':
model = DFANet(input_shape=(IMG_SIZE, IMG_SIZE, 3),
cls_num=CLS_NUM,
size_factor=2)
elif name == 'enet':
model = ENet(input_shape=(IMG_SIZE, IMG_SIZE, 3), cls_num=CLS_NUM)
elif name == 'mobilenet':
model = MobileNet(input_shape=(IMG_SIZE, IMG_SIZE, 3), cls_num=CLS_NUM)
else:
raise NameError("No corresponding model...")
return model
def CreateModel(args):
if args.model == 'enet':
model = ENet(num_classes=args.num_classes).cuda()
optimizer = optim.Adam(
model.parameters(), lr=args.learning_rate, weight_decay=args.weight_decay)
elif args.model == 'deeplab':
model = Deeplab(num_classes=args.num_classes).cuda()
optimizer = None
elif args.model == 'frrnet':
model = FRRNet(out_channels=args.num_classes).cuda()
optimizer = optim.Adam(
model.parameters(), lr=args.learning_rate)
"""
elif args.model == 'fcn8s':
model = FCN8s(pretrained_net=None, n_class=args.num_classes).cuda()
optimizer = optim.SGD(
model.parameters(), lr=args.learning_rate)
"""
else:
class SegmentationModel:
WIDTH = 1240
HEIGHT = 376
def __init__(self):
self.net = ENet(only_encode=True)
# # encoder_weight = torch.load('./model/encoder_ep_497_mIoU_0.5098.pth', map_location='cpu')
model_path = '/Users/yuxuanliu/Desktop/enet.pytorch/ckpt/kitti_checkpoint_19-02-19_22-04-01_encoder_ENet_city_[320, 640]_lr_0.0005.pth'
# model_path = '/Users/yuxuanliu/Desktop/enet.pytorch/ckpt/kitti_checkpoint_19-02-19_22-04-01_encoder_ENet_city_[320, 640]_lr_0.0005.pth'
encoder_weight = torch.load(model_path, map_location='cpu')
self.net.encoder.load_state_dict(encoder_weight)
mean_std = cfg.DATA.MEAN_STD
self.transform = standard_transforms.Compose([
standard_transforms.ToTensor(),
standard_transforms.Normalize(*mean_std)
])
# self.test_model()
# def test_model(self):
# img = torch.zeros([5, 3, 100, 100], dtype=torch.float)
# seg = self.net.forward(img).data
def prepare_image(self, img):
pil_img = Image.fromarray(img.astype(np.uint8))
img_tensor = self.transform(pil_img)
img_tensor = img_tensor[:3, :SegmentationModel.
HEIGHT, :SegmentationModel.WIDTH]
img_tensor = img_tensor.unsqueeze(0)
return img_tensor
def segment_image(self, img):
img_tensor = self.prepare_image(img)
seg = self.net.forward(img_tensor)
max_vals, classes = torch.max(seg, 1)
classes = classes.squeeze(0)
print(torch.sum(classes == 9))
np_classes = classes.data.numpy()
color_mask = colorize_mask(np_classes)
seg_img = color_mask / 255
seg_img_bgr = seg_img[:, :, ::-1]
cv2.imshow('img', img)
cv2.imshow('seg', seg_img_bgr)
cv2.waitKey(1)
return np_classes, seg_img
def find_seg_class(self, class_mask, position):
mask_height, mask_width = class_mask.shape
sample_x = int(position[0] / SegmentationModel.WIDTH * mask_width)
sample_y = int(position[1] / SegmentationModel.HEIGHT * mask_height)
return class_mask[sample_y, sample_x]
def __init__(self):
super().__init__()
self.model = ENet(64)
self.conv_logit = nn.Conv2d(64, 2, 1, 1)
self.conv_embedding = nn.Conv2d(64, 3, 1, 1)
shuffle=True,
num_workers=num_workers)
val_loader = DataLoader(val_dataset,
batch_size=batch_size,
num_workers=num_workers)
test_loader = DataLoader(test_dataset,
batch_size=batch_size,
num_workers=num_workers)
"""Get Class Weighting"""
class_weights = torch.from_numpy(
weighing_class(train_loader, num_classes=12)).float().to(device)
unlabeled_idx = list(color_encoding).index('unlabeled')
class_weights[unlabeled_idx] = 0
print("class_weights : ", class_weights)
"""Get model,loss,optimizer,lr_scheduler"""
model = ENet(num_classes=12).to(device)
loss_fn = nn.CrossEntropyLoss(weight=class_weights)
optimizer = torch.optim.Adam(model.parameters(),
lr=lr,
weight_decay=weight_decay)
lr_scheduler = lr_scheduler.StepLR(optimizer, lr_decay_epochs, lr_decay)
"""Training"""
best_val_iou = 0
for epoch in range(epochs):
#Train
train_loss = train_batch(train_loader,
model,
loss_fn,
optimizer,
device=device)
lr_scheduler.step()
def prepare_network():
ENet_model = ENet(len(cityscapes_labels))
checkpoint = torch.load('model/ENet')
ENet_model.load_state_dict(checkpoint['state_dict'])
return ENet_model.eval().cuda()