def main(num_classes, context_path, check_points):
os.environ['CUDA_VISIBLE_DEVICES'] = '7'
model = BiSeNet(num_classes, context_path).cuda()
#if torch.cuda.is_available():
#model.load_state_dict(torch.load(check_points))
#model = torch.nn.DataParallel(model).cuda()
#model.load_state_dict(torch.load(check_points))
# load pretrained model if exists
#model.eval()
# load pretrained model if exists
#print('load model from %s ...' % check_points)
#model.load_state_dict(torch.load(check_points))
print('load model from %s ...' % check_points)
model.load_state_dict(torch.load(check_points))
print('load model success')
input_var = torch.randn(1, 3, 480, 640, device='cuda')
#output_var = model(input_var)
input_names = ["input1"]
output_names = ["output1"]
torch.onnx.export(model,
input_var,
'Bisenet_nearest.onnx',
verbose=True,
input_names=input_names,
output_names=output_names)
def main(params):
# basic parameters
parser = argparse.ArgumentParser()
parser.add_argument('--data', type=str, default='/home/disk2/xs/sun', help='path of training data')
parser.add_argument('--num_epochs', type=int, default=300, help='Number of epochs to train for')
parser.add_argument('--epoch_start_i', type=int, default=0, help='Start counting epochs from this number')
parser.add_argument('--checkpoint_step', type=int, default=10, help='How often to save checkpoints (epochs)')
parser.add_argument('--validation_step', type=int, default=1, help='How often to perform validation (epochs)')
parser.add_argument('--dataset', type=str, default='SUN', help='Dataset you are using.')
parser.add_argument('--crop_height', type=int, default=480, help='Height of cropped/resized input image to network')
parser.add_argument('--crop_width', type=int, default=640, help='Width of cropped/resized input image to network')
parser.add_argument('--batch_size', type=int, default=5, help='Number of images in each batch')
parser.add_argument('--context_path', type=str, default="resnet101", help='The context path model you are using.')
parser.add_argument('--learning_rate', type=float, default=0.001, help='learning rate used for train')
parser.add_argument('--num_workers', type=int, default=4, help='num of workers')
parser.add_argument('--num_classes', type=int, default=38, help='num of object classes (with void)')
parser.add_argument('--cuda', type=str, default='2', help='GPU ids used for training')
parser.add_argument('--use_gpu', type=bool, default=True, help='whether to user gpu for training')
parser.add_argument('--pretrained_model_path', type=str, default=None, help='path to pretrained model')
parser.add_argument('--save_model_path', type=str, default='./checkpoints', help='path to save model')
parser.add_argument('--csv_path', type=str, default='/home/disk2/xs/sun/seg37_class_dict.csv', help='Path to label info csv file')
args = parser.parse_args(params)
# create dataset and dataloader
train_img_path = os.path.join(args.data, 'train/image')
train_depth_path = os.path.join(args.data, 'train/depth')
train_label_path = os.path.join(args.data, 'train/label')
csv_path = os.path.join(args.data, 'seg37_class_dict.csv')
dataset_train = SUN(train_img_path, train_depth_path, train_label_path, csv_path, scale=(args.crop_height, args.crop_width), mode='train')
dataloader_train = DataLoader(
dataset_train,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers
)
# build model
os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda
model = BiSeNet(args.num_classes, args.context_path)
if torch.cuda.is_available() and args.use_gpu:
model = torch.nn.DataParallel(model).cuda()
# build optimizer
optimizer = torch.optim.RMSprop(model.parameters(), args.learning_rate)
# load pretrained model if exists
if args.pretrained_model_path is not None:
print('load model from %s ...' % args.pretrained_model_path)
model.module.load_state_dict(torch.load(args.pretrained_model_path))
print('Done!')
# train
train(args, model, optimizer, dataloader_train, csv_path)
def main(params):
# basic parameters
parser = argparse.ArgumentParser()
parser.add_argument('--checkpoint_path',
type=str,
default=None,
required=True,
help='The path to the pretrained weights of model')
parser.add_argument(
'--crop_height',
type=int,
default=640,
help='Height of cropped/resized input image to network')
parser.add_argument('--crop_width',
type=int,
default=640,
help='Width of cropped/resized input image to network')
parser.add_argument('--data',
type=str,
default='/path/to/data',
help='Path of training data')
parser.add_argument('--batch_size',
type=int,
default=1,
help='Number of images in each batch')
parser.add_argument('--context_path',
type=str,
default="resnet101",
help='The context path model you are using.')
parser.add_argument('--cuda',
type=str,
default='0',
help='GPU ids used for training')
parser.add_argument('--use_gpu',
type=bool,
default=True,
help='Whether to user gpu for training')
parser.add_argument('--num_classes',
type=int,
default=2,
help='num of object classes (with void)')
args = parser.parse_args(params)
# create dataset and dataloader
dataloader = DataLoader(test(input_transform, target_transform),
num_workers=1,
batch_size=1,
shuffle=True)
# build model
os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda
model = BiSeNet(args.num_classes, args.context_path)
if torch.cuda.is_available() and args.use_gpu:
model = torch.nn.DataParallel(model).cuda()
# load pretrained model if exists
print('load model from %s ...' % args.checkpoint_path)
model.module.load_state_dict(torch.load(args.checkpoint_path))
print('Done!')
# test
eval(model, dataloader, args)
def main(params):
# parse the parameters
parser = argparse.ArgumentParser()
parser = add_arguments(parser)
args = parser.parse_args(params)
print("Training with following arguments:")
pprint(vars(args), indent=4, compact=True)
print("Running on: {}".format(device if args.use_gpu else torch.device('cpu')))
# create dataset and dataloader
train_path = args.data
train_transform, val_transform = get_transform(
args.random_crop_size, args.further_data_aug)
dataset_train = VOC(train_path, image_set="train",
transform=train_transform)
dataset_val = VOC(train_path, image_set="val", transform=val_transform)
dataloader_train = DataLoader(dataset_train, batch_size=args.batch_size,
shuffle=True, num_workers=args.num_workers, drop_last=True, )
dataloader_val = DataLoader(dataset_val, batch_size=args.batch_size,
shuffle=True, num_workers=args.num_workers, )
# build model
os.environ["CUDA_VISIBLE_DEVICES"] = args.cuda
model = BiSeNet(args.num_classes, args.context_path)
if args.use_gpu:
model = model.to(device)
# build optimizer
optimizer = get_optim(args, model)
# load pretrained model if exists
if args.pretrained_model_path is not None:
print("load model from %s ..." % args.pretrained_model_path)
model.load_state_dict(torch.load(args.pretrained_model_path))
print("Done!")
scaler = amp.GradScaler() if args.use_amp else None
# train
train(args, model, optimizer, dataloader_train, dataloader_val, scaler)
print("Training completed.", datetime.now().strftime("%m/%d/%Y, %H:%M:%S"))
parser.add_argument('--crop_height', type=int, default=480, help='Height of cropped/resized input image to network')
parser.add_argument('--crop_width', type=int, default=640, help='Width of cropped/resized input image to network')
parser.add_argument('--cuda', type=str, default='1', help='GPU ids used for training')
parser.add_argument('--use_gpu', type=bool, default=True, help='Whether to user gpu for training')
parser.add_argument('--csv_path', type=str, default='/home/disk1/xs/ADEChallengeData2016/ade150.csv', help='Path to label info csv file')
args = parser.parse_args()
# read csv label path
label_info = get_label_info(args.csv_path)
scale = (args.crop_height, args.crop_width)
# build model
os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda
model = BiSeNet(args.num_classes, args.context_path)
# load pretrained model if exists
print('load model from %s ...' % args.checkpoint_path)
if torch.cuda.is_available() and args.use_gpu:
model = torch.nn.DataParallel(model).cuda()
model.module.load_state_dict(torch.load(args.checkpoint_path)) # GPU -> GPU
else:
model.load_state_dict(torch.load(args.checkpoint_path, map_location=lambda storage, loc: storage)) # GPU -> CPU
print('Done!')
resize_img = transforms.Resize(scale, Image.BILINEAR)
resize_depth = transforms.Resize(scale, Image.NEAREST)
to_tensor = transforms.ToTensor()
def main(params):
# basic parameters
parser = argparse.ArgumentParser()
parser.add_argument('--checkpoint_path',
type=str,
default=None,
required=True,
help='The path to the pretrained weights of model')
parser.add_argument(
'--crop_height',
type=int,
default=720,
help='Height of cropped/resized input image to network')
parser.add_argument('--crop_width',
type=int,
default=960,
help='Width of cropped/resized input image to network')
parser.add_argument('--data',
type=str,
default='/path/to/data',
help='Path of training data')
parser.add_argument('--batch_size',
type=int,
default=1,
help='Number of images in each batch')
parser.add_argument('--context_path',
type=str,
default="resnet101",
help='The context path model you are using.')
parser.add_argument('--cuda',
type=str,
default='0',
help='GPU ids used for training')
parser.add_argument('--use_gpu',
type=bool,
default=True,
help='Whether to user gpu for training')
parser.add_argument('--num_classes',
type=int,
default=32,
help='num of object classes (with void)')
parser.add_argument('--loss',
type=str,
default='dice',
help='loss function, dice or crossentropy')
args = parser.parse_args(params)
# build model
os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda
model = BiSeNet(args.num_classes, args.context_path)
if torch.cuda.is_available() and args.use_gpu:
model = torch.nn.DataParallel(model).cuda()
# load pretrained model if exists
print('load model from %s ...' % args.checkpoint_path)
model.module.load_state_dict(torch.load(args.checkpoint_path))
print('Done!')
rospy.init_node('obj_seg', anonymous=True)
Seg = object_segmentation(model)
rospy.spin()
def main(params):
# basic parameters
parser = argparse.ArgumentParser()
parser.add_argument('--image',
action='store_true',
default=False,
help='predict on image')
parser.add_argument('--video',
action='store_true',
default=False,
help='predict on video')
parser.add_argument('--checkpoint_path',
type=str,
default=None,
help='The path to the pretrained weights of model')
parser.add_argument('--context_path',
type=str,
default="resnet101",
help='The context path model you are using.')
parser.add_argument('--num_classes',
type=int,
default=12,
help='num of object classes (with void)')
parser.add_argument('--data',
type=str,
default=None,
help='Path to image or video for prediction')
parser.add_argument(
'--crop_height',
type=int,
default=720,
help='Height of cropped/resized input image to network')
parser.add_argument('--crop_width',
type=int,
default=960,
help='Width of cropped/resized input image to network')
parser.add_argument('--cuda',
type=str,
default='0',
help='GPU ids used for training')
parser.add_argument('--use_gpu',
type=bool,
default=True,
help='Whether to user gpu for training')
parser.add_argument('--csv_path',
type=str,
default=None,
required=True,
help='Path to label info csv file')
parser.add_argument('--save_path',
type=str,
default=None,
required=True,
help='Path to save predict image')
args = parser.parse_args(params)
# build model
os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda
net = BiSeNet(args.num_classes, args.context_path)
if torch.cuda.is_available() and args.use_gpu:
net = torch.nn.DataParallel(net).cuda()
# load pretrained model if exists
print('load model from %s ...' % args.checkpoint_path)
# model.module.load_state_dict(torch.load(args.checkpoint_path))
net.load_state_dict(torch.load(args.checkpoint_path, map_location='cpu'))
print('Done!')
# predict on image
if args.image:
predict_on_image(net, args)
# predict on video
if args.video:
pass
def main(params):
# basic parameters
parser = argparse.ArgumentParser()
parser.add_argument('--num_epochs',
type=int,
default=300,
help='Number of epochs to train for')
parser.add_argument('--epoch_start_i',
type=int,
default=0,
help='Start counting epochs from this number')
parser.add_argument('--checkpoint_step',
type=int,
default=10,
help='How often to save checkpoints (epochs)')
parser.add_argument('--validation_step',
type=int,
default=10,
help='How often to perform validation (epochs)')
parser.add_argument('--batch_size',
type=int,
default=1,
help='Number of images in each batch')
parser.add_argument(
'--context_path',
type=str,
default="resnet101",
help='The context path model you are using, resnet18, resnet101.')
parser.add_argument('--learning_rate',
type=float,
default=0.01,
help='learning rate used for train')
parser.add_argument('--data',
type=str,
default='data',
help='path of training data')
parser.add_argument('--num_workers',
type=int,
default=4,
help='num of workers')
parser.add_argument('--num_classes',
type=int,
default=32,
help='num of object classes (with void)')
parser.add_argument('--cuda',
type=str,
default='0',
help='GPU ids used for training')
parser.add_argument('--use_gpu',
type=bool,
default=True,
help='whether to user gpu for training')
parser.add_argument('--pretrained_model_path',
type=str,
default=None,
help='path to pretrained model')
parser.add_argument('--save_model_path',
type=str,
default="checkpoints",
help='path to save model')
parser.add_argument('--optimizer',
type=str,
default='rmsprop',
help='optimizer, support rmsprop, sgd, adam')
parser.add_argument('--loss',
type=str,
default='crossentropy',
help='loss function, dice or crossentropy')
# settiamo i nostri parametri
args = parser.parse_args(params)
# create dataset and dataloader
train_path = args.data
train_transform, val_transform = get_transform()
# creiamo un oggetto di tipo VOC per il training
dataset_train = VOC(train_path,
image_set="train",
transform=train_transform)
dataloader_train = DataLoader(dataset_train,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
drop_last=True)
# creiamo un oggetto di tipo VOC per la validation
dataset_val = VOC(train_path, image_set="val", transform=val_transform)
dataloader_val = DataLoader(dataset_val,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers)
# build model
os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda
model = BiSeNet(args.num_classes, args.context_path)
if torch.cuda.is_available() and args.use_gpu:
model = model.cuda()
# build optimizer
if args.optimizer == 'rmsprop':
optimizer = torch.optim.RMSprop(model.parameters(), args.learning_rate)
elif args.optimizer == 'sgd':
optimizer = torch.optim.SGD(model.parameters(),
args.learning_rate,
momentum=0.9,
weight_decay=1e-4)
elif args.optimizer == 'adam':
optimizer = torch.optim.Adam(model.parameters(), args.learning_rate)
else: # rmsprop
print('not supported optimizer \n')
return None
# load pretrained model if exists
# Non ce l'abbiamo
if args.pretrained_model_path is not None:
print('load model from %s ...' % args.pretrained_model_path)
model.load_state_dict(torch.load(args.pretrained_model_path))
print('Done!')
# train
# funzioni presenti in questo file
train(args, model, optimizer, dataloader_train, dataloader_val)
val(args, model, dataloader_val)
def main(params):
parser = argparse.ArgumentParser()
parser.add_argument('--save_model_path',
type=str,
default=None,
help='path to save model')
parser.add_argument('--num_classes',
type=int,
default=32,
help='num of object classes (with void)')
parser.add_argument(
'--context_path',
type=str,
default="resnet18",
help='The context path model you are using, resnet18, resnet101.')
args = parser.parse_args(params)
model = BiSeNet(args.num_classes, args.context_path)
model.load_state_dict(
torch.load(os.path.join(args.save_model_path, 'best_dice_loss.pth')))
model.eval()
img = Image.open('./CamVid/test/Seq05VD_f00660.png')
transform = transforms.Compose([
transforms.Resize([720, 960]),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
img = transform(img).unsqueeze(dim=0)
imresult = np.zeros([img.shape[2], img.shape[3], 3], dtype=np.uint8)
with torch.no_grad():
predict = model(img).squeeze()
predict = reverse_one_hot(predict)
predict = np.array(predict)
imresult[:, :][predict == 0] = [255, 51, 255]
imresult[:, :][predict == 1] = [255, 0, 0]
imresult[:, :][predict == 2] = [0, 255, 0]
imresult[:, :][predict == 3] = [0, 0, 255]
imresult[:, :][predict == 4] = [255, 255, 0]
imresult[:, :][predict == 5] = [255, 0, 255]
imresult[:, :][predict == 6] = [0, 255, 255]
imresult[:, :][predict == 7] = [10, 200, 128]
imresult[:, :][predict == 8] = [125, 18, 78]
imresult[:, :][predict == 9] = [205, 128, 8]
imresult[:, :][predict == 10] = [144, 208, 18]
imresult[:, :][predict == 11] = [5, 88, 198]
cv2.imwrite('result.png', imresult)
print('inference done')
summary(model, (3, 720, 960), 1, "cpu")
macs, params = profile(model, inputs=(img, ))
print('macs', macs)
print('params', params)
log = open('log.txt', 'w')
EXPORTONNXNAME = 'nit-bisenet.onnx'
try:
torch.onnx.export(
model,
img,
EXPORTONNXNAME,
export_params=True,
do_constant_folding=True,
input_names=['data'],
# output_names = ['output']
output_names=['output'])
except Exception:
traceback.print_exc(file=log)
print('export done')
def main(params):
# basic parameters
parser = argparse.ArgumentParser()
parser.add_argument('--num_epochs',
type=int,
default=30,
help='Number of epochs to train for')
parser.add_argument('--epoch_start_i',
type=int,
default=0,
help='Start counting epochs from this number')
parser.add_argument('--checkpoint_step',
type=int,
default=1,
help='How often to save checkpoints (epochs)')
parser.add_argument('--validation_step',
type=int,
default=1,
help='How often to perform validation (epochs)')
parser.add_argument('--dataset',
type=str,
default="CamVid",
help='Dataset you are using.')
parser.add_argument(
'--crop_height',
type=int,
default=640,
help='Height of cropped/resized input image to network')
parser.add_argument('--crop_width',
type=int,
default=640,
help='Width of cropped/resized input image to network')
parser.add_argument('--batch_size',
type=int,
default=1,
help='Number of images in each batch')
parser.add_argument('--context_path',
type=str,
default="resnet101",
help='The context path model you are using.')
parser.add_argument('--learning_rate',
type=float,
default=0.01,
help='learning rate used for train')
parser.add_argument('--data',
type=str,
default='/path/to/data',
help='path of training data')
parser.add_argument('--num_workers',
type=int,
default=4,
help='num of workers')
parser.add_argument('--num_classes',
type=int,
default=2,
help='num of object classes (with void)')
parser.add_argument('--cuda',
type=str,
default='0',
help='GPU ids used for training')
parser.add_argument('--use_gpu',
type=bool,
default=True,
help='whether to user gpu for training')
parser.add_argument('--pretrained_model_path',
type=str,
default=None,
help='path to pretrained model')
parser.add_argument('--save_model_path',
type=str,
default=None,
help='path to save model')
args = parser.parse_args(params)
# create dataset and dataloader
dataloader_train = DataLoader(train(input_transform, target_transform),
num_workers=1,
batch_size=2,
shuffle=True)
# build model
os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda
model = BiSeNet(args.num_classes, args.context_path)
if torch.cuda.is_available() and args.use_gpu:
model = torch.nn.DataParallel(model).cuda()
para = sum([np.prod(list(p.size())) for p in model.parameters()])
print('Model {} : params: {:4f}M'.format(model._get_name(),
para * 4 / 1000 / 1000))
# build optimizer
optimizer = torch.optim.RMSprop(model.parameters(), args.learning_rate)
# load pretrained model if exists
if args.pretrained_model_path is not None:
print('load model from %s ...' % args.pretrained_model_path)
model.module.load_state_dict(torch.load(args.pretrained_model_path))
print('Done!')
# train
train(args, model, optimizer, dataloader_train)
def main(params):
# basic parameters
parser = argparse.ArgumentParser()
parser.add_argument('--num_epochs', type=int, default=300, help='Number of epochs to train for')
parser.add_argument('--epoch_start_i', type=int, default=0, help='Start counting epochs from this number')
parser.add_argument('--checkpoint_step', type=int, default=1000, help='How often to save checkpoints (epochs)')
parser.add_argument('--validation_step', type=int, default=100, help='How often to perform validation (epochs)')
parser.add_argument('--dataset', type=str, default="CamVid", help='Dataset you are using.')
parser.add_argument('--crop_height', type=int, default=720, help='Height of cropped/resized input image to network')
parser.add_argument('--crop_width', type=int, default=960, help='Width of cropped/resized input image to network')
parser.add_argument('--batch_size', type=int, default=1, help='Number of images in each batch')
parser.add_argument('--context_path', type=str, default="resnet101",
help='The context path model you are using, resnet18, resnet101.')
parser.add_argument('--learning_rate', type=float, default=0.01, help='learning rate used for train')
parser.add_argument('--data', type=str, default='', help='path of training data')
parser.add_argument('--num_workers', type=int, default=4, help='num of workers')
parser.add_argument('--num_classes', type=int, default=32, help='num of object classes (with void)')
parser.add_argument('--cuda', type=str, default='0', help='GPU ids used for training')
parser.add_argument('--use_gpu', type=bool, default=True, help='whether to user gpu for training')
parser.add_argument('--pretrained_model_path', type=str, default=None, help='path to pretrained model')
parser.add_argument('--save_model_path', type=str, default=None, help='path to save model')
parser.add_argument('--optimizer', type=str, default='rmsprop', help='optimizer, support rmsprop, sgd, adam')
parser.add_argument('--loss', type=str, default='dice', help='loss function, dice or crossentropy')
args = parser.parse_args(params)
# create dataset and dataloader
train_path = [os.path.join(args.data, 'train'), os.path.join(args.data, 'val')]
train_label_path = [os.path.join(args.data, 'train_labels'), os.path.join(args.data, 'val_labels')]
test_path = os.path.join(args.data, 'test')
test_label_path = os.path.join(args.data, 'test_labels')
csv_path = os.path.join(args.data, 'class_dict.csv')
dataset_train = CamVid(train_path, train_label_path, csv_path, scale=(args.crop_height, args.crop_width),
loss=args.loss, mode='train')
dataloader_train = DataLoader(
dataset_train,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
drop_last=True
)
dataset_val = CamVid(test_path, test_label_path, csv_path, scale=(args.crop_height, args.crop_width),
loss=args.loss, mode='test')
dataloader_val = DataLoader(
dataset_val,
# this has to be 1
batch_size=1,
shuffle=True,
num_workers=args.num_workers
)
# build model
os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda
model = BiSeNet(args.num_classes, args.context_path)
if torch.cuda.is_available() and args.use_gpu:
print('Training using a GPU')
model = torch.nn.DataParallel(model).cuda()
# build optimizer
if args.optimizer == 'rmsprop':
optimizer = torch.optim.RMSprop(model.parameters(), args.learning_rate)
elif args.optimizer == 'sgd':
optimizer = torch.optim.SGD(model.parameters(), args.learning_rate, momentum=0.9, weight_decay=1e-4)
elif args.optimizer == 'adam':
optimizer = torch.optim.Adam(model.parameters(), args.learning_rate)
else: # rmsprop
print('not supported optimizer \n')
return None
# load pretrained model if exists
if args.pretrained_model_path is not None:
print('load model from %s ...' % args.pretrained_model_path)
model.module.load_state_dict(torch.load(args.pretrained_model_path))
print('Done!')
# train
train(args, model, optimizer, dataloader_train, dataloader_val)
def main(params):
# basic parameters
parser = argparse.ArgumentParser()
parser.add_argument('--checkpoint_path',
type=str,
default=None,
required=True,
help='The path to the pretrained weights of model')
parser.add_argument(
'--crop_height',
type=int,
default=512,
help='Height of cropped/resized input image to network')
parser.add_argument('--crop_width',
type=int,
default=512,
help='Width of cropped/resized input image to network')
parser.add_argument('--data',
type=str,
default='/path/to/data',
help='Path of training data')
parser.add_argument('--batch_size',
type=int,
default=1,
help='Number of images in each batch')
parser.add_argument('--context_path',
type=str,
default="resnet101",
help='The context path model you are using.')
parser.add_argument('--cuda',
type=str,
default='0',
help='GPU ids used for training')
parser.add_argument('--use_gpu',
type=bool,
default=True,
help='Whether to user gpu for training')
parser.add_argument('--num_classes',
type=int,
default=32,
help='num of object classes (with void)')
parser.add_argument('--loss',
type=str,
default='dice',
help='loss function, dice or crossentropy')
args = parser.parse_args(params)
# create dataset and dataloader
test_path = os.path.join(args.data, 'test')
# test_path = os.path.join(args.data, 'train')
test_label_path = os.path.join(args.data, 'test_label')
# test_label_path = os.path.join(args.data, 'train_labels')
print(test_path, test_label_path)
csv_path = os.path.join(args.data, 'class_dict.csv')
csv_path = 'potsdam_512_IRRG/class_dict_potsdam.csv'
dataset = potsdam(test_path,
test_label_path,
csv_path,
scale=(args.crop_height, args.crop_width),
mode='test')
dataloader = DataLoader(
dataset,
batch_size=1,
shuffle=True,
num_workers=8,
)
# build model
os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda
model = BiSeNet(args.num_classes, args.context_path)
if torch.cuda.is_available() and args.use_gpu:
model = torch.nn.DataParallel(model).cuda()
# load pretrained model if exists
print('load model from %s ...' % args.checkpoint_path)
model.module.load_state_dict(torch.load(args.checkpoint_path))
print('Done!')
# get label info
# label_info = get_label_info(csv_path)
# test
eval(model, dataloader, args, csv_path)
def main(params):
# basic parameters
parser = argparse.ArgumentParser()
parser.add_argument('--num_epochs',
type=int,
default=300,
help='Number of epochs to train for')
parser.add_argument('--epoch_start_i',
type=int,
default=0,
help='Start counting epochs from this number')
parser.add_argument('--checkpoint_step',
type=int,
default=5,
help='How often to save checkpoints (epochs)')
parser.add_argument('--validation_step',
type=int,
default=1,
help='How often to perform validation (epochs)')
parser.add_argument('--dataset',
type=str,
default="CamVid",
help='Dataset you are using.')
parser.add_argument(
'--crop_height',
type=int,
default=640,
help='Height of cropped/resized input image to network')
parser.add_argument('--crop_width',
type=int,
default=640,
help='Width of cropped/resized input image to network')
parser.add_argument('--train_batch_size',
type=int,
default=1,
help='Number of images in each batch')
parser.add_argument('--val_batch_size',
type=int,
default=1,
help='Number of images in each batch')
parser.add_argument('--context_path',
type=str,
default="resnet101",
help='The context path model you are using.')
parser.add_argument('--learning_rate',
type=float,
default=0.01,
help='learning rate used for train')
parser.add_argument('--num_workers',
type=int,
default=4,
help='num of workers')
parser.add_argument('--num_classes',
type=int,
default=32,
help='num of object classes (with void)')
parser.add_argument('--num_char',
type=int,
default=8,
help='num of lincense chars (include background)')
parser.add_argument('--cuda',
type=str,
default='0',
help='GPU ids used for training')
parser.add_argument('--use_gpu',
type=bool,
default=True,
help='whether to user gpu for training')
parser.add_argument('--log_path',
type=str,
default=None,
help='tensorboard path')
parser.add_argument('--pretrained_model_path',
type=str,
default=None,
help='path to pretrained model')
parser.add_argument('--save_model_path',
type=str,
default=None,
help='path to save model')
args = parser.parse_args(params)
if not os.path.exists(args.log_path):
os.makedirs(args.log_path)
if not os.path.exists(args.save_model_path):
os.makedirs(args.save_model_path)
dataset_train = License_Real_seg_pos_train(split='train_without_night',
num_epochs=args.num_epochs)
dataloader_train = DataLoader(dataset_train,
batch_size=args.train_batch_size,
shuffle=True,
pin_memory=True,
num_workers=args.num_workers)
dataset_val = License_Real_seg_pos_val(split='val_without_night',
num_epochs=1)
dataloader_val = DataLoader(dataset_val,
batch_size=args.val_batch_size,
shuffle=False,
pin_memory=True,
num_workers=args.num_workers)
# build model
model = BiSeNet(args.num_classes, args.num_char, args.context_path)
if torch.cuda.is_available() and args.use_gpu:
model = torch.nn.DataParallel(model).cuda()
# build optimizer
# optimizer = torch.optim.RMSprop(model.parameters(), args.learning_rate)
optimizer = torch.optim.Adam(model.parameters(), args.learning_rate)
criterion = torch.nn.CrossEntropyLoss(ignore_index=255)
# load pretrained model if exists
if args.pretrained_model_path is not None:
print('load model from %s ...' % args.pretrained_model_path)
model.module.load_state_dict(torch.load(args.pretrained_model_path),
False)
print('Done!')
# train
train(args, model, optimizer, criterion, dataloader_train, dataloader_val)
def main(params):
# basic parameters
parser = argparse.ArgumentParser()
parser.add_argument('--image',
action='store_true',
default=False,
help='predict on image')
parser.add_argument('--video',
action='store_true',
default=False,
help='predict on video')
parser.add_argument('--checkpoint_path',
type=str,
default=None,
help='The path to the pretrained weights of model')
parser.add_argument('--context_path',
type=str,
default="resnet101",
help='The context path model you are using.')
parser.add_argument('--num_classes',
type=int,
default=12,
help='num of object classes (with void)')
parser.add_argument('--data',
default=None,
help='Path to image or video for prediction')
parser.add_argument(
'--crop_height',
type=int,
default=720,
help='Height of cropped/resized input image to network')
parser.add_argument('--crop_width',
type=int,
default=960,
help='Width of cropped/resized input image to network')
parser.add_argument('--cuda',
type=str,
default='0',
help='GPU ids used for training')
parser.add_argument('--use_gpu',
type=bool,
default=True,
help='Whether to use gpu for training')
parser.add_argument('--csv_path',
type=str,
default=None,
required=True,
help='Path to label info csv file')
parser.add_argument('--save_path',
type=str,
default=None,
required=True,
help='Path to save predict image')
args = parser.parse_args(params)
# build model
os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda
model = BiSeNet(args.num_classes, args.context_path)
if torch.cuda.is_available() and args.use_gpu:
model = torch.nn.DataParallel(model).cuda()
# load pretrained model if exists
print('load model from %s ...' % args.checkpoint_path)
model.module.load_state_dict(torch.load(args.checkpoint_path))
print('Done!')
# predict on image
if args.image:
# read image
image = cv2.imread(str(args.data), -1)
# run model
res_image = predict_on_image(model, args, image)
cv2.imwrite(args.save_path, res_image, cv2.COLOR_RGB2BGR)
# display the result
cv2.imshow("BiSeNet window", res_image)
cv2.waitKey(0)
# predict on video
if args.video:
cap = cv2.VideoCapture(args.data)
# while video source is available, loop over the frames
while cap.isOpened():
ret, frame = cap.read()
# run model
res_frame = predict_on_image(model, args, frame)
# display frames
cv2.imshow("img", res_frame)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
cap.release()
cv2.destroyAllWindows()
def main(params):
# basic parameters
parser = argparse.ArgumentParser()
parser.add_argument('--num_epochs', type=int, default=300, help='Number of epochs to train for')
parser.add_argument('--epoch_start_i', type=int, default=0, help='Start counting epochs from this number')
parser.add_argument('--checkpoint_step', type=int, default=10, help='How often to save checkpoints (epochs)')
parser.add_argument('--validation_step', type=int, default=2, help='How often to perform validation (epochs)')
parser.add_argument('--dataset', type=str, default="CamVid", help='Dataset you are using.')
parser.add_argument('--crop_height', type=int, default=720, help='Height of cropped/resized input image to network')
parser.add_argument('--crop_width', type=int, default=960, help='Width of cropped/resized input image to network')
parser.add_argument('--batch_size', type=int, default=32, help='Number of images in each batch')
parser.add_argument('--context_path', type=str, default="resnet101",
help='The context path model you are using, resnet18, resnet101.')
parser.add_argument('--learning_rate_G', type=float, default=0.01, help='learning rate for G')
parser.add_argument('--learning_rate_D', type=float, default=0.01, help='learning rate for D')#add lr_D 1e-4
parser.add_argument('--data_CamVid', type=str, default='', help='path of training data_CamVid')
parser.add_argument('--data_IDDA', type=str, default='', help='path of training data_IDDA')
parser.add_argument('--num_workers', type=int, default=4, help='num of workers')
parser.add_argument('--num_classes', type=int, default=32, help='num of object classes (with void)')
parser.add_argument('--cuda', type=str, default='0', help='GPU ids used for training')
parser.add_argument('--use_gpu', type=bool, default=True, help='whether to user gpu for training')
parser.add_argument('--pretrained_model_path', type=str, default=None, help='path to pretrained model')
parser.add_argument('--save_model_path', type=str, default=None, help='path to save model')
parser.add_argument('--optimizer_G', type=str, default='rmsprop', help='optimizer_G, support rmsprop, sgd, adam')
parser.add_argument('--optimizer_D', type=str, default='rmsprop', help='optimizer_D, support rmsprop, sgd, adam')
parser.add_argument('--loss', type=str, default='dice', help='loss function, dice or crossentropy')
parser.add_argument('--loss_G', type=str, default='dice', help='loss function, dice or crossentropy')
parser.add_argument('--lambda_adv', type=float, default=0.01, help='lambda coefficient for adversarial loss')
args = parser.parse_args(params)
# create dataset and dataloader for CamVid
CamVid_train_path = [os.path.join(args.data_CamVid, 'train'), os.path.join(args.data_CamVid, 'val')]
CamVid_train_label_path = [os.path.join(args.data_CamVid, 'train_labels'),
os.path.join(args.data_CamVid, 'val_labels')]
CamVid_test_path = os.path.join(args.data_CamVid, 'test')
CamVid_test_label_path = os.path.join(args.data_CamVid, 'test_labels')
CamVid_csv_path = os.path.join(args.data_CamVid, 'class_dict.csv')
CamVid_dataset_train = CamVid(CamVid_train_path, CamVid_train_label_path, CamVid_csv_path,
scale=(args.crop_height, args.crop_width),
loss=args.loss, mode='train')
CamVid_dataloader_train = DataLoader(
CamVid_dataset_train,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
drop_last=True
)
CamVid_dataset_val = CamVid(CamVid_test_path, CamVid_test_label_path, CamVid_csv_path,
scale=(args.crop_height, args.crop_width),
loss=args.loss, mode='test')
CamVid_dataloader_val = DataLoader(
CamVid_dataset_val,
# this has to be 1
batch_size=1,
shuffle=True,
num_workers=args.num_workers
)
# create dataset and dataloader for IDDA
IDDA_path = os.path.join(args.data_IDDA, 'rgb')
IDDA_label_path = os.path.join(args.data_IDDA, 'labels')
IDDA_info_path = os.path.join(args.data_IDDA, 'classes_info.json')
IDDA_dataset = IDDA(IDDA_path, IDDA_label_path, IDDA_info_path, CamVid_csv_path, scale=(args.crop_height, args.crop_width), loss=args.loss)
IDDA_dataloader = DataLoader(
IDDA_dataset,
batch_size=args.batch_size,
shuffle=True,
num_workers=args.num_workers,
drop_last=True
)
# build model_G
os.environ['CUDA_VISIBLE_DEVICES'] = args.cuda
model_G = BiSeNet(args.num_classes, args.context_path)
if torch.cuda.is_available() and args.use_gpu:
model_G = torch.nn.DataParallel(model_G).cuda()
#build model_D
model_D = DW_Discriminator(args.num_classes)
if torch.cuda.is_available() and args.use_gpu:
model_D = torch.nn.DataParallel(model_D).cuda()
# build optimizer G
if args.optimizer_G == 'rmsprop':
optimizer_G = torch.optim.RMSprop(model_G.parameters(), args.learning_rate_G)
elif args.optimizer_G == 'sgd':
optimizer_G = torch.optim.SGD(model_G.parameters(), args.learning_rate_G, momentum=0.9, weight_decay=1e-4)
elif args.optimizer_G == 'adam':
optimizer_G = torch.optim.Adam(model_G.parameters(), args.learning_rate_G)
else: # rmsprop
print('not supported optimizer \n')
return None
# build optimizer D
if args.optimizer_D == 'rmsprop':
optimizer_D = torch.optim.RMSprop(model_D.parameters(), args.learning_rate_D)
elif args.optimizer_D == 'sgd':
optimizer_D = torch.optim.SGD(model_D.parameters(), args.learning_rate_D, momentum=0.9, weight_decay=1e-4)
elif args.optimizer_D == 'adam':
optimizer_D = torch.optim.Adam(model_D.parameters(), args.learning_rate_D)
else: # rmsprop
print('not supported optimizer \n')
return None
curr_epoch = 0
max_miou = 0
# load pretrained model if exists
if args.pretrained_model_path is not None:
print('load model from %s ...' % args.pretrained_model_path)
state = torch.load(os.path.realpath(args.pretrained_model_path)) # upload the pretrained MODEL_G
model_G.module.load_state_dict(state['model_G_state'])
optimizer_G.load_state_dict(state['optimizer_G'])
model_D.module.load_state_dict(state['model_D_state']) # upload the pretrained MODEL_D
optimizer_D.load_state_dict(state['optimizer_D'])
curr_epoch = state["epoch"]
max_miou = state["max_miou"]
print(str(curr_epoch - 1) + " already trained")
print("start training from epoch " + str(curr_epoch))
print('Done!')
# train
train (args, model_G, model_D, optimizer_G, optimizer_D, CamVid_dataloader_train, CamVid_dataloader_val, IDDA_dataloader, curr_epoch, max_miou)
delt = rmax - img_width
rmax = img_width
rmin -= delt
if cmax > img_length:
delt = cmax - img_length
cmax = img_length
cmin -= delt
return rmin, rmax, cmin, cmax
####################################################################################################
################################### load BiSeNet parameters ########################################
####################################################################################################
print('load BiseNet')
start_time = time.time()
bise_model = BiSeNet(opt.num_classes, opt.context_path)
bise_model = bise_model.cuda()
bise_model.load_state_dict(torch.load(opt.checkpoint_path))
global bise_model
print('Done!')
print("Load time : {}".format(time.time() - start_time))
#####################################################################################################
######################## load Densefusion Netwopy4thork, 3d model #############################
#####################################################################################################
print('load densefusion network')
start_time = time.time()
estimator = PoseNet(num_points=num_points, num_obj=num_obj)
estimator.cuda()
estimator.load_state_dict(torch.load(opt.model))
estimator.eval()
def main():
# Call Python's garbage collector, and empty torch's CUDA cache. Just in case
gc.collect()
torch.cuda.empty_cache()
# Enable cuDNN in benchmark mode. For more info see:
# https://discuss.pytorch.org/t/what-does-torch-backends-cudnn-benchmark-do/5936
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
# Load Bisenet generator
generator = BiSeNet(NUM_CLASSES, CONTEXT_PATH).cuda()
# generator.load_state_dict(torch.load('./checkpoint_101_adversarial_both_augmentation_epoch_len_IDDA/37_Generator.pth'))
generator.train()
# Build discriminator
discriminator = Discriminator(NUM_CLASSES).cuda()
# discriminator.load_state_dict(torch.load('./checkpoint_101_adversarial_both_augmentation_epoch_len_IDDA/37_Discriminator.pth'))
discriminator.train()
# Load source dataset
source_dataset = IDDA(image_path=IDDA_PATH,
label_path=IDDA_LABEL_PATH,
classes_info_path=JSON_IDDA_PATH,
scale=(CROP_HEIGHT, CROP_WIDTH),
loss=LOSS,
mode='train')
source_dataloader = DataLoader(source_dataset,
batch_size=BATCH_SIZE_IDDA,
shuffle=True,
num_workers=NUM_WORKERS,
drop_last=True,
pin_memory=True)
# Load target dataset
target_dataset = CamVid(image_path=CAMVID_PATH,
label_path=CAMVID_LABEL_PATH,
csv_path=CSV_CAMVID_PATH,
scale=(CROP_HEIGHT, CROP_WIDTH),
loss=LOSS,
mode='adversarial_train')
target_dataloader = DataLoader(target_dataset,
batch_size=BATCH_SIZE_CAMVID,
shuffle=True,
num_workers=NUM_WORKERS,
drop_last=True,
pin_memory=True)
optimizer_BiSeNet = torch.optim.SGD(generator.parameters(),
lr=LEARNING_RATE_SEGMENTATION,
momentum=MOMENTUM,
weight_decay=WEIGHT_DECAY)
optimizer_discriminator = torch.optim.Adam(discriminator.parameters(),
lr=LEARNING_RATE_DISCRIMINATOR,
betas=(0.9, 0.99))
# Loss for discriminator training
# Sigmoid layer + BCELoss
bce_loss = nn.BCEWithLogitsLoss()
# Loss for segmentation loss
# Log-softmax layer + 2D Cross Entropy
cross_entropy_loss = CrossEntropy2d()
# for epoch in range(NUM_STEPS):
for epoch in range(1, 51):
source_dataloader_iter = iter(source_dataloader)
target_dataloader_iter = iter(target_dataloader)
print(f'begin epoch {epoch}')
# Initialize gradients=0 for Generator and Discriminator
optimizer_BiSeNet.zero_grad()
optimizer_discriminator.zero_grad()
# Setting losses equal to 0
l_seg_to_print_acc, l_adv_to_print_acc, l_d_to_print_acc = 0, 0, 0
# Compute learning rate for this epoch
adjust_learning_rate(optimizer_BiSeNet, LEARNING_RATE_SEGMENTATION,
epoch, NUM_STEPS, POWER)
adjust_learning_rate(optimizer_discriminator,
LEARNING_RATE_DISCRIMINATOR, epoch, NUM_STEPS,
POWER)
for i in tqdm(range(len(target_dataloader))):
optimizer_BiSeNet.zero_grad()
optimizer_discriminator.zero_grad()
l_seg_to_print, l_adv_to_print, l_d_to_print = minibatch(
source_dataloader_iter, target_dataloader_iter, generator,
discriminator, cross_entropy_loss, bce_loss, source_dataloader,
target_dataloader)
l_seg_to_print_acc += l_seg_to_print
l_adv_to_print_acc += l_adv_to_print
l_d_to_print_acc += l_d_to_print
# Run optimizers using the gradient obtained via backpropagations
optimizer_BiSeNet.step()
optimizer_discriminator.step()
# Output at each epoch
print(
f'epoch = {epoch}/{NUM_STEPS}, loss_seg = {l_seg_to_print_acc:.3f}, loss_adv = {l_adv_to_print_acc:.3f}, loss_D = {l_d_to_print_acc:.3f}'
)
# Save intermediate generator (checkpoint)
if epoch % CHECKPOINT_STEP == 0 and epoch != 0:
# If the directory does not exists create it
if not os.path.isdir(CHECKPOINT_PATH):
os.mkdir(CHECKPOINT_PATH)
# Save the parameters of the generator (segmentation network) and discriminator
generator_checkpoint_path = os.path.join(
CHECKPOINT_PATH, f"{BETA}_{epoch}_Generator.pth")
torch.save(generator.state_dict(), generator_checkpoint_path)
discriminator_checkpoint_path = os.path.join(
CHECKPOINT_PATH, f"{BETA}_{epoch}_Discriminator.pth")
torch.save(discriminator.state_dict(),
discriminator_checkpoint_path)
print(
f"saved:\n{generator_checkpoint_path}\n{discriminator_checkpoint_path}"
)
