def main():
# Training settings
parser = argparse.ArgumentParser(description='Segmentation')
parser.add_argument('--test_batch_size', type=int, default=16, metavar='N',
help='input batch size for test (default: 64)')
parser.add_argument('--nlevels', type=int, default=4, help="number of polygon levels, higher->finer")
parser.add_argument('--feat', type=int, default=256, help="number of base feature layers")
parser.add_argument('--no_cuda', action='store_true', default=False,
help='disables CUDA')
parser.add_argument('--seed', type=int, default=1, metavar='S',
help='random seed (default: 1)')
parser.add_argument('--data_folder', type=str, default="mres_processed_data",
help='path to data folder (default: processed_data)')
parser.add_argument('--ckpt', type=str, default='checkpoint/checkpoint_polygonnet_best.pth.tar', help="path to checkpoint to load")
parser.add_argument('--transpose', action='store_true', help="transpose target")
parser.add_argument('--workers', default=12, type=int, help="number of data loading workers")
args = parser.parse_args()
use_cuda = not args.no_cuda and torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
# PYTORCH VERSION > 1.0.0
assert(float(torch.__version__.split('.')[-3]) > 0)
torch.manual_seed(args.seed)
# get training / valid sets
args.img_mean = [0.485, 0.456, 0.406]
args.img_std = [0.229, 0.224, 0.225]
normalize = torchvision.transforms.Normalize(mean=args.img_mean,
std=args.img_std)
# DO NOT include horizontal and vertical flips in the composed transforms!
transform = torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
normalize,
])
testset = CityScapeLoader(args.data_folder, "test", transforms=transform, RandomHorizontalFlip=0.0, RandomVerticalFlip=0.0, mres=False, transpose=args.transpose)
test_loader = DataLoader(testset, batch_size=args.test_batch_size, shuffle=True, drop_last=False, num_workers=args.workers, pin_memory=True)
# initialize model
model = PolygonNet(nlevels=args.nlevels, dropout=False, feat=args.feat)
model = nn.DataParallel(model)
if os.path.isfile(args.ckpt):
print("=> loading checkpoint '{}'".format(args.ckpt))
checkpoint = torch.load(args.ckpt)
args.best_miou = checkpoint['best_miou']
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.ckpt, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.ckpt))
model.to(device)
print("{} paramerters in total".format(sum(x.numel() for x in model.parameters())))
evaluate(args, model, test_loader, device)
def main():
# Training settings
parser = argparse.ArgumentParser(description='Segmentation')
parser.add_argument('--batch_size', type=int, default=32, metavar='N',
help='input batch size for training (default: 32)')
parser.add_argument('--val_batch_size', type=int, default=32, metavar='N',
help='input batch size for validation (default: 32)')
parser.add_argument('--loss_type', type=str, choices=['l1', 'l2'], default='l1', help='type of loss for raster loss computation')
parser.add_argument('--decay', action="store_true", help="switch to decay learning rate")
parser.add_argument('--nlevels', type=int, default=5, help="number of polygon levels, higher->finer")
parser.add_argument('--feat', type=int, default=256, help="number of base feature layers")
parser.add_argument('--dropout', action='store_true', help="dropout during training")
parser.add_argument('--epochs', type=int, default=100, metavar='N',
help='number of epochs to train (default: 10)')
parser.add_argument('--lr', type=float, default=1e-2, metavar='LR',
help='learning rate (default: 0.01)')
parser.add_argument('--no_cuda', action='store_true', default=False,
help='disables CUDA training')
parser.add_argument('--seed', type=int, default=1, metavar='S',
help='random seed (default: 1)')
parser.add_argument('--data_folder', type=str, default="mres_processed_data",
help='path to data folder (default: mres_processed_data)')
parser.add_argument('--log_interval', type=int, default=20, metavar='N',
help='how many batches to wait before logging training status')
parser.add_argument('--log_dir', type=str, default="log",
help='log directory for run')
parser.add_argument('--resume', type=str, default=None, help="path to checkpoint if resume is needed")
parser.add_argument('--timestamp', action='store_true', help="add timestamp to log_dir name")
parser.add_argument('--multires', action='store_true', help="multiresolution loss")
parser.add_argument('--transpose', action='store_true', help="transpose target")
parser.add_argument('--smooth_loss', default=1.0, type=float, help="smoothness loss multiplier (0 for none)")
parser.add_argument('--uniform_loss', action='store_true', help="use same loss regardless of category frequency")
parser.add_argument('--workers', default=12, type=int, help="number of data loading workers")
parser.add_argument('--check', action='store_true', help="gradient checks")
args = parser.parse_args()
use_cuda = not args.no_cuda and torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
args.TRAIN_GLOB_STEP = 0
args.VAL_GLOB_STEP = 0
args.label_names = ["car", "truck", "train", "bus", "motorcycle", "bicycle", "rider", "person"]
if not args.uniform_loss:
args.instances = torch.tensor([30246, 516, 76, 325, 658, 3307, 1872, 19563]).double().to(device)
else:
args.instances = torch.tensor([1, 1, 1, 1, 1, 1, 1, 1]).double().to(device)
args.nclass = len(args.label_names)
args.weights = 1 / torch.log(args.instances / args.instances.sum() + 1.02)
args.weights /= args.weights.sum()
# PYTORCH VERSION > 1.0.0
assert(float(torch.__version__.split('.')[-3]) > 0)
# boiler-plate
if args.timestamp:
args.log_dir += '_' + time.strftime("%Y_%m_%d_%H_%M_%S")
logger = initialize_logger(args)
torch.manual_seed(args.seed)
# TensorboardX writer
args.tblogdir = os.path.join(args.log_dir, "tensorboard_log")
if not os.path.exists(args.tblogdir):
os.makedirs(args.tblogdir)
writer = SummaryWriter(log_dir=args.tblogdir)
# get training / valid sets
args.img_mean = [0.485, 0.456, 0.406]
args.img_std = [0.229, 0.224, 0.225]
normalize = torchvision.transforms.Normalize(mean=args.img_mean,
std=args.img_std)
# DO NOT include horizontal and vertical flips in the composed transforms!
transform = torchvision.transforms.Compose([
torchvision.transforms.ColorJitter(hue=.1, saturation=.1),
torchvision.transforms.ToTensor(),
normalize,
])
trainset = CityScapeLoader(args.data_folder, "train", transforms=transform, RandomHorizontalFlip=0.5, RandomVerticalFlip=0.0, mres=args.multires, transpose=args.transpose)
valset = CityScapeLoader(args.data_folder, "val", transforms=transform, RandomHorizontalFlip=0.0, RandomVerticalFlip=0.0, mres=args.multires, transpose=args.transpose)
train_loader = DataLoader(trainset, batch_size=args.batch_size, shuffle=True, drop_last=True, num_workers=args.workers, pin_memory=True)
val_loader = DataLoader(valset, batch_size=args.val_batch_size, shuffle=True, drop_last=False, num_workers=args.workers, pin_memory=True)
# initialize and parallelize model
model = PolygonNet(nlevels=args.nlevels, dropout=args.dropout, feat=args.feat)
model = nn.DataParallel(model)
model.to(device)
# Multi-Resolution
n = model.module.npoints
if args.multires:
args.res = [224, 112, 56, 28]
args.npt = [n, int(n/2), int(n/4), int(n/8)]
args.levels = [1, 2, 4, 8]
else:
args.res = [224]
args.npt = [n]
args.levels = [1]
# initialize optimizer
optimizer = optim.Adam(model.parameters(), lr=args.lr)
args.start_epoch = -1
args.best_miou = 0
if args.resume:
if os.path.isfile(args.resume):
print("=> loading checkpoint '{}'".format(args.resume))
checkpoint = torch.load(args.resume)
args.start_epoch = checkpoint['epoch']
args.best_miou = checkpoint['best_miou']
model.load_state_dict(checkpoint['state_dict'])
optimizer.load_state_dict(checkpoint['optimizer'])
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.resume, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.resume))
logger.info("{} paramerters in total".format(sum(x.numel() for x in model.parameters())))
if args.decay:
scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=10, gamma=0.75)
checkpoint_path = os.path.join(args.log_dir, 'checkpoint')
rres = [True, False, False, False] if args.multires else [True]
criterion = [BatchedRasterLoss2D(npoints=n, res=r, loss=args.loss_type, return_raster=rr).to(device) for n, r, rr in zip(args.npt, args.res, rres)]
criterion_smooth = SmoothnessLoss()
# training loop
for epoch in range(args.start_epoch + 1, args.epochs):
if args.decay:
scheduler.step(epoch)
train(args, model, train_loader, criterion, criterion_smooth, optimizer, epoch, device, logger, writer)
miou = validate(args, model, val_loader, criterion, criterion_smooth, epoch, device, logger, writer)
if miou > args.best_miou:
args.best_miou = miou
is_best = True
else:
is_best = False
save_checkpoint({
'epoch': epoch,
'state_dict': model.state_dict(),
'best_miou': args.best_miou,
'optimizer': optimizer.state_dict(),
}, is_best, epoch, checkpoint_path, "_polygonnet", logger)
def deep_rnn_annotate():
parser = argparse.ArgumentParser(description='manual to test script')
# parser.add_argument('--gpu_id', nargs='+', type=int)
parser.add_argument('--batch_size', type=int, default=1)
parser.add_argument('--model', type=str, default='./Polygon_deep_RNN.pth')
args = parser.parse_args()
devices = [0]
batch_size = args.batch_size
root_path = "../images"
img_file = os.path.join(root_path, 'save.jpg')
# mapping_location = {'cuda:0': 'cuda:' + str(devices[0])}
img = Image.open(img_file).convert('RGB')
img_len = img.size[0]
img_width = img.size[1]
transform = transforms.Compose([
transforms.Resize((224, 224)), # 只能对PIL图片进行裁剪
transforms.ToTensor(),
])
img = transform(img)
img = img.unsqueeze(0)
# torch.cuda.set_device(devices[0])
net = PolygonNet(load_vgg=False)
# net = nn.DataParallel(net, device_ids=devices).cuda()
net.load_gpu_version(torch.load(args.model, map_location='cpu'))
# print('finished')
net.eval()
dtype = torch.FloatTensor
dtype_t = torch.LongTensor
# output_path = "D:\\MyProject\\annotation_helper\\output\\output.txt"
# os.system("rm -rf {}".format(output_path))
# selected_classes = ['bicycle', 'bus', 'person', 'train', 'truck', 'motorcycle', 'car', 'rider']
total = []
total_wto_error = []
error_num = 0
error_flag = False
with torch.no_grad():
result = net.test(img.type(dtype), 60)
labels_p = result.cpu().numpy()
# print(labels_p)
for i in range(result.shape[0]):
vertices1 = []
k = 0
for k in range(len(labels_p[i])):
if labels_p[i][k] == 784:
k += 1
else:
break
for p in range(k, len(labels_p[i])):
if labels_p[i][p] == 784:
break
vertex = (
round(((labels_p[i][p] % 28) * 8.0 + 4) * img_len / 224),
float(
round((max(0,
(int(labels_p[i][p] / 28) - 1)) * 8.0 + 4) *
img_width / 224)))
vertices1.append(vertex)
# with open(output_path, 'a') as outfile:
# print(vertices1, file=outfile)
# print(type(vertices1))
ret = []
for i in range(len(vertices1)):
ret.append(vertices1[i][0])
ret.append(vertices1[i][1])
# print(ret)
return ret
parser = argparse.ArgumentParser(description='manual to this script')
parser.add_argument('--gpu_id', nargs='+', type=int)
parser.add_argument('--batch_size', type=int, default=8)
parser.add_argument('--num', type=int, default=45000)
parser.add_argument('--test_mode', type=str, default='full')
parser.add_argument('--model', type=str, default='./save/model.pth')
args = parser.parse_args()
devices = args.gpu_id
print(devices)
batch_size = args.batch_size
num = args.num
torch.cuda.set_device(devices[0])
net = PolygonNet(load_vgg=False)
net = nn.DataParallel(net, device_ids=devices)
net.load_state_dict(torch.load(args.model))
net.cuda()
print('finished')
dtype = torch.cuda.FloatTensor
dtype_t = torch.cuda.LongTensor
if args.test_mode == 'small':
Dataloader = load_data(num, 'trainval', 600, batch_size)
len_dl = len(Dataloader)
print(len_dl)
nu = 0
def main():
# Test settings
parser = argparse.ArgumentParser(description='Segmentation')
parser.add_argument('--ckpt', type=str, default='checkpoint/checkpoint_polygonnet_best.pth.tar', help="path to checkpoint to load")
parser.add_argument('--nlevels', type=int, default=4, help="number of polygon levels, higher->finer")
parser.add_argument('--feat', type=int, default=256, help="number of base feature layers")
parser.add_argument('--dropout', action='store_true', help="dropout during training")
parser.add_argument('--raw_img_dir', type=str, default='leftImg8bit', help='path to raw image directory')
parser.add_argument('--polyrnn2_dir', type=str, default='polyrnn-pp-pytorch-small')
parser.add_argument('--gpuid', type=int, default=0, help='gpu id to use for evaluation')
parser.add_argument('--output_dir', type=str, default='output_vis_full', help='directory to output images')
parser.add_argument('--nsamples', type=int, default=10, help='number of samples to produce. 0 for all.')
parser.add_argument('--export_gt', action='store_true', help='export ground truth images also')
parser.add_argument('--skip_multicomponents', action='store_true', help='skip multiple components')
parser.add_argument('--rand_samp', action='store_true', help='randomly draw samples to save')
args = parser.parse_args()
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpuid)
device = torch.device("cuda")
# create output directory
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
# get training / valid sets
args.img_mean = [0.485, 0.456, 0.406]
args.img_std = [0.229, 0.224, 0.225]
normalize = torchvision.transforms.Normalize(mean=args.img_mean,
std=args.img_std)
# DO NOT include horizontal and vertical flips in the composed transforms!
transform = torchvision.transforms.Compose([
torchvision.transforms.ColorJitter(hue=.1, saturation=.1),
torchvision.transforms.ToTensor(),
normalize,
])
# load model
model = PolygonNet(nlevels=args.nlevels, dropout=args.dropout, feat=args.feat)
model = nn.DataParallel(model)
model.to(device)
if os.path.isfile(args.ckpt):
print("=> loading checkpoint '{}'".format(args.ckpt))
checkpoint = torch.load(args.ckpt)
args.best_miou = checkpoint['best_miou']
try:
model.load_state_dict(checkpoint['state_dict'])
except:
# simple hack for loading the old model
sdict = checkpoint['state_dict']
sdict['module.projection.1.weight'] = sdict.pop('module.projection.0.weight')
sdict['module.projection.1.bias'] = sdict.pop('module.projection.0.bias')
model.load_state_dict(sdict)
print("=> loaded checkpoint '{}' (epoch {})"
.format(args.ckpt, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.ckpt))
# create and save visualizations
args.json_dir = os.path.join(args.polyrnn2_dir, 'data', 'cityscapes_final_v5')
vis = visualizer(model=model, transform=transform, json_dir=args.json_dir, img_dir=args.raw_img_dir, skip_multicomponents=args.skip_multicomponents, export_gt=args.export_gt, rand_samp=args.rand_samp)
if args.nsamples < 0:
args.nsamples = len(vis)
for i in tqdm(range(args.nsamples)):
if args.export_gt:
pd_img, gt_img, imgpath = vis[i]
else:
pd_img, imgpath = vis[i]
img_path = imgpath.split('/')[-1].split('.')[0]
pd_save = os.path.join(args.output_dir, img_path+'_PD.png')
plt.imsave(pd_save, pd_img)
if args.export_gt:
gt_save = os.path.join(args.output_dir, img_path+'_GT.png')
plt.imsave(gt_save, gt_img)
def main():
# Training settings
parser = argparse.ArgumentParser(description='Segmentation')
parser.add_argument('--test_batch_size',
type=int,
default=64,
metavar='N',
help='input batch size for test (default: 64)')
parser.add_argument('--loss_type',
type=str,
choices=['l1', 'l2'],
default='l1',
help='type of loss for raster loss computation')
parser.add_argument('--nlevels',
type=int,
default=4,
help="number of polygon levels, higher->finer")
parser.add_argument('--feat',
type=int,
default=256,
help="number of base feature layers")
parser.add_argument('--dropout',
action='store_true',
help="dropout during training")
parser.add_argument('--no_cuda',
action='store_true',
default=False,
help='disables CUDA training')
parser.add_argument('--seed',
type=int,
default=1,
metavar='S',
help='random seed (default: 1)')
parser.add_argument(
'--data_folder',
type=str,
default="mres_processed_data",
help='path to data folder (default: mres_processed_data)')
parser.add_argument(
'--ckpt',
type=str,
default='checkpoint/checkpoint_polygonnet_best.pth.tar',
help="path to checkpoint to load")
parser.add_argument('--output_dir',
type=str,
default='output_vis_ins',
help="directory to output visualizations")
parser.add_argument('--nsamples',
type=int,
default=0,
help='number of samples to produce.')
parser.add_argument('--multires',
action='store_true',
help="multiresolution loss")
parser.add_argument('--transpose',
action='store_true',
help="transpose target")
parser.add_argument('--smooth_loss',
default=1.0,
type=float,
help="smoothness loss multiplier (0 for none)")
parser.add_argument('--uniform_loss',
action='store_true',
help="use same loss regardless of category frequency")
parser.add_argument('--workers',
default=12,
type=int,
help="number of data loading workers")
args = parser.parse_args()
use_cuda = not args.no_cuda and torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
args.label_names = [
"car", "truck", "train", "bus", "motorcycle", "bicycle", "rider",
"person"
]
args.instances = torch.tensor(
[30246, 516, 76, 325, 658, 3307, 1872, 19563]).double().to(device)
args.nclass = len(args.label_names)
args.weights = 1 / torch.log(args.instances / args.instances.sum() + 1.02)
args.weights /= args.weights.sum()
if not os.path.exists(args.output_dir) and args.nsamples > 0:
os.makedirs(args.output_dir)
# PYTORCH VERSION > 1.0.0
assert (float(torch.__version__.split('.')[-3]) > 0)
torch.manual_seed(args.seed)
# get training / valid sets
args.img_mean = [0.485, 0.456, 0.406]
args.img_std = [0.229, 0.224, 0.225]
normalize = torchvision.transforms.Normalize(mean=args.img_mean,
std=args.img_std)
# DO NOT include horizontal and vertical flips in the composed transforms!
transform = torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
normalize,
])
testset = CityScapeLoader(args.data_folder,
"test",
transforms=transform,
RandomHorizontalFlip=0.0,
RandomVerticalFlip=0.0,
mres=args.multires,
transpose=args.transpose)
test_loader = DataLoader(testset,
batch_size=args.test_batch_size,
shuffle=True,
drop_last=False,
num_workers=args.workers,
pin_memory=True)
# initialize and parallelize model
model = PolygonNet(nlevels=args.nlevels,
dropout=args.dropout,
feat=args.feat)
model = nn.DataParallel(model)
model.to(device)
# Multi-Resolution
n = model.module.npoints
if args.multires:
args.res = [224, 112, 56, 28]
args.npt = [n, int(n / 2), int(n / 4), int(n / 8)]
args.levels = [1, 2, 4, 8]
else:
args.res = [224]
args.npt = [n]
args.levels = [1]
if os.path.isfile(args.ckpt):
print("=> loading checkpoint '{}'".format(args.ckpt))
checkpoint = torch.load(args.ckpt)
args.best_miou = checkpoint['best_miou']
model.load_state_dict(checkpoint['state_dict'])
print("=> loaded checkpoint '{}' (epoch {})".format(
args.ckpt, checkpoint['epoch']))
else:
print("=> no checkpoint found at '{}'".format(args.ckpt))
print("{} paramerters in total".format(
sum(x.numel() for x in model.parameters())))
rres = [True, False, False, False] if args.multires else [True]
criterion = [
BatchedRasterLoss2D(npoints=n,
res=r,
loss=args.loss_type,
return_raster=rr).to(device)
for n, r, rr in zip(args.npt, args.res, rres)
]
criterion_smooth = SmoothnessLoss()
evaluate(args, model, test_loader, criterion, criterion_smooth,
checkpoint['epoch'], device)
batch_size = args.batch_size
root_path = "/Disk8/kevin/Cityscapes/"
mapping_location = {
'cuda:1': 'cuda:' + str(devices[0]),
'cuda:2': 'cuda:' + str(devices[0]),
'cuda:3': 'cuda:' + str(devices[0]),
'cuda:4': 'cuda:' + str(devices[0]),
'cuda:5': 'cuda:' + str(devices[0]),
'cuda:6': 'cuda:' + str(devices[0]),
'cuda:7': 'cuda:' + str(devices[0]),
'cuda:0': 'cuda:' + str(devices[0])
}
torch.cuda.set_device(devices[0])
net = PolygonNet(load_vgg=False)
net = nn.DataParallel(net, device_ids=devices).cuda()
net.load_state_dict(torch.load(args.model, map_location=mapping_location))
print('finished')
net.eval()
dtype = torch.cuda.FloatTensor
dtype_t = torch.cuda.LongTensor
output_path = "/home/kevin/polygon-release/test_log"
os.system("rm -rf {}".format(output_path))
selected_classes = [
'bicycle', 'bus', 'person', 'train', 'truck', 'motorcycle', 'car', 'rider'
]
parser.add_argument('--gpu_id', nargs='+', type=int)
parser.add_argument(
'--num',
type=int,
)
parser.add_argument('--model', type=str)
parser.add_argument('--dataset', type=str)
parser.add_argument('--config', dest='config_file', help='Config File')
args = parser.parse_args()
config_from_args = args.__dict__
config_file = config_from_args.pop('config_file')
config = get_config('test', config_from_args, config_file)
devices = config['gpu_id']
num = config['num']
dataset = config['dataset']
model = config['model']
print('gpus: {}'.format(devices))
torch.cuda.set_device(devices[0])
net = PolygonNet(load_vgg=False)
net = nn.DataParallel(net, device_ids=devices)
net.load_state_dict(torch.load(config['model']))
net.cuda()
print('Loading completed!')
iou_score = test(net, dataset, num)
print('iou score:{}'.format(iou_score))
parser.add_argument('--num', type=int, default=45000)
parser.add_argument('--lr', type=float, default=0.0001)
args = parser.parse_args()
devices = args.gpu_id
print(devices)
batch_size = args.batch_size
num = args.num
lr = args.lr
torch.cuda.set_device(devices[0])
Dataloader = load_data(num, 'train', 60, batch_size)
len_dl = len(Dataloader)
print(len_dl)
net = PolygonNet()
net = nn.DataParallel(net, device_ids=devices)
if args.pretrained == "True":
net.load_state_dict(torch.load('save/model' + '_' + str(num) + '.pth'))
net.cuda()
print('finished')
loss_function = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=lr)
scheduler = optim.lr_scheduler.MultiStepLR(optimizer,
milestones=[4000, 100000],
gamma=0.1)
writer = SummaryWriter()
dtype = torch.cuda.FloatTensor
batch_size = args.batch_size
num = args.num
lr = args.lr
opt_step = args.opt_step
lbd = 0.8
epoch_r = 5
loss_step = 100
root_path = "/Disk8/kevin/Cityscapes/"
user_path = "/home/kevin/polygon-release"
Dataloader = load_data(num, 'train', 60, root_path, batch_size)
len_dl = len(Dataloader)
torch.cuda.set_device(devices[0])
net = PolygonNet().cuda()
net = nn.DataParallel(net, device_ids=devices)
mapping_location = {'cuda:1': 'cuda:' + str(devices[0]), 'cuda:2': 'cuda:' + str(devices[0]),
'cuda:3': 'cuda:' + str(devices[0]), 'cuda:4': 'cuda:' + str(devices[0]),
'cuda:5': 'cuda:' + str(devices[0]), 'cuda:6': 'cuda:' + str(devices[0]),
'cuda:7': 'cuda:' + str(devices[0]), 'cuda:0': 'cuda:' + str(devices[0])}
if args.pretrained == 'True':
print('Loading model...')
net.load_state_dict(torch.load(root_path + 'save/4_50000.pth', map_location=mapping_location))
else:
os.system('rm -rf {}/log'.format(user_path))
if not os.path.isfile('{}/control'.format(user_path)):
os.system('touch {}/control'.format(user_path))
def deep_rnn_annotate():
parser = argparse.ArgumentParser(description='manual to test script')
parser.add_argument('--batch_size', type=int, default=1)
parser.add_argument('--model', type=str, default='./Polygon_deep_RNN.pth')
args = parser.parse_args()
devices = [0]
batch_size = args.batch_size
root_path = "../images"
img_file = os.path.join(root_path, 'save.jpg')
img = Image.open(img_file).convert('RGB')
img_len = img.size[0]
img_width = img.size[1]
transform = transforms.Compose([
transforms.Resize((224, 224)),
transforms.ToTensor(),
])
img = transform(img)
img = img.unsqueeze(0)
net = PolygonNet(load_vgg=False)
net.load_gpu_version(torch.load(args.model, map_location='cpu'))
net.eval()
dtype = torch.FloatTensor
dtype_t = torch.LongTensor
total = []
total_wto_error = []
error_num = 0
error_flag = False
with torch.no_grad():
result = net.test(img.type(dtype), 60)
labels_p = result.cpu().numpy()
for i in range(result.shape[0]):
vertices1 = []
k = 0
for k in range(len(labels_p[i])):
if labels_p[i][k] == 784:
k += 1
else:
break
for p in range(k, len(labels_p[i])):
if labels_p[i][p] == 784:
break
vertex = (
round(((labels_p[i][p] % 28) * 8.0 + 4) * img_len / 224),
float(
round((max(0,
(int(labels_p[i][p] / 28) - 1)) * 8.0 + 4) *
img_width / 224)))
vertices1.append(vertex)
ret = []
for i in range(len(vertices1)):
ret.append(vertices1[i][0])
ret.append(vertices1[i][1])
return ret
from glob import glob
parser = argparse.ArgumentParser(description='manual to this script')
parser.add_argument('--gpu_id', nargs='+', type=int)
parser.add_argument('--batch_size', type=int, default=8)
parser.add_argument('--num', type=int, default=45000)
parser.add_argument('--test_mode', type=str, default='full')
args = parser.parse_args()
devices = args.gpu_id
print(devices)
batch_size = args.batch_size
num = args.num
torch.cuda.set_device(devices[0])
net = PolygonNet(load_vgg=False)
net = nn.DataParallel(net, device_ids=devices)
net.load_state_dict(torch.load('../polygon/save/third3.pth'))
net.cuda()
print('finished')
dtype = torch.cuda.FloatTensor
dtype_t = torch.cuda.LongTensor
if args.test_mode == 'small':
Dataloader = load_data(num, 'trainval', 600, batch_size)
len_dl = len(Dataloader)
print(len_dl)
nu = 0
def train(config, pretrained=None):
devices = config['gpu_id']
batch_size = config['batch_size']
lr = config['lr']
log_dir = config['log_dir']
prefix = config['prefix']
num = config['num']
print('Using gpus: {}'.format(devices))
torch.cuda.set_device(devices[0])
Dataloader = load_data(num, 'train', 60, batch_size)
len_dl = len(Dataloader)
print(len_dl)
net = PolygonNet()
net = nn.DataParallel(net, device_ids=devices)
if pretrained:
net.load_state_dict(torch.load(pretrained))
net.cuda()
print('Loading completed!')
loss_function = nn.CrossEntropyLoss()
optimizer = optim.Adam(net.parameters(), lr=lr)
scheduler = optim.lr_scheduler.MultiStepLR(optimizer,
milestones=[4000, 100000],
gamma=0.1)
writer = SummaryWriter(log_dir)
dtype = torch.cuda.FloatTensor
dtype_t = torch.cuda.LongTensor
epoch_r = int(300000 / len_dl)
for epoch in range(epoch_r):
for step, data in enumerate(Dataloader):
scheduler.step()
x = Variable(data[0].type(dtype))
x1 = Variable(data[1].type(dtype))
x2 = Variable(data[2].type(dtype))
x3 = Variable(data[3].type(dtype))
ta = Variable(data[4].type(dtype_t))
optimizer.zero_grad()
r = net(x, x1, x2, x3)
result = r.contiguous().view(-1, 28 * 28 + 3)
target = ta.contiguous().view(-1)
loss = loss_function(result, target)
loss.backward()
result_index = torch.argmax(result, 1)
correct = (target == result_index).type(dtype).sum().item()
acc = correct * 1.0 / target.shape[0]
# scheduler.step(loss)
optimizer.step()
writer.add_scalar('train/loss', loss, epoch * len_dl + step)
writer.add_scalar('train/accuracy', acc, epoch * len_dl + step)
if step % 100 == 0:
torch.save(net.state_dict(), prefix + '_' + str(num) + '.pth')
# for param_group in optimizer.param_groups:
# print(
# 'epoch{} step{}:{}'.format(epoch, step,
# param_group['lr']))
train_iou = test(net, 'train', 10)
val_iou = test(net, 'val', 10)
for key, val in train_iou.items():
writer.add_scalar('train/iou_{}'.format(key), val, epoch * len_dl)
for key, val in val_iou.items():
writer.add_scalar('val/iou_{}'.format(key), val, epoch * len_dl)
print('iou score on training set:{}'.format(train_iou))
print('iou score on test set:{}'.format(val_iou))
if epoch % 5 == 0 and len_dl > 200:
torch.save(net.state_dict(),
prefix + str(epoch) + '_' + str(num) + '.pth')
writer.close()