def main():
config = Config()
if os.path.exists(config.SAVE_PATH):
shutil.rmtree(config.SAVE_PATH)
os.makedirs(config.SAVE_PATH, exist_ok=True)
trainF = open(os.path.join(config.SAVE_PATH, "train.csv"), 'w')
testF = open(os.path.join(config.SAVE_PATH, "test.csv"), 'w')
train_img_path = os.path.abspath('../ICDAR_2015/train_img')
train_gt_path = os.path.abspath('../ICDAR_2015/train_gt')
val_img_path = os.path.abspath('../ICDAR_2015/test_img')
val_gt_path = os.path.abspath('../ICDAR_2015/test_gt')
kwargs = {'num_workers': 2, 'pin_memory': True} if torch.cuda.is_available() else {}
train_dataset = custom_dataset(train_img_path, train_gt_path)
train_loader = data.DataLoader(train_dataset, batch_size=config.TRAIN_BATCH*len(device_list), \
shuffle=True, drop_last=True, **kwargs)
val_dataset = custom_dataset(val_img_path, val_gt_path)
val_loader = data.DataLoader(val_dataset, batch_size=config.TRAIN_BATCH*len(device_list), \
shuffle=True, drop_last=True, **kwargs)
net = EAST()
if torch.cuda.is_available():
net = net.cuda(device=device_list[0])
net = torch.nn.DataParallel(net, device_ids=device_list)
optimizer = torch.optim.Adam(net.parameters(), lr=config.BASE_LR, weight_decay=config.WEIGHT_DECAY)
for epoch in range(config.EPOCHS):
train(net, epoch, train_loader, optimizer, trainF, config)
test(net, epoch, val_loader, testF, config)
if epoch != 0 and epoch % config.SAVE_INTERVAL == 0:
torch.save({'state_dict': net.state_dict()}, os.path.join(os.getcwd(), config.SAVE_PATH, "laneNet{}.pth.tar".format(epoch)))
trainF.close()
testF.close()
torch.save({'state_dict': net.state_dict()}, os.path.join(os.getcwd(), config.SAVE_PATH, "finalNet.pth.tar"))
trained_model_file = config['trained_model_file']
test_mini_batch_size = config['test_mini_batch_size']
score_threshold = config['score_threshold']
iou_threshold = config['iou_threshold']
max_boxes = config['max_boxes']
representation = geometry + "_" + label_method
test_images_dir = os.path.join(test_data_dir, "images")
test_images_pred_dir = os.path.join(test_data_dir, "images_pred")
test_annotations_pred_dir = os.path.join(test_data_dir, "annotations_pred")
model = EAST(geometry=geometry)
if cuda:
model.cuda()
model.load_state_dict(torch.load(trained_model_file))
model.eval()
if not os.path.exists(test_images_pred_dir):
os.mkdir(test_images_pred_dir)
if not os.path.exists(test_annotations_pred_dir):
os.mkdir(test_annotations_pred_dir)
testset = ImageTestDataSet(test_images_dir)
test_loader = torch.utils.data.DataLoader(testset,
batch_size=test_mini_batch_size,
shuffle=True)
with torch.no_grad():
boxes_pred = []