├── README.md ├── data_loader.py ├── evaluation.py ├── main.py ├── models ├── classification_modern_model.py └── segmentation_unet.py └── utils.py /README.md: -------------------------------------------------------------------------------- 1 | ## Classification and Segmentation of Longitudinal Road Marking using Convolutional Neural Networks for Dynamic Retroreflection Estimation 2 | 3 | ### Requirements 4 | glob2==0.7 5 | numpy==1.16.5 6 | pandas==1.0.4 7 | torch==1.2.0 8 | torchsummary==1.5.1 9 | torchvision==0.4.0a0+6b959ee 10 | tqdm==4.32.1 11 | 12 | ### Abstract 13 | The dynamic retroreflection estimation method of the longitudinal road marking from the luminance camera using convolutional neural networks (CNNs). From the image captured by the luminance camera, a classification and regression CNN model is proposed to find out whether the longitudinal road marking was accurately acquired. If the longitudinal road marking is present in the captured image, a segmentation model that appropriately presents the longitudinal road marking and the reference plate is also proposed. 14 | 15 | ### Overall Structure 16 | First, it is determined whether retroreflection prediction is possible from the luminance image. If possible, the luminance image is cropped as the small image to find the retroreflection. After that, the reference plate and the longitudinal road marking are precisely segmented. 17 | 18 | 19 | ### Dynamic Retroreflection Estimation 20 | The dynamic retroreflection was predicted during actual drivingon the road. Note that when the classification model is determined to be negative, the ratio is 0. If the captured luminance image has a blue frame, it is a result predicted as positive, and if it is a red frame, it is a result predicted as negative. It is possible to confirm that only the retroreflection of the longituidnal road marking was extracted. 21 | ![fig5_1](https://user-images.githubusercontent.com/23445222/92069511-7e914f00-ede4-11ea-9b80-fd0438bd6d3b.png) 22 | 23 | ### Acknowledgement 24 | This research was supported by a grant from Technology Business Innovation Program (TBIP) funded by Ministry of Land, Infrastructure and Transport of Korean government (No. 18TBIP-C144255-01) [Development of Road Damage Information Technology based on Artificial Intelligence]. -------------------------------------------------------------------------------- /data_loader.py: -------------------------------------------------------------------------------- 1 | import pandas as pd 2 | import glob 3 | import tqdm 4 | import numpy as np 5 | import cv2 6 | 7 | import torch 8 | from torch.utils.data import Dataset 9 | from torchvision import transforms 10 | from torch.utils.data import DataLoader 11 | 12 | class ImagevDatasetForClassi(Dataset): 13 | def __init__(self, mode='train', transform=None, width=640, height=480): 14 | self.mode = mode 15 | self.transform = transform 16 | self.width = width 17 | self.height = height 18 | 19 | path = '../db/' 20 | if self.mode == 'train': 21 | path += 'train/' 22 | else: 23 | path += 'val/' 24 | 25 | # positive files 26 | self.positive_imgfiles = sorted(glob.glob(path + '*/ok/*.csv')) 27 | #self.positive_hfiles = sorted(glob.glob(path + '*/npy/*.npy')) 28 | 29 | # negative files 30 | self.negative_imgfiles = sorted(glob.glob(path + '*/not_ok/*.csv')) 31 | 32 | self.x_data, self.y_data = [], [] 33 | for imgfile in tqdm.tqdm(self.positive_imgfiles): 34 | npyfile = imgfile.replace('/ok', '/npy') 35 | npyfile = npyfile.replace('csv', 'npy') 36 | df_org = pd.read_csv(imgfile) 37 | ok_h = np.load(npyfile) 38 | self.x_data.append(df_org.values[:,1:641]) 39 | self.y_data.append([1, ok_h / 480.0]) 40 | 41 | for imgfile in tqdm.tqdm(self.negative_imgfiles): 42 | df_org = pd.read_csv(imgfile) 43 | self.x_data.append(df_org.values[:,1:641]) 44 | self.y_data.append([0, 0]) 45 | 46 | self.x_data = np.array(self.x_data) 47 | self.y_data = np.array(self.y_data) 48 | print(self.x_data.shape, self.y_data.shape) 49 | 50 | def __len__(self): 51 | return len(self.x_data) 52 | 53 | def __getitem__(self, idx): 54 | sample = torch.FloatTensor(self.x_data[idx]).view(1, self.height, self.width), torch.as_tensor(self.y_data[idx][0], dtype=torch.long), torch.as_tensor(self.y_data[idx][1]) 55 | if self.transform: 56 | sample = self.transform(sample) 57 | 58 | return sample 59 | 60 | class ImagevDatasetForSegment(Dataset): 61 | def __init__(self, mode='train', transform=None, width=640, height=192): 62 | self.mode = mode 63 | self.transform = transform 64 | self.width = width 65 | self.height = height 66 | 67 | path = '../db/' 68 | if self.mode == 'train': 69 | path += 'train/' 70 | else: 71 | path += 'val/' 72 | 73 | # positive files 74 | self.positive_imgfiles = sorted(glob.glob(path + '*/ok/*.csv')) 75 | self.positive_lblfiles = sorted(glob.glob(path + '*/labeled/*.png')) 76 | 77 | self.x_data = [] 78 | for imgfile in tqdm.tqdm(self.positive_imgfiles): 79 | npyfile = imgfile.replace('/ok', '/npy') 80 | npyfile = npyfile.replace('csv', 'npy') 81 | df_org = pd.read_csv(imgfile) 82 | ok_h = np.load(npyfile) 83 | self.x_data.append(df_org.values[ok_h:ok_h+self.height, 1:641]) 84 | 85 | self.y_data = [] 86 | for lblfile in tqdm.tqdm(self.positive_lblfiles): 87 | img = np.array(cv2.imread(lblfile, cv2.IMREAD_GRAYSCALE)) 88 | #img = img.reshape(1, self.height, self.width) 89 | self.y_data.append(img) 90 | 91 | self.x_data = np.array(self.x_data) 92 | self.y_data = np.array(self.y_data) 93 | print(self.x_data.shape, self.y_data.shape) 94 | 95 | def __len__(self): 96 | return len(self.x_data) 97 | 98 | def __getitem__(self, idx): 99 | sample = torch.FloatTensor(self.x_data[idx]).view(1, self.height, self.width), torch.LongTensor(self.y_data[idx]).view(self.height, self.width) 100 | if self.transform: 101 | sample = self.transform(sample) 102 | 103 | return sample 104 | 105 | 106 | class ImagevDatasetForRetroEsti(Dataset): 107 | def __init__(self, mode='val2', transform=None, width=640, height=480): 108 | self.mode = mode 109 | self.transform = transform 110 | self.width = width 111 | self.height = height 112 | 113 | path = '../db/' 114 | path += 'val2/' 115 | 116 | # iamge files 117 | self.imgfiles = sorted(glob.glob(path + 'pixels/*.csv')) 118 | 119 | self.x_data, self.y_data = [], [] 120 | for imgfile in tqdm.tqdm(self.imgfiles): 121 | jpgfile = imgfile.replace('pixels', 'original') 122 | jpgfile = jpgfile.replace('csv', 'jpg') 123 | 124 | df_org = pd.read_csv(imgfile) 125 | self.x_data.append(df_org.values[:,1:641]) 126 | img = np.array(cv2.imread(jpgfile)) 127 | self.y_data.append(img) 128 | 129 | self.x_data = np.array(self.x_data) 130 | self.y_data = np.array(self.y_data) 131 | print(self.x_data.shape, self.y_data.shape) 132 | 133 | def __len__(self): 134 | return len(self.x_data) 135 | 136 | def __getitem__(self, idx): 137 | sample = torch.FloatTensor(self.x_data[idx]).view(1, self.height, self.width), torch.LongTensor(self.y_data[idx]).view(3, self.height, self.width) 138 | if self.transform: 139 | sample = self.transform(sample) 140 | 141 | return sample 142 | 143 | 144 | if __name__ == "__main__": 145 | ''' 146 | train_dataset = ImagevDatasetForClassi(mode='val') 147 | train_loader = DataLoader(train_dataset, batch_size=4, shuffle=False, num_workers=0) 148 | for batch_idx, samples in enumerate(train_loader): 149 | data, target1, target2 = samples 150 | print(data.shape, target1.shape, target2.shape) 151 | 152 | break 153 | ''' 154 | 155 | train_dataset = ImagevDatasetForRetroEsti(mode='val2') 156 | train_loader = DataLoader(train_dataset, batch_size=1, shuffle=False, num_workers=0) 157 | for batch_idx, samples in enumerate(train_loader): 158 | data, target = samples 159 | print(data.shape, target.shape) 160 | print(torch.max(target), torch.min(target)) 161 | 162 | break 163 | -------------------------------------------------------------------------------- /evaluation.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | 4 | def compute_iou(pred, target, n_classes=3): 5 | ious = [] 6 | pred = torch.argmax(pred, dim=1).view(-1) 7 | target = target.view(-1) 8 | 9 | for cls in range(n_classes): 10 | y_true = np.multiply(target.eq(cls).detach().cpu().numpy(), 1) 11 | y_pred = np.multiply(pred.eq(cls).detach().cpu().numpy(), 1) 12 | 13 | intersection = np.sum(y_true * y_pred) 14 | union = np.sum(y_true) + np.sum(y_pred) - intersection 15 | 16 | iou = intersection / (union + 1e-6) 17 | ious.append(iou) 18 | 19 | return np.array(ious) -------------------------------------------------------------------------------- /main.py: -------------------------------------------------------------------------------- 1 | import argparse 2 | import tqdm 3 | import numpy as np 4 | 5 | import torch 6 | from torch.optim import Adam 7 | from torch.optim.lr_scheduler import StepLR 8 | from torch import nn 9 | 10 | from torch.utils.data import DataLoader 11 | 12 | from evaluation import compute_iou 13 | from models.classification_modern_model import cnn_model 14 | from models.segmentation_unet import UNet 15 | from data_loader import ImagevDatasetForClassi, ImagevDatasetForSegment, ImagevDatasetForRetroEsti 16 | from utils import generate_img_from_mask, write_image, calculate_luminance_ratio 17 | 18 | USE_CUDA = torch.cuda.is_available() 19 | DEVICE = torch.device('cuda' if USE_CUDA else 'cpu') 20 | 21 | CLASSI_MODEL_PATH = '../pth/classi/classi_' 22 | SEGMENT_MODEL_PATH = '../pth/segment/segment_' 23 | 24 | SEGMENT_TRUE_PATH = '../output/true/true_' 25 | SEGMENT_PRED_PATH = '../output/pred/pred_' 26 | 27 | BOTH_PRED_PATH = '../output/both/output_' 28 | 29 | def train(model, train_loader, optimizer): 30 | model.train() 31 | correct = 0 32 | 33 | criterion1 = nn.CrossEntropyLoss().to(DEVICE) 34 | criterion2 = nn.MSELoss().to(DEVICE) 35 | 36 | train_loss = 0 37 | for batch_idx, samples in enumerate(train_loader): 38 | data, target1, target2 = samples 39 | #print(data.shape, target.shape) 40 | 41 | data = data.to(DEVICE) 42 | target1 = target1.to(DEVICE) 43 | target2 = target2.to(DEVICE) 44 | 45 | output = model(data) 46 | #loss = criterion(m(output), target) 47 | loss1 = criterion1(output[:,0:2], target1) 48 | loss2 = criterion2(output[:,2], target2) 49 | loss = loss1 + loss2 50 | train_loss += loss.item() 51 | 52 | classi_output = output[:, 0:2] 53 | prediction = classi_output.max(1, keepdim=True)[1] 54 | correct += prediction.eq(target1.view_as(prediction)).sum().item() 55 | 56 | optimizer.zero_grad() 57 | loss.backward() 58 | optimizer.step() 59 | 60 | print('Batch Index: {}\tLoss: {:.6f}\tLoss1: {:.6f}\tLoss2: {:.6f}'.format(batch_idx, loss.item(), loss1.item(), loss2.item())) 61 | train_loss /= len(train_loader.dataset) 62 | train_acc = 100. * correct / len(train_loader.dataset) 63 | 64 | return train_loss, train_acc 65 | 66 | def train2(model, train_loader, optimizer): 67 | model.train() 68 | train_loss = 0 69 | criterion = nn.CrossEntropyLoss().to(DEVICE) 70 | 71 | for batch_idx, samples in enumerate(train_loader): 72 | data, target = samples 73 | #print(data.shape, target.shape) 74 | 75 | data = data.to(DEVICE) 76 | target = target.to(DEVICE) 77 | 78 | output = model(data) 79 | loss = criterion(output, target) 80 | train_loss += loss.item() 81 | 82 | optimizer.zero_grad() 83 | loss.backward() 84 | optimizer.step() 85 | 86 | print('Batch Index: {}\tLoss: {:.6f}'.format(batch_idx, loss.item())) 87 | train_loss /= len(train_loader.dataset) 88 | 89 | return train_loss 90 | 91 | def evaluate(model, test_loader): 92 | model.eval() 93 | test_loss = 0 94 | correct = 0 95 | 96 | criterion1 = nn.CrossEntropyLoss().to(DEVICE) 97 | criterion2 = nn.L1Loss().to(DEVICE) 98 | 99 | with torch.no_grad(): 100 | for samples in test_loader: 101 | data, target1, target2 = samples 102 | 103 | data = data.to(DEVICE) 104 | target1 = target1.to(DEVICE) 105 | target2 = target2.to(DEVICE) 106 | 107 | output = model(data) 108 | loss1 = criterion1(output[:, 0:2], target1) 109 | loss2 = criterion2(output[:, 2], target2) 110 | #loss = loss1 + loss2 111 | loss = loss2 112 | test_loss += loss.item() 113 | 114 | classi_output = output[:, 0:2] 115 | prediction = classi_output.max(1, keepdim=True)[1] 116 | correct += prediction.eq(target1.view_as(prediction)).sum().item() 117 | 118 | test_loss /= len(test_loader.dataset) 119 | test_acc = 100. * correct / len(test_loader.dataset) 120 | return test_loss, test_acc 121 | 122 | def evaluate2(model, test_loader): 123 | model.eval() 124 | test_loss = 0 125 | criterion = nn.CrossEntropyLoss().to(DEVICE) 126 | 127 | total_ious = [] 128 | with torch.no_grad(): 129 | for samples in test_loader: 130 | data, target = samples 131 | 132 | data = data.to(DEVICE) 133 | target = target.to(DEVICE) 134 | 135 | output = model(data) 136 | loss = criterion(output, target) 137 | test_loss += loss.item() 138 | 139 | ious = compute_iou(output, target) 140 | total_ious.append(ious) 141 | 142 | test_loss /= len(test_loader.dataset) 143 | total_ious = np.mean(np.array(total_ious), axis=0) 144 | 145 | return test_loss, total_ious 146 | 147 | def predict2(model, test_loader): 148 | model.eval() 149 | 150 | cnt = 0 151 | with torch.no_grad(): 152 | for samples in tqdm.tqdm(test_loader): 153 | data, target = samples 154 | 155 | data = data.to(DEVICE) 156 | target = target.to(DEVICE) 157 | output = model(data) 158 | 159 | y_true = target 160 | y_pred = output.max(1, keepdim=False)[1] 161 | 162 | y_true = y_true.cpu().detach().numpy() 163 | y_pred = y_pred.cpu().detach().numpy() 164 | 165 | for i in range(np.size(y_true, 0)): 166 | img_true = generate_img_from_mask(y_true[i], np.size(y_true, 1), np.size(y_true, 2)) 167 | img_pred = generate_img_from_mask(y_pred[i], np.size(y_pred, 1), np.size(y_pred, 2)) 168 | 169 | filepath_true = SEGMENT_TRUE_PATH + str(cnt).zfill(4) + '.png' 170 | filepath_pred = SEGMENT_PRED_PATH + str(cnt).zfill(4) + '.png' 171 | 172 | write_image(filepath_true, img_true) 173 | write_image(filepath_pred, img_pred) 174 | 175 | cnt += 1 176 | 177 | print('prediction done!') 178 | 179 | 180 | def predict3(classi_model, segment_model, test_loader): 181 | classi_model.eval() 182 | segment_model.eval() 183 | 184 | cnt = 0 185 | res = [] 186 | with torch.no_grad(): 187 | for samples in tqdm.tqdm(test_loader): 188 | csv, img = samples 189 | 190 | csv = csv.to(DEVICE) 191 | img = img.to(DEVICE) 192 | 193 | output = classi_model(csv) 194 | 195 | classi_output = output[:, 0:2] 196 | prediction = classi_output.max(1, keepdim=False)[1] 197 | prediction = prediction.cpu().detach().numpy() 198 | h = int(480 * np.squeeze(output[:, 2].cpu().detach().numpy())) 199 | crop_csv = csv[:, :, h:h + 192, :] 200 | 201 | if prediction == 1: 202 | segment_output = segment_model(crop_csv) 203 | segment_output = segment_output.max(1, keepdim=False)[1] 204 | segment_output = segment_output.cpu().detach().numpy() 205 | 206 | for i in range(np.size(segment_output, 0)): 207 | img_pred = generate_img_from_mask(segment_output[i], np.size(segment_output, 1), np.size(segment_output, 2)) 208 | filepath_pred = BOTH_PRED_PATH + str(cnt).zfill(4) + '.png' 209 | write_image(filepath_pred, img_pred) 210 | 211 | ratio = calculate_luminance_ratio(crop_csv.cpu().detach().numpy(), segment_output[i], np.size(segment_output, 1), np.size(segment_output, 2)) 212 | res.append(ratio) 213 | 214 | cnt += 1 215 | else: 216 | res.append(0) 217 | 218 | cnt += 1 219 | 220 | res = np.array(res) 221 | np.savetxt('../output/output.txt', res, fmt='%1.6f') # use exponential notation 222 | print('prediction done!') 223 | 224 | 225 | 226 | def save_model(modelpath, model, optimizer, scheduler): 227 | state = { 228 | 'model': model.state_dict(), 229 | 'optimizer': optimizer.state_dict(), 230 | 'scheduler': scheduler.state_dict() 231 | } 232 | torch.save(state, modelpath) 233 | 234 | print('model saved') 235 | 236 | 237 | def load_model(modelpath, model, optimizer=None, scheduler=None): 238 | state = torch.load(modelpath, map_location=torch.device(DEVICE)) 239 | model.load_state_dict(state['model']) 240 | if optimizer is not None: 241 | optimizer.load_state_dict(state['optimizer']) 242 | if scheduler is not None: 243 | scheduler.load_state_dict(state['scheduler']) 244 | 245 | print('model loaded') 246 | 247 | 248 | def main(args): 249 | # train 250 | if args.mode == 'train': 251 | train_dataset = ImagevDatasetForClassi(mode='train') 252 | train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=args.shuffle, num_workers=0) 253 | 254 | val_dataset = ImagevDatasetForClassi(mode='val') 255 | val_dataloader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=0) 256 | print(train_dataloader, val_dataloader) 257 | 258 | modelpath = CLASSI_MODEL_PATH 259 | if args.model == 'resnet18': 260 | modelpath += 'resnet18.pth' 261 | elif args.model == 'vgg16': 262 | modelpath += 'vgg16.pth' 263 | elif args.model == 'alexnet': 264 | modelpath += 'alexnet.pth' 265 | elif args.model == 'mobilenet': 266 | modelpath += 'mobilenet.pth' 267 | elif args.model == 'shufflenet': 268 | modelpath += 'shufflenet.pth' 269 | 270 | model = cnn_model(args.model) 271 | if torch.cuda.device_count() > 1: 272 | print('multi gpu used!') 273 | model = nn.DataParallel(model) 274 | model = model.to(DEVICE) 275 | 276 | # set optimizer 277 | optimizer = Adam( 278 | [param for param in model.parameters() if param.requires_grad], 279 | lr=args.lr) 280 | scheduler = StepLR(optimizer, step_size=10, gamma=0.8) 281 | #modelpath = CLASSI_MODEL_PATH 282 | #load_model(modelpath, model, optimizer, scheduler) 283 | 284 | acc_prev = 0.0 285 | for epoch in range(args.epoch): 286 | # train set 287 | loss, acc = train(model, train_dataloader, optimizer) 288 | # validate set 289 | val_loss, val_acc = evaluate(model, val_dataloader) 290 | 291 | print('Epoch:{}\tTrain Loss:{:.6f}\tTrain Acc:{:2.4f}'.format(epoch, loss, acc)) 292 | print('Val Loss:{:.6f}\tVal Acc:{:2.4f}'.format(val_loss, val_acc)) 293 | 294 | if val_acc > acc_prev: 295 | acc_prev = val_acc 296 | #modelpath = '../pth/20200726/model-{:d}-{:.6f}-{:2.4f}.pth'.format(epoch, val_loss, val_acc) 297 | save_model(modelpath, model, optimizer, scheduler) 298 | 299 | # scheduler update 300 | scheduler.step() 301 | # evaluate 302 | elif args.mode == 'evaluate': 303 | val_dataset = ImagevDatasetForClassi(mode='val') 304 | val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=0) 305 | print(val_dataloader) 306 | 307 | model = cnn_model(args.model) 308 | if torch.cuda.device_count() > 1: 309 | print('multi gpu used!') 310 | model = nn.DataParallel(model) 311 | model = model.to(DEVICE) 312 | 313 | modelpath = CLASSI_MODEL_PATH 314 | if args.model == 'resnet18': 315 | modelpath += 'resnet18.pth' 316 | elif args.model == 'vgg16': 317 | modelpath += 'vgg16.pth' 318 | elif args.model == 'alexnet': 319 | modelpath += 'alexnet.pth' 320 | elif args.model == 'mobilenet': 321 | modelpath += 'mobilenet.pth' 322 | elif args.model == 'shufflenet': 323 | modelpath += 'shufflenet.pth' 324 | load_model(modelpath, model) 325 | 326 | test_loss, test_acc = evaluate(model, val_dataloader) 327 | print('Test Loss:{:.6f}\tTest Acc:{:2.4f}'.format(test_loss, test_acc)) 328 | # predict 329 | elif args.mode == 'predict': 330 | print('there is no prediction module in classification.') 331 | 332 | def main2(args): 333 | # train 334 | if args.mode == 'train': 335 | train_dataset = ImagevDatasetForSegment(mode='train') 336 | train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=args.shuffle, num_workers=0) 337 | 338 | val_dataset = ImagevDatasetForSegment(mode='val') 339 | val_dataloader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=0) 340 | print(train_dataloader, val_dataloader) 341 | 342 | model = UNet() 343 | if torch.cuda.device_count() > 1: 344 | print('multi gpu used!') 345 | model = nn.DataParallel(model) 346 | model = model.to(DEVICE) 347 | 348 | modelpath = SEGMENT_MODEL_PATH 349 | modelpath += 'unet.pth' 350 | # set optimizer 351 | optimizer = Adam( 352 | [param for param in model.parameters() if param.requires_grad], 353 | lr=args.lr) 354 | scheduler = StepLR(optimizer, step_size=10, gamma=0.8) 355 | # modelpath = CLASSI_MODEL_PATH 356 | # load_model(modelpath, model, optimizer, scheduler) 357 | 358 | iou_prev = 0.0 359 | for epoch in range(args.epoch): 360 | # train set 361 | loss = train2(model, train_dataloader, optimizer) 362 | train_loss, train_ious = evaluate2(model, train_dataloader) 363 | train_mean_ious = np.mean(train_ious) 364 | # validate set 365 | val_loss, val_ious = evaluate2(model, val_dataloader) 366 | val_mean_iou = np.mean(val_ious) 367 | 368 | print('Epoch:{}\tTrain Loss:{:.6f}'.format(epoch, loss)) 369 | print('Train Loss:{:.6f}\tTrain Mean IoU:{:2.4f}\tTrain IoU1:{:2.4f}\tTrain IoU2:{:2.4f}\tTrain IoU3:{:2.4f}'.format(train_loss, train_mean_ious, train_ious[0], train_ious[1], train_ious[2])) 370 | print('Val Loss:{:.6f}\tVal Mean IoU:{:2.4f}\tVal IoU1:{:2.4f}\tVal IoU2:{:2.4f}\tVal IoU3:{:2.4f}'.format(val_loss, val_mean_iou, val_ious[0], val_ious[1], val_ious[2])) 371 | 372 | if val_mean_iou > iou_prev: 373 | iou_prev = val_mean_iou 374 | # modelpath = '../pth/20200726/model-{:d}-{:.6f}-{:2.4f}.pth'.format(epoch, val_loss, val_acc) 375 | save_model(modelpath, model, optimizer, scheduler) 376 | 377 | # scheduler update 378 | scheduler.step() 379 | # evaluate 380 | elif args.mode == 'evaluate': 381 | val_dataset = ImagevDatasetForSegment(mode='val') 382 | val_dataloader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=0) 383 | print(val_dataloader) 384 | 385 | model = UNet() 386 | if torch.cuda.device_count() > 1: 387 | print('multi gpu used!') 388 | model = nn.DataParallel(model) 389 | model = model.to(DEVICE) 390 | 391 | modelpath = SEGMENT_MODEL_PATH 392 | modelpath += 'unet.pth' 393 | load_model(modelpath, model) 394 | 395 | val_loss, val_ious = evaluate2(model, val_dataloader) 396 | val_mean_iou = np.mean(val_ious) 397 | print('Val Loss:{:.6f}\tVal Mean IoU:{:2.4f}\tVal IoU1:{:2.4f}\tVal IoU2:{:2.4f}\tVal IoU3:{:2.4f}'.format(val_loss, val_mean_iou, val_ious[0], val_ious[1], val_ious[2])) 398 | # predict 399 | elif args.mode == 'predict': 400 | val_dataset = ImagevDatasetForSegment(mode='val') 401 | val_dataloader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, num_workers=0) 402 | print(val_dataloader) 403 | 404 | model = UNet() 405 | if torch.cuda.device_count() > 1: 406 | print('multi gpu used!') 407 | model = nn.DataParallel(model) 408 | model = model.to(DEVICE) 409 | 410 | modelpath = SEGMENT_MODEL_PATH 411 | modelpath += 'unet.pth' 412 | load_model(modelpath, model) 413 | 414 | predict2(model, val_dataloader) 415 | 416 | # only for prediction using both classification and segmentation 417 | def main3(args): 418 | # predict 419 | if args.mode == 'predict': 420 | val_dataset = ImagevDatasetForRetroEsti(mode='val2') 421 | val_dataloader = DataLoader(val_dataset, batch_size=1, shuffle=False, num_workers=0) 422 | print(val_dataloader) 423 | 424 | classi_model = cnn_model(args.model) 425 | if torch.cuda.device_count() > 1: 426 | print('multi gpu used!') 427 | classi_model = nn.DataParallel(classi_model) 428 | classi_model = classi_model.to(DEVICE) 429 | 430 | modelpath = CLASSI_MODEL_PATH 431 | if args.model == 'resnet18': 432 | modelpath += 'resnet18.pth' 433 | elif args.model == 'vgg16': 434 | modelpath += 'vgg16.pth' 435 | elif args.model == 'alexnet': 436 | modelpath += 'alexnet.pth' 437 | elif args.model == 'mobilenet': 438 | modelpath += 'mobilenet.pth' 439 | elif args.model == 'shufflenet': 440 | modelpath += 'shufflenet.pth' 441 | load_model(modelpath, classi_model) 442 | 443 | segment_model = UNet() 444 | if torch.cuda.device_count() > 1: 445 | print('multi gpu used!') 446 | segment_model = nn.DataParallel(segment_model) 447 | segment_model = segment_model.to(DEVICE) 448 | 449 | modelpath = SEGMENT_MODEL_PATH 450 | modelpath += 'unet.pth' 451 | load_model(modelpath, segment_model) 452 | 453 | predict3(classi_model, segment_model, val_dataloader) 454 | 455 | 456 | 457 | 458 | if __name__ == '__main__': 459 | parser = argparse.ArgumentParser() 460 | parser.add_argument( 461 | '--batch_size', 462 | help='the number of samples in mini-batch', 463 | default=128, 464 | type=int) 465 | parser.add_argument( 466 | '--epoch', 467 | help='the number of training iterations', 468 | default=50, 469 | type=int) 470 | parser.add_argument( 471 | '--lr', 472 | help='learning rate', 473 | default=0.001, 474 | type=float) 475 | parser.add_argument( 476 | '--shuffle', 477 | help='True, or False', 478 | default=True, 479 | type=bool) 480 | parser.add_argument( 481 | '--mode', 482 | help='train, evaluate, or predict', 483 | default='evaluate', 484 | type=str) 485 | parser.add_argument( 486 | '--model', 487 | help='resnet18, vgg16, alexnet, mobilenet, or shufflenet', 488 | default='resnet18', 489 | type=str) 490 | parser.add_argument( 491 | '--type', 492 | help='classification, segmentation, or both', 493 | default='classification', 494 | type=str) 495 | 496 | args = parser.parse_args() 497 | print(args) 498 | 499 | if args.type == 'classification': 500 | main(args) 501 | elif args.type == 'segmentation': 502 | main2(args) 503 | elif args.type == 'both': 504 | main3(args) -------------------------------------------------------------------------------- /models/classification_modern_model.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torchvision 3 | from torchsummary import summary 4 | 5 | USE_CUDA = torch.cuda.is_available() 6 | DEVICE = torch.device('cuda' if USE_CUDA else 'cpu') 7 | 8 | def cnn_model(name='resnet18', pretrained=True, num_output=3): 9 | if name == 'resnet18': 10 | net = torchvision.models.resnet18(pretrained=pretrained).to(DEVICE) 11 | net.conv1 = torch.nn.Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False) 12 | net.fc = torch.nn.Sequential( 13 | torch.nn.Linear(512, 256), 14 | torch.nn.ReLU(), 15 | torch.nn.Linear(256, 128), 16 | torch.nn.ReLU(), 17 | torch.nn.Linear(128, num_output)) 18 | elif name == 'vgg16': 19 | net = torchvision.models.vgg16(pretrained=pretrained).to(DEVICE) 20 | net.features[0] = torch.nn.Conv2d(1, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False) 21 | net.classifier = torch.nn.Sequential( 22 | torch.nn.Linear(512 * 7 * 7, 256), 23 | torch.nn.ReLU(), 24 | torch.nn.Linear(256, 128), 25 | torch.nn.ReLU(), 26 | torch.nn.Linear(128, num_output)) 27 | elif name == 'alexnet': 28 | net = torchvision.models.alexnet(pretrained=pretrained).to(DEVICE) 29 | net.features[0] = torch.nn.Conv2d(1, 64, kernel_size=(11, 11), stride=(4, 4), padding=(2, 2), bias=False) 30 | net.classifier = torch.nn.Sequential( 31 | torch.nn.Linear(256 * 6 * 6, 256), 32 | torch.nn.ReLU(), 33 | torch.nn.Linear(256, 128), 34 | torch.nn.ReLU(), 35 | torch.nn.Linear(128, num_output)) 36 | elif name == 'mobilenet': 37 | net = torchvision.models.mobilenet_v2(pretrained=pretrained).to(DEVICE) 38 | net.features[0][0] = torch.nn.Conv2d(1, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False) 39 | net.classifier = torch.nn.Sequential( 40 | torch.nn.Linear(1280, 256), 41 | torch.nn.ReLU(), 42 | torch.nn.Linear(256, 128), 43 | torch.nn.ReLU(), 44 | torch.nn.Linear(128, num_output)) 45 | elif name == 'shufflenet': 46 | net = torchvision.models.shufflenet_v2_x0_5(pretrained=pretrained).to(DEVICE) 47 | net.conv1[0] = torch.nn.Conv2d(1, 24, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False) 48 | net.fc = torch.nn.Sequential( 49 | torch.nn.Linear(1024, 256), 50 | torch.nn.ReLU(), 51 | torch.nn.Linear(256, 128), 52 | torch.nn.ReLU(), 53 | torch.nn.Linear(128, num_output)) 54 | 55 | else: 56 | print('wrong cnn model!') 57 | 58 | return net 59 | 60 | def test(): 61 | net = cnn_model(name='shufflenet') 62 | 63 | summary(net, (1, 480, 640)) 64 | #y = net(torch.randn(1, 3, 480, 640)) 65 | print(net) 66 | 67 | 68 | if __name__ == "__main__": 69 | test() 70 | -------------------------------------------------------------------------------- /models/segmentation_unet.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import torch.nn as nn 3 | from torchsummary import summary 4 | 5 | def double_conv(in_channels, out_channels): 6 | return nn.Sequential( 7 | nn.Conv2d(in_channels, out_channels, 5, padding=2), 8 | nn.ReLU(inplace=True), 9 | nn.Conv2d(out_channels, out_channels, 5, padding=2), 10 | nn.ReLU(inplace=True) 11 | ) 12 | 13 | class UNet(nn.Module): 14 | def __init__(self, cc=16, num_class=3): 15 | super().__init__() 16 | 17 | self.dconv_down1 = double_conv(1, cc) 18 | self.dconv_down2 = double_conv(cc, 2*cc) 19 | self.dconv_down3 = double_conv(2*cc, 4*cc) 20 | self.dconv_down4 = double_conv(4*cc, 8*cc) 21 | self.dconv_down5 = double_conv(8*cc, 16*cc) 22 | 23 | self.maxpool = nn.MaxPool2d(2) 24 | self.upsample = nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True) 25 | 26 | self.dconv_up4 = double_conv(16*cc + 8*cc, 8*cc) 27 | self.dconv_up3 = double_conv(8*cc + 4*cc, 4*cc) 28 | self.dconv_up2 = double_conv(4*cc + 2*cc, 2*cc) 29 | self.dconv_up1 = double_conv(2*cc + cc, cc) 30 | 31 | self.conv_last = nn.Conv2d(in_channels=cc, out_channels=num_class, kernel_size=1) 32 | 33 | def forward(self, x): 34 | conv1 = self.dconv_down1(x) 35 | x = self.maxpool(conv1) 36 | 37 | conv2 = self.dconv_down2(x) 38 | x = self.maxpool(conv2) 39 | 40 | conv3 = self.dconv_down3(x) 41 | x = self.maxpool(conv3) 42 | 43 | conv4 = self.dconv_down4(x) 44 | x = self.maxpool(conv4) 45 | 46 | conv5 = self.dconv_down5(x) 47 | #x = self.maxpool(conv5) 48 | 49 | x = self.upsample(conv5) 50 | x = torch.cat([x, conv4], dim=1) 51 | 52 | x = self.dconv_up4(x) 53 | x = self.upsample(x) 54 | x = torch.cat([x, conv3], dim=1) 55 | 56 | x = self.dconv_up3(x) 57 | x = self.upsample(x) 58 | x = torch.cat([x, conv2], dim=1) 59 | 60 | x = self.dconv_up2(x) 61 | x = self.upsample(x) 62 | x = torch.cat([x, conv1], dim=1) 63 | 64 | x = self.dconv_up1(x) 65 | 66 | out = self.conv_last(x) 67 | 68 | return out 69 | 70 | if __name__ == '__main__': 71 | model = UNet() 72 | summary(model, (1, 192, 640)) 73 | 74 | print(model) -------------------------------------------------------------------------------- /utils.py: -------------------------------------------------------------------------------- 1 | import numpy as np 2 | import cv2 3 | 4 | def generate_img_from_mask(mask, height, width): 5 | b = np.float32(np.where(mask == 2, 255, 0)) 6 | g = np.float32(np.where(mask == 1, 255, 0)) 7 | r = np.float32(np.zeros([height, width, 1])) 8 | res = cv2.merge((b, g, r)) 9 | 10 | return res 11 | 12 | def write_image(filepath, img): 13 | cv2.imwrite(filepath, img) 14 | 15 | def calculate_luminance_ratio(img, mask, height, width): 16 | reference_plate = np.float32(np.where(mask == 1, 1, 0)) 17 | road_marking = np.float32(np.where(mask == 2, 1, 0)) 18 | 19 | img = img.reshape(1, height, width) 20 | reference_plate = reference_plate.reshape(1, height, width) 21 | road_marking = road_marking.reshape(1, height, width) 22 | 23 | rp_mean = np.mean(np.multiply(img, reference_plate).reshape(-1)) 24 | rm_mean = np.mean(np.multiply(img, road_marking).reshape(-1)) 25 | 26 | ratio = rm_mean / (rp_mean + 0.01) 27 | if ratio > 2.0: 28 | ratio = 2.0; 29 | 30 | return ratio 31 | 32 | --------------------------------------------------------------------------------