def load_model(self, weight_file, use_half, device_id):
# import pytorch
# self._log.info("load_model does not use device, or use_half in MRCNN")
# print("Device in load model: ", device)
try:
import torch
except:
raise RuntimeError("could not load pytorch!")
# import model
try:
from modeling.model_builder import Generalized_RCNN as ubMRCNN
# from ubmrcnn import ubMRCNN
except:
raise RuntimeError(
"could not load ubMRCNN model. did you remember" +
" to setup everything?")
# if "cuda" not in device and "cpu" not in device:
# raise ValueError("invalid device name [{}]. Must str with name \
# \"cpu\" or \"cuda:X\" where X=device number")
self._log = logging.getLogger(self.idname())
# self.device = torch.device(device)
self.model = Generalized_RCNN()
#checkpoint = torch.load(weight_file, map_location=lambda storage, loc: storage)
locations = {}
for x in range(6):
locations["cuda:%d" % (x)] = "cpu"
checkpoint = torch.load(weight_file, map_location=locations)
# self.device = device_id
# if device_id==None:
# self.device = torch.device("cpu")
# else:
# self.device = torch.device("cuda:%d"%(device_id))
net_utils.load_ckpt(self.model, checkpoint['model'])
# self.model = mynn.DataParallel(self.model, cpu_keywords=['im_info', 'roidb'],
# minibatch=True, device_ids=[0],
# output_device=0) # only support single GPU
self.model = mynn.DataSingular(self.model,
cpu_keywords=['im_info', 'roidb'],
minibatch=True,
device_id=[cfg.MODEL.DEVICE])
self.model.eval()
def initialize_model_from_cfg(args, gpu_id=0):
"""Initialize a model from the global cfg. Loads test-time weights and
set to evaluation mode.
"""
model = Generalized_RCNN()
model.eval()
if args.cuda:
model.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
logger.info("loading checkpoint %s", load_name)
checkpoint = torch.load(load_name,
map_location=lambda storage, loc: storage)
net_utils.load_ckpt(model, checkpoint['model'])
# model.load_state_dict(checkpoint['model'])
if args.load_detectron:
logger.info("loading detectron weights %s", args.load_detectron)
load_detectron_weight(model, args.load_detectron)
model = mynn.DataParallel(model,
cpu_keywords=['im_info', 'roidb'],
minibatch=True)
return model
def __init__(self):
args = parse_args()
# Historical code, because we have trained on the Baidu Apoloscape dataset.
dataset = WAD_CVPR2018(args.dataset_dir)
cfg.MODEL.NUM_CLASSES = len(
dataset.eval_class) + 1 # with a background class
print('load cfg from file: {}'.format(args.cfg_file))
cfg_from_file(osp.join(this_dir, args.cfg_file))
if args.nms_soft:
cfg.TEST.SOFT_NMS.ENABLED = True
else:
cfg.TEST.NMS = args.nms
cfg.RESNETS.IMAGENET_PRETRAINED = False # Don't need to load imagenet pretrained weights
assert_and_infer_cfg()
maskRCNN = Generalized_RCNN()
maskRCNN.cuda()
if args.load_ckpt:
load_name = osp.join(this_dir, args.load_ckpt)
print("loading checkpoint %s" % (load_name))
checkpoint = torch.load(load_name,
map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
maskRCNN = mynn.DataParallel(maskRCNN,
cpu_keywords=['im_info', 'roidb'],
minibatch=True,
device_ids=[0])
maskRCNN.eval()
self.model = maskRCNN
self.dataset = dataset
def parse_dataset(path_to_dataset):
'''
Parse every image contained in 'path_to_dataset' (a path to the training
or testing set), classify each image and save in a csv file the resulting
assignment
dataset_result: dict structure returned by the function. It contains the
label of each image
'''
llist = listdir(path_to_dataset)
dataset_result = {}
maskRCNN = Generalized_RCNN()
maskRCNN.cuda()
maskRCNN = mynn.DataParallel(maskRCNN,
cpu_keywords=['im_info', 'roidb'],
minibatch=True,
device_ids=[0]) # only support single GPU
maskRCNN.eval()
for name in tqdm(llist):
dn = join(path_to_dataset, name)
timers = defaultdict(Timer)
dataset_result[dn] = localize_obj_in_image(dn, maskRCNN, timers)
return dataset_result
def main():
"""main function"""
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
args = parse_args()
print('Called with args:')
print(args)
dataset = WAD_CVPR2018(args.dataset_dir)
cfg.MODEL.NUM_CLASSES = len(
dataset.eval_class) + 1 # with a background class
print('load cfg from file: {}'.format(args.cfg_file))
cfg_from_file(args.cfg_file)
cfg.RESNETS.IMAGENET_PRETRAINED = False # Don't need to load imagenet pretrained weights
assert_and_infer_cfg()
maskRCNN = Generalized_RCNN()
maskRCNN.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
print("loading checkpoint %s" % (load_name))
checkpoint = torch.load(load_name,
map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
maskRCNN = mynn.DataParallel(maskRCNN,
cpu_keywords=['im_info', 'roidb'],
minibatch=True,
device_ids=[0]) # only support single GPU
maskRCNN.eval()
imglist = misc_utils.get_imagelist_from_dir(dataset.test_image_dir)
num_images = len(imglist)
output_dir = os.path.join(('/').join(args.load_ckpt.split('/')[:-2]),
'Images')
if not os.path.exists(output_dir):
os.makedirs(output_dir)
for i in tqdm(xrange(num_images)):
im = cv2.imread(imglist[i])
assert im is not None
timers = defaultdict(Timer)
im_name, _ = os.path.splitext(os.path.basename(imglist[i]))
args.current_im_name = im_name
cls_boxes, cls_segms, prediction_row = im_detect_all(args,
maskRCNN,
im,
dataset,
timers=timers)
thefile = open(os.path.join(output_dir, im_name + '.txt'), 'w')
for item in prediction_row:
thefile.write("%s\n" % item)
def GetRCNNModel():
""" RCNN model factory """
if cfg.CONTEXT_ROI_POOLING.ENABLED:
maskRCNN = Context_RCNN()
elif cfg.CONTRASTED_CONTEXT.ENABLED:
maskRCNN = CC_RCNN()
elif cfg.LESION.USE_3DCE:
maskRCNN = CE3D_RCNN()
elif cfg.LESION.MULTI_MODALITY:
maskRCNN = MULTI_MODALITY_RCNN()
elif cfg.LESION.MM_TEST:
maskRCNN = MM_TEST_RCNN()
else:
maskRCNN = Generalized_RCNN()
return maskRCNN
def __init__(self):
cfg_from_file("maskrcnn/configs/baselines/e2e_mask_rcnn_X-152-32x8d-FPN-IN5k_1.44x.yaml")
cfg.RESNETS.IMAGENET_PRETRAINED = False # Don't need to load imagenet pretrained weights
assert_and_infer_cfg()
maskRCNN = Generalized_RCNN()
maskRCNN.cuda()
pretrained_path = "maskrcnn/data/X-152-32x8d-FPN-IN5k.pkl"
print("loading detectron weights %s" % pretrained_path)
load_detectron_weight(maskRCNN, pretrained_path)
maskRCNN.eval() # Note this step
self.maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'],
minibatch=True, device_ids=[0]) # only support single GPU
def main():
"""main function"""
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
args = parse_args()
print('Called with args:')
print(args)
assert args.image_dir or args.images
assert bool(args.image_dir) ^ bool(args.images)
if args.dataset == "pascal_parts_heads":
dataset = datasets.get_head_dataset()
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "scuthead_a":
dataset = datasets.get_head_dataset()
cfg.MODEL.NUM_CLASSES = 2
else:
raise ValueError('Unexpected dataset name: {}'.format(args.dataset))
print('load cfg from file: {}'.format(args.cfg_file))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
assert bool(args.load_ckpt) ^ bool(args.load_detectron), \
'Exactly one of --load_ckpt and --load_detectron should be specified.'
cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights
assert_and_infer_cfg()
maskRCNN = Generalized_RCNN()
if args.cuda:
maskRCNN.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
print("loading checkpoint %s" % (load_name))
checkpoint = torch.load(load_name,
map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.load_detectron:
print("loading detectron weights %s" % args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
maskRCNN = mynn.DataParallel(maskRCNN,
cpu_keywords=['im_info', 'roidb'],
minibatch=True,
device_ids=[0]) # only support single GPU
maskRCNN.eval()
if args.image_dir:
imglist = misc_utils.get_imagelist_from_dir(args.image_dir)
else:
imglist = args.images
num_images = len(imglist)
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
for i in xrange(num_images):
im = cv2.imread(imglist[i])
assert im is not None
timers = defaultdict(Timer)
cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN,
im,
timers=timers)
im_name, _ = os.path.splitext(os.path.basename(imglist[i]))
outputfile = os.path.join(args.output_dir, im_name)
head_boxes = cls_boxes[1]
print('img :', i, ' num_heads :', len(head_boxes), ' img_path :',
imglist[i])
np.save(outputfile, head_boxes)
# -----------------------------------------------------------------------------------
# Model setup
# -----------------------------------------------------------------------------------
cfg.TEST.SCALE = 800
cfg.TEST.MAX_SIZE = 1333
cfg.MODEL.NUM_CLASSES = 2
cfg.TEST.NMS = NMS_THRESH
print('load cfg from file: {}'.format(args.cfg_file))
cfg_from_file(args.cfg_file)
assert bool(args.load_ckpt) ^ bool(args.load_detectron), \
'Exactly one of --load_ckpt and --load_detectron should be specified.'
cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False
assert_and_infer_cfg()
net = Generalized_RCNN()
if args.cuda:
net.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
print("loading checkpoint %s" % (load_name))
checkpoint = torch.load(load_name,
map_location=lambda storage, loc: storage)
net_utils.load_ckpt(net, checkpoint['model'])
if args.load_detectron:
print("loading detectron weights %s" % args.load_detectron)
load_detectron_weight(net, args.load_detectron)
def main():
"""main function"""
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
args = parse_args()
print('Called with args:')
print(args)
assert args.image_dir or args.images
assert bool(args.image_dir) ^ bool(args.images)
if args.dataset.startswith("coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = len(dataset.classes)
elif args.dataset.startswith("keypoints_coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "rsna2018":
dataset = datasets.get_rsna_dataset()
cfg.MODEL.NUM_CLASSES = 2
else:
raise ValueError('Unexpected dataset name: {}'.format(args.dataset))
print('load cfg from file: {}'.format(args.cfg_file))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
assert bool(args.load_ckpt) ^ bool(args.load_detectron), \
'Exactly one of --load_ckpt and --load_detectron should be specified.'
cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights
assert_and_infer_cfg()
maskRCNN = Generalized_RCNN()
if args.cuda:
maskRCNN.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
print("loading checkpoint %s" % (load_name))
checkpoint = torch.load(load_name,
map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.load_detectron:
print("loading detectron weights %s" % args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
maskRCNN = mynn.DataParallel(maskRCNN,
cpu_keywords=['im_info', 'roidb'],
minibatch=True,
device_ids=[0]) # only support single GPU
maskRCNN.eval()
if args.image_dir:
imglist = misc_utils.get_imagelist_from_dir(args.image_dir)
else:
imglist = args.images
num_images = len(imglist)
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
submission_dict = {}
vis_dir = os.path.join(args.output_dir, 'vis')
for i in xrange(num_images):
print('img', i)
predictionstring = ''
im = cv2.imread(imglist[i])
assert im is not None
timers = defaultdict(Timer)
cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN,
im,
timers=timers)
im_name, _ = os.path.splitext(os.path.basename(imglist[i]))
boxes = []
# vis_utils.vis_one_image(
# im[:, :, ::-1], # BGR -> RGB for visualization
# im_name,
# vis_dir,
# cls_boxes,
# cls_segms,
# cls_keyps,
# dataset=dataset,
# box_alpha=0.3,
# show_class=True,
# thresh=args.submit_threshold,
# kp_thresh=2,
# ext='png'
# )
box_list = [b for b in cls_boxes if len(b) > 0]
if len(box_list) > 0:
boxes = np.concatenate(box_list)
for box in boxes:
bbox = box[:4]
score = box[-1]
if score < args.submit_threshold:
continue
x = bbox[0]
y = bbox[1]
height = bbox[3] - bbox[1]
width = bbox[2] - bbox[0]
predictionstring += str(score) + ' ' + str(x) + ' ' + str(
y) + ' ' + str(width) + ' ' + str(height) + ' '
filename = imglist[i].split('.')[0]
filename = filename.split('/')[-1]
submission_dict[filename] = predictionstring
assert len(submission_dict) <= num_images
sub = pd.DataFrame.from_dict(submission_dict, orient='index')
sub.index.name = 'patientId'
sub.columns = ['PredictionString']
sub.to_csv(os.path.join(args.output_dir, "submission.csv"))
def test_net_on_dataset(args):
dataset = WAD_CVPR2018(args.dataset_dir)
cfg.MODEL.NUM_CLASSES = len(dataset.eval_class) + 1 # with a background class
print('load cfg from file: {}'.format(args.cfg_file))
cfg_from_file(osp.join(this_dir,args.cfg_file))
if args.nms_soft:
cfg.TEST.SOFT_NMS.ENABLED = True
else:
cfg.TEST.NMS = args.nms
cfg.RESNETS.IMAGENET_PRETRAINED = False # Don't need to load imagenet pretrained weights
assert_and_infer_cfg()
maskRCNN = Generalized_RCNN()
maskRCNN.cuda()
if args.load_ckpt:
load_name = osp.join(this_dir, args.load_ckpt)
print("loading checkpoint %s" % (load_name))
checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'], minibatch=True, device_ids=[0]) # only support single GPU
maskRCNN.eval()
if args.nms_soft:
output_dir = os.path.join(('/').join(args.load_ckpt.split('/')[:-2]), 'Images_' + str(cfg.TEST.SCALE)+'_SOFT_NMS')
elif args.nms:
output_dir = os.path.join(('/').join(args.load_ckpt.split('/')[:-2]), 'Images_' + str(cfg.TEST.SCALE)+'_NMS_%.2f'%args.nms)
else:
output_dir = os.path.join(('/').join(args.load_ckpt.split('/')[:-2]), 'Images_' + str(cfg.TEST.SCALE))
if cfg.TEST.BBOX_AUG.ENABLED:
output_dir += '_TEST_AUG'
#if args.cls_boxes_confident_threshold < 0.5:
output_dir += '_cls_boxes_confident_threshold_%.1f' % args.cls_boxes_confident_threshold
if not os.path.exists(output_dir):
os.makedirs(output_dir)
args.output_img_dir = os.path.join(output_dir, 'Image_Masks')
if not os.path.exists(args.output_img_dir):
os.makedirs(args.output_img_dir)
output_list_dir = '/home/stevenwudi/PycharmProjects/Udacity/CarND-Vehicle-Detection/output_images/car_detection_maskrcnn'
if not os.path.exists(output_list_dir):
os.makedirs(output_list_dir)
output_vis_dir = output_list_dir
imglist = ['/home/stevenwudi/PycharmProjects/Udacity/CarND-Vehicle-Detection/test_images/test3.jpg',
'/home/stevenwudi/PycharmProjects/Udacity/CarND-Vehicle-Detection/test_images/test4.jpg',
'/home/stevenwudi/PycharmProjects/Udacity/CarND-Vehicle-Detection/test_images/test5.jpg']
if not args.range:
start = 0
end = len(imglist)
else:
start = args.range[0]
end = args.range[1]
for i in tqdm(xrange(start, end)):
im = cv2.imread(imglist[i])
assert im is not None
timers = defaultdict(Timer)
im_name, _ = os.path.splitext(os.path.basename(imglist[i]))
args.current_im_name = im_name
cls_boxes, cls_segms, prediction_row = im_detect_all(args, maskRCNN, im, dataset, timers=timers)
im_name, _ = os.path.splitext(os.path.basename(imglist[i]))
print(im_name)
if args.vis:
vis_utils.vis_one_image_cvpr2018_wad(
im[:, :, ::-1], # BGR -> RGB for visualization
im_name,
output_vis_dir,
cls_boxes,
cls_segms,
None,
dataset=dataset,
box_alpha=0.3,
show_class=True,
thresh=0.99,
kp_thresh=2
)
thefile = open(os.path.join(output_list_dir, im_name + '.txt'), 'w')
for item in prediction_row:
thefile.write("%s\n" % item)
def main():
"""Main function"""
args = parse_args()
print('Called with args:')
print(args)
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
if args.cuda or cfg.NUM_GPUS > 0:
cfg.CUDA = True
else:
raise ValueError("Need Cuda device to run !")
if args.dataset == "coco2017":
cfg.TRAIN.DATASETS = ('coco_2017_train',)
cfg.MODEL.NUM_CLASSES = 81
elif args.dataset == "keypoints_coco2017":
cfg.TRAIN.DATASETS = ('keypoints_coco_2017_train',)
cfg.MODEL.NUM_CLASSES = 2
else:
raise ValueError("Unexpected args.dataset: {}".format(args.dataset))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
### Adaptively adjust some configs ###
original_batch_size = cfg.NUM_GPUS * cfg.TRAIN.IMS_PER_BATCH
original_ims_per_batch = cfg.TRAIN.IMS_PER_BATCH
original_num_gpus = cfg.NUM_GPUS
if args.batch_size is None:
args.batch_size = original_batch_size
cfg.NUM_GPUS = torch.cuda.device_count()
assert (args.batch_size % cfg.NUM_GPUS) == 0, \
'batch_size: %d, NUM_GPUS: %d' % (args.batch_size, cfg.NUM_GPUS)
cfg.TRAIN.IMS_PER_BATCH = args.batch_size // cfg.NUM_GPUS
effective_batch_size = args.iter_size * args.batch_size
print('effective_batch_size = batch_size * iter_size = %d * %d' % (args.batch_size, args.iter_size))
print('Adaptive config changes:')
print(' effective_batch_size: %d --> %d' % (original_batch_size, effective_batch_size))
print(' NUM_GPUS: %d --> %d' % (original_num_gpus, cfg.NUM_GPUS))
print(' IMS_PER_BATCH: %d --> %d' % (original_ims_per_batch, cfg.TRAIN.IMS_PER_BATCH))
### Adjust learning based on batch size change linearly
# For iter_size > 1, gradients are `accumulated`, so lr is scaled based
# on batch_size instead of effective_batch_size
old_base_lr = cfg.SOLVER.BASE_LR
cfg.SOLVER.BASE_LR *= args.batch_size / original_batch_size
print('Adjust BASE_LR linearly according to batch_size change:\n'
' BASE_LR: {} --> {}'.format(old_base_lr, cfg.SOLVER.BASE_LR))
### Adjust solver steps
step_scale = original_batch_size / effective_batch_size
old_solver_steps = cfg.SOLVER.STEPS
old_max_iter = cfg.SOLVER.MAX_ITER
cfg.SOLVER.STEPS = list(map(lambda x: int(x * step_scale + 0.5), cfg.SOLVER.STEPS))
cfg.SOLVER.MAX_ITER = int(cfg.SOLVER.MAX_ITER * step_scale + 0.5)
print('Adjust SOLVER.STEPS and SOLVER.MAX_ITER linearly based on effective_batch_size change:\n'
' SOLVER.STEPS: {} --> {}\n'
' SOLVER.MAX_ITER: {} --> {}'.format(old_solver_steps, cfg.SOLVER.STEPS,
old_max_iter, cfg.SOLVER.MAX_ITER))
# Scale FPN rpn_proposals collect size (post_nms_topN) in `collect` function
# of `collect_and_distribute_fpn_rpn_proposals.py`
#
# post_nms_topN = int(cfg[cfg_key].RPN_POST_NMS_TOP_N * cfg.FPN.RPN_COLLECT_SCALE + 0.5)
if cfg.FPN.FPN_ON and cfg.MODEL.FASTER_RCNN:
cfg.FPN.RPN_COLLECT_SCALE = cfg.TRAIN.IMS_PER_BATCH / original_ims_per_batch
print('Scale FPN rpn_proposals collect size directly propotional to the change of IMS_PER_BATCH:\n'
' cfg.FPN.RPN_COLLECT_SCALE: {}'.format(cfg.FPN.RPN_COLLECT_SCALE))
if args.num_workers is not None:
cfg.DATA_LOADER.NUM_THREADS = args.num_workers
print('Number of data loading threads: %d' % cfg.DATA_LOADER.NUM_THREADS)
### Overwrite some solver settings from command line arguments
if args.optimizer is not None:
cfg.SOLVER.TYPE = args.optimizer
if args.lr is not None:
cfg.SOLVER.BASE_LR = args.lr
if args.lr_decay_gamma is not None:
cfg.SOLVER.GAMMA = args.lr_decay_gamma
assert_and_infer_cfg()
timers = defaultdict(Timer)
### Dataset ###
timers['roidb'].tic()
roidb, ratio_list, ratio_index = combined_roidb_for_training(
cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES)
timers['roidb'].toc()
roidb_size = len(roidb)
logger.info('{:d} roidb entries'.format(roidb_size))
logger.info('Takes %.2f sec(s) to construct roidb', timers['roidb'].average_time)
# Effective training sample size for one epoch
train_size = roidb_size // args.batch_size * args.batch_size
batchSampler = BatchSampler(
sampler=MinibatchSampler(ratio_list, ratio_index),
batch_size=args.batch_size,
drop_last=True
)
dataset = RoiDataLoader(
roidb,
cfg.MODEL.NUM_CLASSES,
training=True)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_sampler=batchSampler,
num_workers=cfg.DATA_LOADER.NUM_THREADS,
collate_fn=collate_minibatch)
dataiterator = iter(dataloader)
### Model ###
maskRCNN = Generalized_RCNN()
if cfg.CUDA:
maskRCNN.cuda()
### Optimizer ###
gn_param_nameset = set()
for name, module in maskRCNN.named_modules():
if isinstance(module, nn.GroupNorm):
gn_param_nameset.add(name+'.weight')
gn_param_nameset.add(name+'.bias')
gn_params = []
gn_param_names = []
bias_params = []
bias_param_names = []
nonbias_params = []
nonbias_param_names = []
nograd_param_names = []
for key, value in maskRCNN.named_parameters():
if value.requires_grad:
if 'bias' in key:
bias_params.append(value)
bias_param_names.append(key)
elif key in gn_param_nameset:
gn_params.append(value)
gn_param_names.append(key)
else:
nonbias_params.append(value)
nonbias_param_names.append(key)
else:
nograd_param_names.append(key)
assert (gn_param_nameset - set(nograd_param_names) - set(bias_param_names)) == set(gn_param_names)
# Learning rate of 0 is a dummy value to be set properly at the start of training
params = [
{'params': nonbias_params,
'lr': 0,
'weight_decay': cfg.SOLVER.WEIGHT_DECAY},
{'params': bias_params,
'lr': 0 * (cfg.SOLVER.BIAS_DOUBLE_LR + 1),
'weight_decay': cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0},
{'params': gn_params,
'lr': 0,
'weight_decay': cfg.SOLVER.WEIGHT_DECAY_GN}
]
# names of paramerters for each paramter
param_names = [nonbias_param_names, bias_param_names, gn_param_names]
if cfg.SOLVER.TYPE == "SGD":
optimizer = torch.optim.SGD(params, momentum=cfg.SOLVER.MOMENTUM)
elif cfg.SOLVER.TYPE == "Adam":
optimizer = torch.optim.Adam(params)
### Load checkpoint
if args.load_ckpt:
load_name = args.load_ckpt
logging.info("loading checkpoint %s", load_name)
checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.resume:
args.start_step = checkpoint['step'] + 1
if 'train_size' in checkpoint: # For backward compatibility
if checkpoint['train_size'] != train_size:
print('train_size value: %d different from the one in checkpoint: %d'
% (train_size, checkpoint['train_size']))
# reorder the params in optimizer checkpoint's params_groups if needed
# misc_utils.ensure_optimizer_ckpt_params_order(param_names, checkpoint)
# There is a bug in optimizer.load_state_dict on Pytorch 0.3.1.
# However it's fixed on master.
optimizer.load_state_dict(checkpoint['optimizer'])
# misc_utils.load_optimizer_state_dict(optimizer, checkpoint['optimizer'])
del checkpoint
torch.cuda.empty_cache()
if args.load_detectron: #TODO resume for detectron weights (load sgd momentum values)
logging.info("loading Detectron weights %s", args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
lr = optimizer.param_groups[0]['lr'] # lr of non-bias parameters, for commmand line outputs.
maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'],
minibatch=True)
### Training Setups ###
args.run_name = misc_utils.get_run_name() + '_step'
output_dir = misc_utils.get_output_dir(args, args.run_name)
args.cfg_filename = os.path.basename(args.cfg_file)
if not args.no_save:
if not os.path.exists(output_dir):
os.makedirs(output_dir)
blob = {'cfg': yaml.dump(cfg), 'args': args}
with open(os.path.join(output_dir, 'config_and_args.pkl'), 'wb') as f:
pickle.dump(blob, f, pickle.HIGHEST_PROTOCOL)
if args.use_tfboard:
from tensorboardX import SummaryWriter
# Set the Tensorboard logger
tblogger = SummaryWriter(output_dir)
### Training Loop ###
maskRCNN.train()
CHECKPOINT_PERIOD = int(cfg.TRAIN.SNAPSHOT_ITERS / cfg.NUM_GPUS)
# Set index for decay steps
decay_steps_ind = None
for i in range(1, len(cfg.SOLVER.STEPS)):
if cfg.SOLVER.STEPS[i] >= args.start_step:
decay_steps_ind = i
break
if decay_steps_ind is None:
decay_steps_ind = len(cfg.SOLVER.STEPS)
training_stats = TrainingStats(
args,
args.disp_interval,
tblogger if args.use_tfboard and not args.no_save else None)
try:
logger.info('Training starts !')
step = args.start_step
for step in range(args.start_step, cfg.SOLVER.MAX_ITER):
# Warm up
if step < cfg.SOLVER.WARM_UP_ITERS:
method = cfg.SOLVER.WARM_UP_METHOD
if method == 'constant':
warmup_factor = cfg.SOLVER.WARM_UP_FACTOR
elif method == 'linear':
alpha = step / cfg.SOLVER.WARM_UP_ITERS
warmup_factor = cfg.SOLVER.WARM_UP_FACTOR * (1 - alpha) + alpha
else:
raise KeyError('Unknown SOLVER.WARM_UP_METHOD: {}'.format(method))
lr_new = cfg.SOLVER.BASE_LR * warmup_factor
net_utils.update_learning_rate(optimizer, lr, lr_new)
lr = optimizer.param_groups[0]['lr']
assert lr == lr_new
elif step == cfg.SOLVER.WARM_UP_ITERS:
net_utils.update_learning_rate(optimizer, lr, cfg.SOLVER.BASE_LR)
lr = optimizer.param_groups[0]['lr']
assert lr == cfg.SOLVER.BASE_LR
# Learning rate decay
if decay_steps_ind < len(cfg.SOLVER.STEPS) and \
step == cfg.SOLVER.STEPS[decay_steps_ind]:
logger.info('Decay the learning on step %d', step)
lr_new = lr * cfg.SOLVER.GAMMA
net_utils.update_learning_rate(optimizer, lr, lr_new)
lr = optimizer.param_groups[0]['lr']
assert lr == lr_new
decay_steps_ind += 1
training_stats.IterTic()
optimizer.zero_grad()
for inner_iter in range(args.iter_size):
try:
input_data = next(dataiterator)
except StopIteration:
dataiterator = iter(dataloader)
input_data = next(dataiterator)
for key in input_data:
if key != 'roidb': # roidb is a list of ndarrays with inconsistent length
input_data[key] = list(map(Variable, input_data[key]))
net_outputs = maskRCNN(**input_data)
training_stats.UpdateIterStats(net_outputs, inner_iter)
loss = net_outputs['total_loss']
loss.backward()
optimizer.step()
training_stats.IterToc()
training_stats.LogIterStats(step, lr)
if (step+1) % CHECKPOINT_PERIOD == 0:
save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer)
# ---- Training ends ----
# Save last checkpoint
save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer)
except (RuntimeError, KeyboardInterrupt):
del dataiterator
logger.info('Save ckpt on exception ...')
save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer)
logger.info('Save ckpt done.')
stack_trace = traceback.format_exc()
print(stack_trace)
finally:
if args.use_tfboard and not args.no_save:
tblogger.close()
def main():
"""Main function"""
args = parse_args()
print('Called with args:')
print(args)
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
if args.cuda or cfg.NUM_GPUS > 0:
cfg.CUDA = True
else:
raise ValueError("Need Cuda device to run !")
if args.dataset == "coco2017":
cfg.TRAIN.DATASETS = ('coco_2017_train', )
cfg.MODEL.NUM_CLASSES = 81
elif args.dataset == "keypoints_coco2017":
cfg.TRAIN.DATASETS = ('keypoints_coco_2017_train', )
cfg.MODEL.NUM_CLASSES = 2
else:
raise ValueError("Unexpected args.dataset: {}".format(args.dataset))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
### Adaptively adjust some configs ###
original_batch_size = cfg.NUM_GPUS * cfg.TRAIN.IMS_PER_BATCH
if args.batch_size is None:
args.batch_size = original_batch_size
cfg.NUM_GPUS = torch.cuda.device_count()
assert (args.batch_size % cfg.NUM_GPUS) == 0, \
'batch_size: %d, NUM_GPUS: %d' % (args.batch_size, cfg.NUM_GPUS)
cfg.TRAIN.IMS_PER_BATCH = args.batch_size // cfg.NUM_GPUS
print('Batch size change from {} (in config file) to {}'.format(
original_batch_size, args.batch_size))
print('NUM_GPUs: %d, TRAIN.IMS_PER_BATCH: %d' %
(cfg.NUM_GPUS, cfg.TRAIN.IMS_PER_BATCH))
if args.num_workers is not None:
cfg.DATA_LOADER.NUM_THREADS = args.num_workers
print('Number of data loading threads: %d' % cfg.DATA_LOADER.NUM_THREADS)
### Adjust learning based on batch size change linearly
old_base_lr = cfg.SOLVER.BASE_LR
cfg.SOLVER.BASE_LR *= args.batch_size / original_batch_size
print('Adjust BASE_LR linearly according to batch size change: {} --> {}'.
format(old_base_lr, cfg.SOLVER.BASE_LR))
### Overwrite some solver settings from command line arguments
if args.optimizer is not None:
cfg.SOLVER.TYPE = args.optimizer
if args.lr is not None:
cfg.SOLVER.BASE_LR = args.lr
if args.lr_decay_gamma is not None:
cfg.SOLVER.GAMMA = args.lr_decay_gamma
assert_and_infer_cfg()
timers = defaultdict(Timer)
### Dataset ###
timers['roidb'].tic()
roidb, ratio_list, ratio_index = combined_roidb_for_training(
cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES)
timers['roidb'].toc()
roidb_size = len(roidb)
logger.info('{:d} roidb entries'.format(roidb_size))
logger.info('Takes %.2f sec(s) to construct roidb',
timers['roidb'].average_time)
# Effective training sample size for one epoch
train_size = roidb_size // args.batch_size * args.batch_size
sampler = MinibatchSampler(ratio_list, ratio_index)
dataset = RoiDataLoader(roidb, cfg.MODEL.NUM_CLASSES, training=True)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=args.batch_size,
drop_last=True,
sampler=sampler,
num_workers=cfg.DATA_LOADER.NUM_THREADS,
collate_fn=collate_minibatch)
dataiterator = iter(dataloader)
### Model ###
maskRCNN = Generalized_RCNN()
if cfg.CUDA:
maskRCNN.cuda()
### Optimizer ###
bias_params = []
nonbias_params = []
for key, value in dict(maskRCNN.named_parameters()).items():
if value.requires_grad:
if 'bias' in key:
bias_params.append(value)
else:
nonbias_params.append(value)
# Learning rate of 0 is a dummy value to be set properly at the start of training
params = [{
'params': nonbias_params,
'lr': 0,
'weight_decay': cfg.SOLVER.WEIGHT_DECAY
}, {
'params':
bias_params,
'lr':
0 * (cfg.SOLVER.BIAS_DOUBLE_LR + 1),
'weight_decay':
cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0
}]
if cfg.SOLVER.TYPE == "SGD":
optimizer = torch.optim.SGD(params, momentum=cfg.SOLVER.MOMENTUM)
elif cfg.SOLVER.TYPE == "Adam":
optimizer = torch.optim.Adam(params)
### Load checkpoint
if args.load_ckpt:
load_name = args.load_ckpt
logging.info("loading checkpoint %s", load_name)
checkpoint = torch.load(load_name,
map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.resume:
args.start_step = checkpoint['step'] + 1
assert checkpoint['train_size'] == train_size
# There is a bug in optimizer.load_state_dict on Pytorch 0.3.1.
# However it's fixed on master.
# optimizer.load_state_dict(checkpoint['optimizer'])
misc_utils.load_optimizer_state_dict(optimizer,
checkpoint['optimizer'])
del checkpoint
torch.cuda.empty_cache()
if args.load_detectron: #TODO resume for detectron weights (load sgd momentum values)
logging.info("loading Detectron weights %s", args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
lr = optimizer.param_groups[0][
'lr'] # lr of non-bias parameters, for commmand line outputs.
maskRCNN = mynn.DataParallel(maskRCNN,
cpu_keywords=['im_info', 'roidb'],
minibatch=True)
### Training Setups ###
run_name = misc_utils.get_run_name()
output_dir = misc_utils.get_output_dir(args, run_name)
if not args.no_save:
if not os.path.exists(output_dir):
os.makedirs(output_dir)
blob = {'cfg': yaml.dump(cfg), 'args': args}
with open(os.path.join(output_dir, 'config_and_args.pkl'), 'wb') as f:
pickle.dump(blob, f, pickle.HIGHEST_PROTOCOL)
if args.use_tfboard:
from tensorboardX import SummaryWriter
# Set the Tensorboard logger
tblogger = SummaryWriter(output_dir)
### Training Loop ###
maskRCNN.train()
CHECKPOINT_PERIOD = int(cfg.TRAIN.SNAPSHOT_ITERS / cfg.NUM_GPUS)
# Set index for decay steps
decay_steps_ind = None
for i in range(1, len(cfg.SOLVER.STEPS)):
if cfg.SOLVER.STEPS[i] > args.start_step:
decay_steps_ind = i
if decay_steps_ind is None:
decay_steps_ind = len(cfg.SOLVER.STEPS)
logger.info('Training starts !')
loss_avg = 0
try:
timers['train_loop'].tic()
for step in range(args.start_step, cfg.SOLVER.MAX_ITER):
# Warm up
if step < cfg.SOLVER.WARM_UP_ITERS:
method = cfg.SOLVER.WARM_UP_METHOD
if method == 'constant':
warmup_factor = cfg.SOLVER.WARM_UP_FACTOR
elif method == 'linear':
alpha = step / cfg.SOLVER.WARM_UP_ITERS
warmup_factor = cfg.SOLVER.WARM_UP_FACTOR * (1 -
alpha) + alpha
else:
raise KeyError(
'Unknown SOLVER.WARM_UP_METHOD: {}'.format(method))
lr_new = cfg.SOLVER.BASE_LR * warmup_factor
net_utils.update_learning_rate(optimizer, lr, lr_new)
lr = lr_new
elif step == cfg.SOLVER.WARM_UP_ITERS:
lr = cfg.SOLVER.BASE_LR
# Learning rate decay
if decay_steps_ind < len(cfg.SOLVER.STEPS) and \
step == cfg.SOLVER.STEPS[decay_steps_ind]:
logger.info('Decay the learning on step %d', step)
lr_new = lr * cfg.SOLVER.GAMMA
net_utils.update_learning_rate(optimizer, lr, lr_new)
lr = lr_new
decay_steps_ind += 1
try:
input_data = next(dataiterator)
except StopIteration:
dataiterator = iter(dataloader)
input_data = next(dataiterator)
for key in input_data:
if key != 'roidb': # roidb is a list of ndarrays with inconsistent length
input_data[key] = list(map(Variable, input_data[key]))
outputs = maskRCNN(**input_data)
rois_label = outputs['rois_label']
cls_score = outputs['cls_score']
bbox_pred = outputs['bbox_pred']
loss_rpn_cls = outputs['loss_rpn_cls'].mean()
loss_rpn_bbox = outputs['loss_rpn_bbox'].mean()
loss_rcnn_cls = outputs['loss_rcnn_cls'].mean()
loss_rcnn_bbox = outputs['loss_rcnn_bbox'].mean()
loss = loss_rpn_cls + loss_rpn_bbox + loss_rcnn_cls + loss_rcnn_bbox
if cfg.MODEL.MASK_ON:
loss_rcnn_mask = outputs['loss_rcnn_mask'].mean()
loss += loss_rcnn_mask
if cfg.MODEL.KEYPOINTS_ON:
loss_rcnn_keypoints = outputs['loss_rcnn_keypoints'].mean()
loss += loss_rcnn_keypoints
loss_avg += loss.data.cpu().numpy()[0]
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (step + 1) % CHECKPOINT_PERIOD == 0:
save_ckpt(output_dir, args, step, train_size, maskRCNN,
optimizer)
if ((step % args.disp_interval == 0 and
(step - args.start_step >= args.disp_interval))
or step == cfg.SOLVER.MAX_ITER - 1):
diff = timers['train_loop'].toc(average=False)
loss_avg /= args.disp_interval
loss_rpn_cls = loss_rpn_cls.data[0]
loss_rpn_bbox = loss_rpn_bbox.data[0]
loss_rcnn_cls = loss_rcnn_cls.data[0]
loss_rcnn_bbox = loss_rcnn_bbox.data[0]
fg_cnt = torch.sum(rois_label.data.ne(0))
bg_cnt = rois_label.data.numel() - fg_cnt
print("[ %s ][ step %d ]" % (run_name, step))
print("\t\tloss: %.4f, lr: %.2e" % (loss_avg, lr))
print("\t\tfg/bg=(%d/%d), time cost: %f" %
(fg_cnt, bg_cnt, diff))
print(
"\t\trpn_cls: %.4f, rpn_bbox: %.4f, rcnn_cls: %.4f, rcnn_bbox %.4f"
% (loss_rpn_cls, loss_rpn_bbox, loss_rcnn_cls,
loss_rcnn_bbox))
print_prefix = "\t\t"
if cfg.MODEL.MASK_ON:
loss_rcnn_mask = loss_rcnn_mask.data[0]
print("%srcnn_mask %.4f" % (print_prefix, loss_rcnn_mask))
print_prefix = ", "
if cfg.MODEL.KEYPOINTS_ON:
loss_rcnn_keypoints = loss_rcnn_keypoints.data[0]
print("%srcnn_keypoints %.4f" %
(print_prefix, loss_rcnn_keypoints))
if args.use_tfboard and not args.no_save:
info = {
'lr': lr,
'loss': loss_avg,
'loss_rpn_cls': loss_rpn_cls,
'loss_rpn_box': loss_rpn_bbox,
'loss_rcnn_cls': loss_rcnn_cls,
'loss_rcnn_box': loss_rcnn_bbox,
}
if cfg.MODEL.MASK_ON:
info['loss_rcnn_mask'] = loss_rcnn_mask
if cfg.MODEL.KEYPOINTS_ON:
info['loss_rcnn_keypoints'] = loss_rcnn_keypoints
for tag, value in info.items():
tblogger.add_scalar(tag, value, step)
loss_avg = 0
timers['train_loop'].tic()
# Save last checkpoint
save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer)
except (RuntimeError, KeyboardInterrupt) as e:
print('Save on exception:', e)
save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer)
stack_trace = traceback.format_exc()
print(stack_trace)
finally:
# ---- Training ends ----
if args.use_tfboard and not args.no_save:
tblogger.close()
def main():
"""main function"""
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
args = parse_args()
print('Called with args:')
print(args)
assert args.image_dir or args.images
assert bool(args.image_dir) ^ bool(args.images)
if args.dataset.startswith("coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = len(dataset.classes)
elif args.dataset.startswith("keypoints_coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = 2
else:
raise ValueError('Unexpected dataset name: {}'.format(args.dataset))
print('load cfg from file: {}'.format(args.cfg_file))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
assert bool(args.load_ckpt) ^ bool(args.load_detectron), \
'Exactly one of --load_ckpt and --load_detectron should be specified.'
cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights
assert_and_infer_cfg()
maskRCNN = Generalized_RCNN()
if args.cuda:
maskRCNN.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
print("loading checkpoint %s" % (load_name))
checkpoint = torch.load(load_name,
map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.load_detectron:
print("loading detectron weights %s" % args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
maskRCNN = mynn.DataParallel(maskRCNN,
cpu_keywords=['im_info', 'roidb'],
minibatch=True,
device_ids=[0]) # only support single GPU
maskRCNN.eval()
if args.image_dir:
imglist = misc_utils.get_imagelist_from_dir(args.image_dir)
else:
imglist = args.images
num_images = len(imglist)
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
for i in xrange(num_images):
print('img', i)
im = cv2.imread(imglist[i])
assert im is not None
timers = defaultdict(Timer)
cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN,
im,
timers=timers)
im_name, _ = os.path.splitext(os.path.basename(imglist[i]))
vis_utils.vis_one_image(
im[:, :, ::-1], # BGR -> RGB for visualization
im_name,
args.output_dir,
cls_boxes,
cls_segms,
cls_keyps,
dataset=dataset,
box_alpha=0.3,
show_class=True,
thresh=0.7,
kp_thresh=2)
if args.merge_pdfs and num_images > 1:
merge_out_path = '{}/results.pdf'.format(args.output_dir)
if os.path.exists(merge_out_path):
os.remove(merge_out_path)
command = "pdfunite {}/*.pdf {}".format(args.output_dir,
merge_out_path)
subprocess.call(command, shell=True)
def main():
"""Main function"""
args = parse_args()
print('Called with args:')
print(args)
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
if args.cuda or cfg.NUM_GPUS > 0:
cfg.CUDA = True
else:
raise ValueError("Need Cuda device to run !")
if args.dataset == "coco2017":
cfg.TRAIN.DATASETS = ('coco_2017_train', )
cfg.MODEL.NUM_CLASSES = 81
elif args.dataset == "keypoints_coco2017":
cfg.TRAIN.DATASETS = ('keypoints_coco_2017_train', )
cfg.MODEL.NUM_CLASSES = 2
# ADE20k as a detection dataset
elif args.dataset == "ade_train":
cfg.TRAIN.DATASETS = ('ade_train', )
cfg.MODEL.NUM_CLASSES = 446
# Noisy CS6+WIDER datasets
elif args.dataset == 'cs6_noise020+WIDER':
cfg.TRAIN.DATASETS = ('cs6_noise020', 'wider_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'cs6_noise050+WIDER':
cfg.TRAIN.DATASETS = ('cs6_noise050', 'wider_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'cs6_noise080+WIDER':
cfg.TRAIN.DATASETS = ('cs6_noise080', 'wider_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'cs6_noise085+WIDER':
cfg.TRAIN.DATASETS = ('cs6_noise085', 'wider_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'cs6_noise090+WIDER':
cfg.TRAIN.DATASETS = ('cs6_noise090', 'wider_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'cs6_noise095+WIDER':
cfg.TRAIN.DATASETS = ('cs6_noise095', 'wider_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'cs6_noise100+WIDER':
cfg.TRAIN.DATASETS = ('cs6_noise100', 'wider_train')
cfg.MODEL.NUM_CLASSES = 2
# Just Noisy CS6 datasets
elif args.dataset == 'cs6_noise020':
cfg.TRAIN.DATASETS = ('cs6_noise020', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'cs6_noise030':
cfg.TRAIN.DATASETS = ('cs6_noise030', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'cs6_noise040':
cfg.TRAIN.DATASETS = ('cs6_noise040', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'cs6_noise050':
cfg.TRAIN.DATASETS = ('cs6_noise050', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'cs6_noise060':
cfg.TRAIN.DATASETS = ('cs6_noise060', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'cs6_noise070':
cfg.TRAIN.DATASETS = ('cs6_noise070', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'cs6_noise080':
cfg.TRAIN.DATASETS = ('cs6_noise080', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'cs6_noise085':
cfg.TRAIN.DATASETS = ('cs6_noise085', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'cs6_noise090':
cfg.TRAIN.DATASETS = ('cs6_noise090', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'cs6_noise095':
cfg.TRAIN.DATASETS = ('cs6_noise095', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'cs6_noise100':
cfg.TRAIN.DATASETS = ('cs6_noise100', )
cfg.MODEL.NUM_CLASSES = 2
# Cityscapes 7 classes
elif args.dataset == "cityscapes":
cfg.TRAIN.DATASETS = ('cityscapes_train', )
cfg.MODEL.NUM_CLASSES = 8
# BDD 7 classes
elif args.dataset == "bdd_any_any_any":
cfg.TRAIN.DATASETS = ('bdd_any_any_any_train', )
cfg.MODEL.NUM_CLASSES = 8
elif args.dataset == "bdd_any_any_daytime":
cfg.TRAIN.DATASETS = ('bdd_any_any_daytime_train', )
cfg.MODEL.NUM_CLASSES = 8
elif args.dataset == "bdd_clear_any_daytime":
cfg.TRAIN.DATASETS = ('bdd_clear_any_daytime_train', )
cfg.MODEL.NUM_CLASSES = 8
# Cistyscapes Pedestrian sets
elif args.dataset == "cityscapes_peds":
cfg.TRAIN.DATASETS = ('cityscapes_peds_train', )
cfg.MODEL.NUM_CLASSES = 2
# Cityscapes Car sets
elif args.dataset == "cityscapes_cars_HPlen3+kitti_car_train":
cfg.TRAIN.DATASETS = ('cityscapes_cars_HPlen3', 'kitti_car_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "cityscapes_cars_HPlen5+kitti_car_train":
cfg.TRAIN.DATASETS = ('cityscapes_cars_HPlen5', 'kitti_car_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "cityscapes_car_train+kitti_car_train":
cfg.TRAIN.DATASETS = ('cityscapes_car_train', 'kitti_car_train')
cfg.MODEL.NUM_CLASSES = 2
# KITTI Car set
elif args.dataset == "kitti_car_train":
cfg.TRAIN.DATASETS = ('kitti_car_train', )
cfg.MODEL.NUM_CLASSES = 2
# BDD pedestrians sets
elif args.dataset == "bdd_peds":
cfg.TRAIN.DATASETS = ('bdd_peds_train',
) # bdd peds: clear_any_daytime
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "bdd_peds_full":
cfg.TRAIN.DATASETS = ('bdd_peds_full_train', ) # bdd peds: any_any_any
cfg.MODEL.NUM_CLASSES = 2
# Pedestrians with constraints
elif args.dataset == "bdd_peds_not_clear_any_daytime":
cfg.TRAIN.DATASETS = ('bdd_peds_not_clear_any_daytime_train', )
cfg.MODEL.NUM_CLASSES = 2
# Ashish's 20k samples videos
elif args.dataset == "bdd_peds_not_clear_any_daytime_20k":
cfg.TRAIN.DATASETS = ('bdd_peds_not_clear_any_daytime_20k_train', )
cfg.MODEL.NUM_CLASSES = 2
# Source domain + Target domain detections
elif args.dataset == "bdd_peds+DETS_20k":
cfg.TRAIN.DATASETS = ('bdd_peds_dets_20k_target_domain',
'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
# Source domain + Target domain detections -- same 18k images as HP18k
elif args.dataset == "bdd_peds+DETS18k":
cfg.TRAIN.DATASETS = ('bdd_peds_dets18k_target_domain',
'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
# Only Dets
elif args.dataset == "DETS20k":
cfg.TRAIN.DATASETS = ('bdd_peds_dets_20k_target_domain', )
cfg.MODEL.NUM_CLASSES = 2
# Only Dets18k - same images as HP18k
elif args.dataset == 'DETS18k':
cfg.TRAIN.DATASETS = ('bdd_peds_dets18k_target_domain', )
cfg.MODEL.NUM_CLASSES = 2
# Only HP
elif args.dataset == 'HP':
cfg.TRAIN.DATASETS = ('bdd_peds_HP_target_domain', )
cfg.MODEL.NUM_CLASSES = 2
# Only HP 18k videos
elif args.dataset == 'HP18k':
cfg.TRAIN.DATASETS = ('bdd_peds_HP18k_target_domain', )
cfg.MODEL.NUM_CLASSES = 2
# Source domain + Target domain HP
elif args.dataset == 'bdd_peds+HP':
cfg.TRAIN.DATASETS = ('bdd_peds_train', 'bdd_peds_HP_target_domain')
cfg.MODEL.NUM_CLASSES = 2
# Source domain + Target domain HP 18k videos
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'bdd_peds+HP18k':
cfg.TRAIN.DATASETS = ('bdd_peds_HP18k_target_domain', 'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
#### Source domain + Target domain with different conf threshold theta ####
elif args.dataset == 'bdd_peds+HP18k_thresh-050':
cfg.TRAIN.DATASETS = ('bdd_HP18k_thresh-050', 'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'bdd_peds+HP18k_thresh-060':
cfg.TRAIN.DATASETS = ('bdd_HP18k_thresh-060', 'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'bdd_peds+HP18k_thresh-070':
cfg.TRAIN.DATASETS = ('bdd_HP18k_thresh-070', 'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'bdd_peds+HP18k_thresh-090':
cfg.TRAIN.DATASETS = ('bdd_HP18k_thresh-090', 'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
##############################
#### Data distillation on BDD -- for rebuttal
elif args.dataset == 'bdd_peds+bdd_data_dist_small':
cfg.TRAIN.DATASETS = ('bdd_data_dist_small', 'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'bdd_peds+bdd_data_dist_mid':
cfg.TRAIN.DATASETS = ('bdd_data_dist_mid', 'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'bdd_peds+bdd_data_dist':
cfg.TRAIN.DATASETS = ('bdd_data_dist', 'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
##############################
#### Source domain + **Labeled** Target domain with varying number of images
elif args.dataset == 'bdd_peds+labeled_100':
cfg.TRAIN.DATASETS = ('bdd_peds_not_clear_any_daytime_train_100',
'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'bdd_peds+labeled_075':
cfg.TRAIN.DATASETS = ('bdd_peds_not_clear_any_daytime_train_075',
'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'bdd_peds+labeled_050':
cfg.TRAIN.DATASETS = ('bdd_peds_not_clear_any_daytime_train_050',
'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'bdd_peds+labeled_025':
cfg.TRAIN.DATASETS = ('bdd_peds_not_clear_any_daytime_train_025',
'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'bdd_peds+labeled_010':
cfg.TRAIN.DATASETS = ('bdd_peds_not_clear_any_daytime_train_010',
'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'bdd_peds+labeled_005':
cfg.TRAIN.DATASETS = ('bdd_peds_not_clear_any_daytime_train_005',
'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'bdd_peds+labeled_001':
cfg.TRAIN.DATASETS = ('bdd_peds_not_clear_any_daytime_train_001',
'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
##############################
# Source domain + Target domain HP tracker bboxes only
elif args.dataset == 'bdd_peds+HP18k_track_only':
cfg.TRAIN.DATASETS = ('bdd_HP18k_track_only', 'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
##### subsets of bdd_HP18k with different constraints
# Source domain + HP tracker images at NIGHT
elif args.dataset == 'bdd_peds+HP18k_any_any_night':
cfg.TRAIN.DATASETS = ('bdd_HP18k_any_any_night', 'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'bdd_peds+HP18k_rainy_any_daytime':
cfg.TRAIN.DATASETS = ('bdd_HP18k_rainy_any_daytime', 'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'bdd_peds+HP18k_rainy_any_night':
cfg.TRAIN.DATASETS = ('bdd_HP18k_rainy_any_night', 'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'bdd_peds+HP18k_overcast,rainy_any_daytime':
cfg.TRAIN.DATASETS = ('bdd_HP18k_overcast,rainy_any_daytime',
'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'bdd_peds+HP18k_overcast,rainy_any_night':
cfg.TRAIN.DATASETS = ('bdd_HP18k_overcast,rainy_any_night',
'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'bdd_peds+HP18k_overcast,rainy,snowy_any_daytime':
cfg.TRAIN.DATASETS = ('bdd_HP18k_overcast,rainy,snowy_any_daytime',
'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
############# end of bdd constraned subsets #####################
# Source domain + Target domain HP18k -- after histogram matching
elif args.dataset == 'bdd_peds+HP18k_remap_hist':
cfg.TRAIN.DATASETS = ('bdd_peds_HP18k_target_domain_remap_hist',
'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'bdd_peds+HP18k_remap_cityscape_hist':
cfg.TRAIN.DATASETS = (
'bdd_peds_HP18k_target_domain_remap_cityscape_hist',
'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
# Source domain + Target domain HP18k -- after histogram matching
elif args.dataset == 'bdd_peds+HP18k_remap_random':
cfg.TRAIN.DATASETS = ('bdd_peds_HP18k_target_domain_remap_random',
'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
# Source+Noisy Target domain -- prevent domain adv from using HP roi info
elif args.dataset == 'bdd_peds+bdd_HP18k_noisy_100k':
cfg.TRAIN.DATASETS = ('bdd_HP18k_noisy_100k', 'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'bdd_peds+bdd_HP18k_noisy_080':
cfg.TRAIN.DATASETS = ('bdd_HP18k_noisy_080', 'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'bdd_peds+bdd_HP18k_noisy_060':
cfg.TRAIN.DATASETS = ('bdd_HP18k_noisy_060', 'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == 'bdd_peds+bdd_HP18k_noisy_070':
cfg.TRAIN.DATASETS = ('bdd_HP18k_noisy_070', 'bdd_peds_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "wider_train":
cfg.TRAIN.DATASETS = ('wider_train', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "cs6-subset":
cfg.TRAIN.DATASETS = ('cs6-subset', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "cs6-subset-score":
cfg.TRAIN.DATASETS = ('cs6-subset-score', )
elif args.dataset == "cs6-subset-gt":
cfg.TRAIN.DATASETS = ('cs6-subset-gt', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "cs6-3013-gt":
cfg.TRAIN.DATASETS = ('cs6-3013-gt',
) # DEBUG: overfit on one video annots
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "cs6-subset-gt+WIDER":
cfg.TRAIN.DATASETS = ('cs6-subset-gt', 'wider_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "cs6-subset+WIDER":
cfg.TRAIN.DATASETS = ('cs6-subset', 'wider_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "cs6-train-gt":
cfg.TRAIN.DATASETS = ('cs6-train-gt', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "cs6-train-gt-noisy-0.3":
cfg.TRAIN.DATASETS = ('cs6-train-gt-noisy-0.3', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "cs6-train-gt-noisy-0.5":
cfg.TRAIN.DATASETS = ('cs6-train-gt-noisy-0.5', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "cs6-train-det-score":
cfg.TRAIN.DATASETS = ('cs6-train-det-score', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "cs6-train-det-score-0.5":
cfg.TRAIN.DATASETS = ('cs6-train-det-score-0.5', )
elif args.dataset == "cs6-train-det":
cfg.TRAIN.DATASETS = ('cs6-train-det', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "cs6-train-det-0.5":
cfg.TRAIN.DATASETS = ('cs6-train-det-0.5', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "cs6-train-hp":
cfg.TRAIN.DATASETS = ('cs6-train-hp', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "cs6-train-easy-gt":
cfg.TRAIN.DATASETS = ('cs6-train-easy-gt', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "cs6-train-easy-gt-sub":
cfg.TRAIN.DATASETS = ('cs6-train-easy-gt-sub', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "cs6-train-easy-hp":
cfg.TRAIN.DATASETS = ('cs6-train-easy-hp', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "cs6-train-easy-det":
cfg.TRAIN.DATASETS = ('cs6-train-easy-det', )
cfg.MODEL.NUM_CLASSES = 2
# Joint training with CS6 and WIDER
elif args.dataset == "cs6-train-easy-gt-sub+WIDER":
cfg.TRAIN.DATASETS = ('cs6-train-easy-gt-sub', 'wider_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "cs6-train-gt+WIDER":
cfg.TRAIN.DATASETS = ('cs6-train-gt', 'wider_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "cs6-train-hp+WIDER":
cfg.TRAIN.DATASETS = ('cs6-train-hp', 'wider_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "cs6-train-dummy+WIDER":
cfg.TRAIN.DATASETS = ('cs6-train-dummy', 'wider_train')
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "cs6-train-det+WIDER":
cfg.TRAIN.DATASETS = ('cs6-train-det', 'wider_train')
cfg.MODEL.NUM_CLASSES = 2
# Dets dataset created by removing tracker results from the HP json
elif args.dataset == "cs6_train_det_from_hp+WIDER":
cfg.TRAIN.DATASETS = ('cs6_train_det_from_hp', 'wider_train')
cfg.MODEL.NUM_CLASSES = 2
# Dataset created by removing det results from the HP json -- HP tracker only
elif args.dataset == "cs6_train_hp_tracker_only+WIDER":
cfg.TRAIN.DATASETS = ('cs6_train_hp_tracker_only', 'wider_train')
cfg.MODEL.NUM_CLASSES = 2
# HP dataset with noisy labels: used to prevent DA from getting any info from HP
elif args.dataset == "cs6_train_hp_noisy_100+WIDER":
cfg.TRAIN.DATASETS = ('cs6_train_hp_noisy_100', 'wider_train')
cfg.MODEL.NUM_CLASSES = 2
else:
raise ValueError("Unexpected args.dataset: {}".format(args.dataset))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
### Adaptively adjust some configs ###
original_batch_size = cfg.NUM_GPUS * cfg.TRAIN.IMS_PER_BATCH
original_ims_per_batch = cfg.TRAIN.IMS_PER_BATCH
original_num_gpus = cfg.NUM_GPUS
if args.batch_size is None:
args.batch_size = original_batch_size
cfg.NUM_GPUS = torch.cuda.device_count()
assert (args.batch_size % cfg.NUM_GPUS) == 0, \
'batch_size: %d, NUM_GPUS: %d' % (args.batch_size, cfg.NUM_GPUS)
cfg.TRAIN.IMS_PER_BATCH = args.batch_size // cfg.NUM_GPUS
effective_batch_size = args.iter_size * args.batch_size
print('effective_batch_size = batch_size * iter_size = %d * %d' %
(args.batch_size, args.iter_size))
print('Adaptive config changes:')
print(' effective_batch_size: %d --> %d' %
(original_batch_size, effective_batch_size))
print(' NUM_GPUS: %d --> %d' %
(original_num_gpus, cfg.NUM_GPUS))
print(' IMS_PER_BATCH: %d --> %d' %
(original_ims_per_batch, cfg.TRAIN.IMS_PER_BATCH))
### Adjust learning based on batch size change linearly
# For iter_size > 1, gradients are `accumulated`, so lr is scaled based
# on batch_size instead of effective_batch_size
old_base_lr = cfg.SOLVER.BASE_LR
cfg.SOLVER.BASE_LR *= args.batch_size / original_batch_size
print('Adjust BASE_LR linearly according to batch_size change:\n'
' BASE_LR: {} --> {}'.format(old_base_lr, cfg.SOLVER.BASE_LR))
### Adjust solver steps
step_scale = original_batch_size / effective_batch_size
old_solver_steps = cfg.SOLVER.STEPS
old_max_iter = cfg.SOLVER.MAX_ITER
cfg.SOLVER.STEPS = list(
map(lambda x: int(x * step_scale + 0.5), cfg.SOLVER.STEPS))
cfg.SOLVER.MAX_ITER = int(cfg.SOLVER.MAX_ITER * step_scale + 0.5)
print(
'Adjust SOLVER.STEPS and SOLVER.MAX_ITER linearly based on effective_batch_size change:\n'
' SOLVER.STEPS: {} --> {}\n'
' SOLVER.MAX_ITER: {} --> {}'.format(old_solver_steps,
cfg.SOLVER.STEPS, old_max_iter,
cfg.SOLVER.MAX_ITER))
# Scale FPN rpn_proposals collect size (post_nms_topN) in `collect` function
# of `collect_and_distribute_fpn_rpn_proposals.py`
#
# post_nms_topN = int(cfg[cfg_key].RPN_POST_NMS_TOP_N * cfg.FPN.RPN_COLLECT_SCALE + 0.5)
if cfg.FPN.FPN_ON and cfg.MODEL.FASTER_RCNN:
cfg.FPN.RPN_COLLECT_SCALE = cfg.TRAIN.IMS_PER_BATCH / original_ims_per_batch
print(
'Scale FPN rpn_proposals collect size directly propotional to the change of IMS_PER_BATCH:\n'
' cfg.FPN.RPN_COLLECT_SCALE: {}'.format(
cfg.FPN.RPN_COLLECT_SCALE))
if args.num_workers is not None:
cfg.DATA_LOADER.NUM_THREADS = args.num_workers
print('Number of data loading threads: %d' % cfg.DATA_LOADER.NUM_THREADS)
### Overwrite some solver settings from command line arguments
if args.optimizer is not None:
cfg.SOLVER.TYPE = args.optimizer
if args.lr is not None:
cfg.SOLVER.BASE_LR = args.lr
if args.lr_decay_gamma is not None:
cfg.SOLVER.GAMMA = args.lr_decay_gamma
assert_and_infer_cfg()
timers = defaultdict(Timer)
"""def _init_fn(worker_id):
random.seed(999)
np.random.seed(999)
torch.cuda.manual_seed(999)
torch.cuda.manual_seed_all(999)
torch.manual_seed(999)
torch.backends.cudnn.deterministic = True
"""
### Dataset ###
timers['roidb'].tic()
roidb, ratio_list, ratio_index = combined_roidb_for_training(
cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES)
timers['roidb'].toc()
roidb_size = len(roidb)
logger.info('{:d} roidb entries'.format(roidb_size))
logger.info('Takes %.2f sec(s) to construct roidb',
timers['roidb'].average_time)
if cfg.TRAIN.JOINT_TRAINING:
if len(cfg.TRAIN.DATASETS) == 2:
print('Joint training on two datasets')
else:
raise NotImplementedError
joint_training_roidb = []
for i, dataset_name in enumerate(cfg.TRAIN.DATASETS):
# ROIDB construction
timers['roidb'].tic()
roidb, ratio_list, ratio_index = combined_roidb_for_training(
(dataset_name), cfg.TRAIN.PROPOSAL_FILES)
timers['roidb'].toc()
roidb_size = len(roidb)
logger.info('{:d} roidb entries'.format(roidb_size))
logger.info('Takes %.2f sec(s) to construct roidb',
timers['roidb'].average_time)
if i == 0:
roidb_size = len(roidb)
batchSampler = BatchSampler(sampler=MinibatchSampler(
ratio_list, ratio_index),
batch_size=args.batch_size,
drop_last=True)
dataset = RoiDataLoader(roidb,
cfg.MODEL.NUM_CLASSES,
training=True)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_sampler=batchSampler,
num_workers=cfg.DATA_LOADER.NUM_THREADS,
collate_fn=collate_minibatch)
#worker_init_fn=_init_fn)
# decrease num-threads when using two dataloaders
dataiterator = iter(dataloader)
joint_training_roidb.append({
'dataloader': dataloader,
'dataiterator': dataiterator,
'dataset_name': dataset_name
})
else:
roidb, ratio_list, ratio_index = combined_roidb_for_training(
cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES)
timers['roidb'].toc()
roidb_size = len(roidb)
logger.info('{:d} roidb entries'.format(roidb_size))
logger.info('Takes %.2f sec(s) to construct roidb',
timers['roidb'].average_time)
batchSampler = BatchSampler(sampler=MinibatchSampler(
ratio_list, ratio_index),
batch_size=args.batch_size,
drop_last=True)
dataset = RoiDataLoader(roidb, cfg.MODEL.NUM_CLASSES, training=True)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_sampler=batchSampler,
num_workers=cfg.DATA_LOADER.NUM_THREADS,
collate_fn=collate_minibatch)
#worker_init_fn=init_fn)
dataiterator = iter(dataloader)
# Effective training sample size for one epoch
train_size = roidb_size // args.batch_size * args.batch_size
if cfg.TRAIN.JOINT_SELECTIVE_FG:
orig_fg_batch_ratio = cfg.TRAIN.FG_FRACTION
### Model ###
maskRCNN = Generalized_RCNN()
if cfg.CUDA:
maskRCNN.cuda()
### Optimizer ###
gn_param_nameset = set()
for name, module in maskRCNN.named_modules():
if isinstance(module, nn.GroupNorm):
gn_param_nameset.add(name + '.weight')
gn_param_nameset.add(name + '.bias')
gn_params = []
gn_param_names = []
bias_params = []
bias_param_names = []
nonbias_params = []
nonbias_param_names = []
nograd_param_names = []
for key, value in dict(maskRCNN.named_parameters()).items():
if value.requires_grad:
if 'bias' in key:
bias_params.append(value)
bias_param_names.append(key)
elif key in gn_param_nameset:
gn_params.append(value)
gn_param_names.append(key)
else:
nonbias_params.append(value)
nonbias_param_names.append(key)
else:
nograd_param_names.append(key)
assert (gn_param_nameset - set(nograd_param_names) -
set(bias_param_names)) == set(gn_param_names)
# Learning rate of 0 is a dummy value to be set properly at the start of training
params = [{
'params': nonbias_params,
'lr': 0,
'weight_decay': cfg.SOLVER.WEIGHT_DECAY
}, {
'params':
bias_params,
'lr':
0 * (cfg.SOLVER.BIAS_DOUBLE_LR + 1),
'weight_decay':
cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0
}, {
'params': gn_params,
'lr': 0,
'weight_decay': cfg.SOLVER.WEIGHT_DECAY_GN
}]
# names of paramerters for each paramter
param_names = [nonbias_param_names, bias_param_names, gn_param_names]
if cfg.SOLVER.TYPE == "SGD":
optimizer = torch.optim.SGD(params, momentum=cfg.SOLVER.MOMENTUM)
elif cfg.SOLVER.TYPE == "Adam":
optimizer = torch.optim.Adam(params)
### Load checkpoint
if args.load_ckpt:
load_name = args.load_ckpt
logging.info("loading checkpoint %s", load_name)
checkpoint = torch.load(load_name,
map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.resume:
args.start_step = checkpoint['step'] + 1
if 'train_size' in checkpoint: # For backward compatibility
if checkpoint['train_size'] != train_size:
print(
'train_size value: %d different from the one in checkpoint: %d'
% (train_size, checkpoint['train_size']))
# reorder the params in optimizer checkpoint's params_groups if needed
# misc_utils.ensure_optimizer_ckpt_params_order(param_names, checkpoint)
# There is a bug in optimizer.load_state_dict on Pytorch 0.3.1.
# However it's fixed on master.
# optimizer.load_state_dict(checkpoint['optimizer'])
misc_utils.load_optimizer_state_dict(optimizer,
checkpoint['optimizer'])
del checkpoint
torch.cuda.empty_cache()
if args.load_detectron: #TODO resume for detectron weights (load sgd momentum values)
logging.info("loading Detectron weights %s", args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
lr = optimizer.param_groups[0][
'lr'] # lr of non-bias parameters, for commmand line outputs.
maskRCNN = mynn.DataParallel(maskRCNN,
cpu_keywords=['im_info', 'roidb'],
minibatch=True)
### Training Setups ###
args.run_name = misc_utils.get_run_name() + '_' + str(
cfg.TRAIN.DATASETS) + '_step'
output_dir = misc_utils.get_output_dir(args, args.run_name)
args.cfg_filename = os.path.basename(args.cfg_file)
if not args.no_save:
if not os.path.exists(output_dir):
os.makedirs(output_dir)
blob = {'cfg': yaml.dump(cfg), 'args': args}
with open(os.path.join(output_dir, 'config_and_args.pkl'), 'wb') as f:
pickle.dump(blob, f, pickle.HIGHEST_PROTOCOL)
if args.use_tfboard:
from tensorboardX import SummaryWriter
# Set the Tensorboard logger
tblogger = SummaryWriter(output_dir)
### Training Loop ###
maskRCNN.train()
CHECKPOINT_PERIOD = int(cfg.TRAIN.SNAPSHOT_ITERS / cfg.NUM_GPUS)
# Set index for decay steps
decay_steps_ind = None
for i in range(1, len(cfg.SOLVER.STEPS)):
if cfg.SOLVER.STEPS[i] >= args.start_step:
decay_steps_ind = i
break
if decay_steps_ind is None:
decay_steps_ind = len(cfg.SOLVER.STEPS)
training_stats = TrainingStats(
args, args.disp_interval,
tblogger if args.use_tfboard and not args.no_save else None)
try:
logger.info('Training starts !')
step = args.start_step
for step in range(args.start_step, cfg.SOLVER.MAX_ITER):
"""
random.seed(cfg.RNG_SEED)
np.random.seed(cfg.RNG_SEED)
torch.cuda.manual_seed(cfg.RNG_SEED)
torch.cuda.manual_seed_all(cfg.RNG_SEED)
torch.manual_seed(cfg.RNG_SEED)
torch.backends.cudnn.deterministic = True
"""
# Warm up
if step < cfg.SOLVER.WARM_UP_ITERS:
method = cfg.SOLVER.WARM_UP_METHOD
if method == 'constant':
warmup_factor = cfg.SOLVER.WARM_UP_FACTOR
elif method == 'linear':
alpha = step / cfg.SOLVER.WARM_UP_ITERS
warmup_factor = cfg.SOLVER.WARM_UP_FACTOR * (1 -
alpha) + alpha
else:
raise KeyError(
'Unknown SOLVER.WARM_UP_METHOD: {}'.format(method))
lr_new = cfg.SOLVER.BASE_LR * warmup_factor
net_utils.update_learning_rate(optimizer, lr, lr_new)
lr = optimizer.param_groups[0]['lr']
assert lr == lr_new
elif step == cfg.SOLVER.WARM_UP_ITERS:
net_utils.update_learning_rate(optimizer, lr,
cfg.SOLVER.BASE_LR)
lr = optimizer.param_groups[0]['lr']
assert lr == cfg.SOLVER.BASE_LR
# Learning rate decay
if decay_steps_ind < len(cfg.SOLVER.STEPS) and \
step == cfg.SOLVER.STEPS[decay_steps_ind]:
logger.info('Decay the learning on step %d', step)
lr_new = lr * cfg.SOLVER.GAMMA
net_utils.update_learning_rate(optimizer, lr, lr_new)
lr = optimizer.param_groups[0]['lr']
assert lr == lr_new
decay_steps_ind += 1
training_stats.IterTic()
optimizer.zero_grad()
for inner_iter in range(args.iter_size):
# use a iter counter for optional alternating batches
if args.iter_size == 1:
iter_counter = step
else:
iter_counter = inner_iter
if cfg.TRAIN.JOINT_TRAINING:
# alternate batches between dataset[0] and dataset[1]
if iter_counter % 2 == 0:
if True: #DEBUG:
print('Dataset: %s' %
joint_training_roidb[0]['dataset_name'])
dataloader = joint_training_roidb[0]['dataloader']
dataiterator = joint_training_roidb[0]['dataiterator']
# NOTE: if available FG samples cannot fill minibatch
# then batchsize will be smaller than cfg.TRAIN.BATCH_SIZE_PER_IM.
else:
if True: #DEBUG:
print('Dataset: %s' %
joint_training_roidb[1]['dataset_name'])
dataloader = joint_training_roidb[1]['dataloader']
dataiterator = joint_training_roidb[1]['dataiterator']
try:
input_data = next(dataiterator)
except StopIteration:
# end of epoch for dataloader
dataiterator = iter(dataloader)
input_data = next(dataiterator)
if cfg.TRAIN.JOINT_TRAINING:
if iter_counter % 2 == 0:
joint_training_roidb[0][
'dataiterator'] = dataiterator
else:
joint_training_roidb[1][
'dataiterator'] = dataiterator
for key in input_data:
if key != 'roidb': # roidb is a list of ndarrays with inconsistent length
input_data[key] = list(map(Variable, input_data[key]))
net_outputs = maskRCNN(**input_data)
training_stats.UpdateIterStats(net_outputs, inner_iter)
loss = net_outputs['total_loss']
# [p.data.get_device() for p in maskRCNN.parameters()]
# [(name, p.data.get_device()) for name, p in maskRCNN.named_parameters()]
loss.backward()
optimizer.step()
training_stats.IterToc()
training_stats.LogIterStats(step, lr)
if (step + 1) % CHECKPOINT_PERIOD == 0:
save_ckpt(output_dir, args, step, train_size, maskRCNN,
optimizer)
# ---- Training ends ----
# Save last checkpoint
save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer)
except (RuntimeError, KeyboardInterrupt):
del dataiterator
logger.info('Save ckpt on exception ...')
save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer)
logger.info('Save ckpt done.')
stack_trace = traceback.format_exc()
print(stack_trace)
finally:
if args.use_tfboard and not args.no_save:
tblogger.close()
p_tensor.copy_(torch.Tensor(src_blobs[d_name]))
def resnet_weights_name_pattern():
pattern = re.compile(r"conv1_w|conv1_gn_[sb]|res_conv1_.+|res\d_\d_.+")
return pattern
if __name__ == '__main__':
"""Testing"""
from pprint import pprint
import sys
sys.path.insert(0, '..')
from modeling.model_builder import Generalized_RCNN
from core.config import cfg, cfg_from_file
cfg.MODEL.NUM_CLASSES = 81
cfg_from_file('../../cfgs/res50_mask.yml')
net = Generalized_RCNN()
# pprint(list(net.state_dict().keys()), width=1)
mapping, orphans = net.detectron_weight_mapping
state_dict = net.state_dict()
for k in mapping.keys():
assert k in state_dict, '%s' % k
rest = set(state_dict.keys()) - set(mapping.keys())
assert len(rest) == 0
def main():
"""main function"""
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
args = parse_args()
print('Called with args:')
print(args)
assert args.image_dir or args.images
assert bool(args.image_dir) ^ bool(args.images)
if args.dataset.startswith("coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = len(dataset.classes)
elif args.dataset.startswith("keypoints_coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset.startswith("gangjin"):
dataset = datasets.get_gangjin_dataset()
cfg.MODEL.NUM_CLASSES = len(dataset.classes)
else:
raise ValueError('Unexpected dataset name: {}'.format(args.dataset))
print('load cfg from file: {}'.format(args.cfg_file))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
assert bool(args.load_ckpt) ^ bool(args.load_detectron), \
'Exactly one of --load_ckpt and --load_detectron should be specified.'
cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights
assert_and_infer_cfg()
maskRCNN = Generalized_RCNN()
if args.cuda:
maskRCNN.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
print("loading checkpoint %s" % (load_name))
checkpoint = torch.load(load_name,
map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.load_detectron:
print("loading detectron weights %s" % args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
maskRCNN = mynn.DataParallel(maskRCNN,
cpu_keywords=['im_info', 'roidb'],
minibatch=True,
device_ids=[0]) # only support single GPU
maskRCNN.eval()
if args.image_dir:
imglist = misc_utils.get_imagelist_from_dir(args.image_dir)
else:
imglist = args.images
num_images = len(imglist)
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
img_ids = []
rects = []
for i in range(num_images):
print('img', i)
im = cv2.imread(imglist[i])
assert im is not None
timers = defaultdict(Timer)
cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN,
im,
timers=timers)
boxes, segms, keypoints, classes = vis_utils.convert_from_cls_format(
cls_boxes, cls_segms, cls_keyps)
if boxes is not None:
for j in range(len(boxes)):
# print(boxes[j][-1])
if float(boxes[j][-1]) < 0.99: # 阀值
continue
xmin = float(boxes[j, 0])
xmax = float(boxes[j, 2])
ymin = float(boxes[j, 1])
ymax = float(boxes[j, 3])
img_ids.append(os.path.basename(imglist[i]))
rects.append(
str(xmin) + " " + str(ymin) + " " + str(xmax) + " " +
str(ymax))
# im_name, _ = os.path.splitext(os.path.basename(imglist[i]))
# vis_utils.vis_one_image(
# im[:, :, ::-1], # BGR -> RGB for visualization
# im_name,
# args.output_dir,
# cls_boxes,
# cls_segms,
# cls_keyps,
# dataset=dataset,
# box_alpha=0.3,
# show_class=False,
# thresh=0.99,
# kp_thresh=2,
# ext="jpg"
# )
result_dict = {"ID": img_ids, "rects": rects}
import pandas as pd
result = pd.DataFrame.from_dict(result_dict)
result.to_csv('submit/submit1.csv', header=None, index=False)
def main():
"""main function"""
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
args = parse_args()
print('Called with args:')
print(args)
assert args.image_dir or args.images
assert bool(args.image_dir) ^ bool(args.images)
prefix_path = args.output_dir + '_results'
if os.path.exists(prefix_path):
shutil.rmtree(prefix_path)
os.mkdir(prefix_path)
else:
os.mkdir(prefix_path)
if args.dataset.startswith("coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = len(dataset.classes)
elif args.dataset.startswith("keypoints_coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = 2
else:
raise ValueError('Unexpected dataset name: {}'.format(args.dataset))
print('load cfg from file: {}'.format(args.cfg_file))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
assert bool(args.load_ckpt) ^ bool(args.load_detectron), \
'Exactly one of --load_ckpt and --load_detectron should be specified.'
cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights
assert_and_infer_cfg()
maskRCNN = Generalized_RCNN()
if args.cuda:
maskRCNN.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
print("loading checkpoint %s" % (load_name))
checkpoint = torch.load(load_name,
map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.load_detectron:
print("loading detectron weights %s" % args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
maskRCNN = mynn.DataParallel(maskRCNN,
cpu_keywords=['im_info', 'roidb'],
minibatch=True,
device_ids=[0]) # only support single GPU
maskRCNN.eval()
if args.image_dir:
imglist = misc_utils.get_imagelist_from_dir(args.image_dir)
else:
imglist = args.images
num_images = len(imglist)
for i in tqdm(range(num_images)):
im = cv2.imread(imglist[i])
assert im is not None
timers = defaultdict(Timer)
cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN,
im,
timers=timers)
im_name, _ = os.path.splitext(os.path.basename(imglist[i]))
boxes, _, _, classes = convert_from_cls_format(cls_boxes, cls_segms,
cls_keyps)
if classes == []:
continue
voc_boxes = np.zeros_like(boxes)
voc_boxes[:, 0:1] = boxes[:, 4:5]
voc_boxes[:, 1:3] = boxes[:, 0:2] + 1
voc_boxes[:, 3:5] = boxes[:, 2:4] + 1
for instance_idx, cls_idx in enumerate(classes):
cls_name = dataset.classes[cls_idx]
if cls_name == 'motorcycle':
cls_name = 'motorbike'
f = open(os.path.join(prefix_path, cls_name + ".txt"), "a+")
f.write("%s " % im_name)
for item in voc_boxes[instance_idx]:
f.write("%f " % item)
f.write("\n")
f.close()
def main():
"""main function"""
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
args = parse_args()
print('Called with args:')
print(args)
assert args.image_dir or args.images
assert bool(args.image_dir) ^ bool(args.images)
prefix_path = args.output_dir
os.makedirs(prefix_path, exist_ok=True)
if args.dataset.startswith("coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = len(dataset.classes)
elif args.dataset.startswith("keypoints_coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = 2
else:
raise ValueError('Unexpected dataset name: {}'.format(args.dataset))
print('load cfg from file: {}'.format(args.cfg_file))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
assert bool(args.load_ckpt) ^ bool(args.load_detectron), \
'Exactly one of --load_ckpt and --load_detectron should be specified.'
cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights
assert_and_infer_cfg()
maskRCNN = Generalized_RCNN()
if args.cuda:
maskRCNN.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
print("loading checkpoint %s" % (load_name))
checkpoint = torch.load(load_name,
map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.load_detectron:
print("loading detectron weights %s" % args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
maskRCNN = mynn.DataParallel(maskRCNN,
cpu_keywords=['im_info', 'roidb'],
minibatch=True,
device_ids=[0]) # only support single GPU
maskRCNN.eval()
if args.image_dir:
imglist = misc_utils.get_imagelist_from_dir(args.image_dir)
else:
imglist = args.images
num_images = len(imglist)
writen_results = []
# validate
demo_im = cv2.imread(imglist[0])
print(np.shape(demo_im))
h, w, _ = np.shape(demo_im)
#print(h)
#print(args.height)
assert h == args.height
assert w == args.width
h_scale = 720 / args.height
w_scale = 1280 / args.width
for i in tqdm(range(num_images)):
im = cv2.imread(imglist[i])
assert im is not None
timers = defaultdict(Timer)
cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN,
im,
timers=timers)
im_name, _ = os.path.splitext(os.path.basename(imglist[i]))
# boxs = [[x1, y1, x2, y2, cls], ...]
boxes, _, _, classes = convert_from_cls_format(cls_boxes, cls_segms,
cls_keyps)
if boxes is None:
continue
# scale
boxes[:, 0] = boxes[:, 0] * w_scale
boxes[:, 2] = boxes[:, 2] * w_scale
boxes[:, 1] = boxes[:, 1] * h_scale
boxes[:, 3] = boxes[:, 3] * h_scale
if classes == []:
continue
for instance_idx, cls_idx in enumerate(classes):
cls_name = dataset.classes[cls_idx]
if cls_name == 'motorcycle':
cls_name = 'motor'
elif cls_name == 'stop sign':
cls_name = 'traffic sign'
elif cls_name == 'bicycle':
cls_name = 'bike'
if cls_name not in bdd_category:
continue
writen_results.append({
"name": imglist[i].split('/')[-1],
"timestamp": 1000,
"category": cls_name,
"bbox": boxes[instance_idx, :4],
"score": boxes[instance_idx, -1]
})
with open(os.path.join(prefix_path, args.name + '.json'),
'w') as outputfile:
json.dump(writen_results, outputfile, cls=MyEncoder)
from __future__ import absolute_import
def main():
"""main function"""
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
args = parse_args()
print('Called with args:')
print(args)
assert args.image_dir or args.images
assert bool(args.image_dir) ^ bool(args.images)
if args.dataset.startswith("coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = len(dataset.classes)
elif args.dataset == "miotcd":
dataset = datasets.get_miotcd_dataset()
cfg.MODEL.NUM_CLASSES = 12
elif args.dataset.startswith("keypoints_coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset.startswith("bogota"):
dataset = datasets.get_bogota_dataset()
cfg.MODEL.NUM_CLASSES = 12
else:
raise ValueError('Unexpected dataset name: {}'.format(args.dataset))
print('load cfg from file: {}'.format(args.cfg_file))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
assert bool(args.load_ckpt) ^ bool(args.load_detectron), \
'Exactly one of --load_ckpt and --load_detectron should be specified.'
cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights
assert_and_infer_cfg()
maskRCNN = Generalized_RCNN()
if args.cuda:
maskRCNN.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
print("loading checkpoint %s" % (load_name))
checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.load_detectron:
print("loading detectron weights %s" % args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'],
minibatch=True, device_ids=[0]) # only support single GPU
maskRCNN.eval()
if args.image_dir:
imglist = misc_utils.get_imagelist_from_dir(args.image_dir)
else:
imglist = args.images
num_images = len(imglist)
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
dataset_result = {}
for i in xrange(num_images):
print('img', i)
im = cv2.imread(imglist[i])
assert im is not None
try:
timers = defaultdict(Timer)
cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN, im, timers=timers)
boxes_, segme_ , keyps_ ,clasies= convert_from_cls_format(cls_boxes,cls_segms,cls_keyps)
im_name, _ = os.path.splitext(os.path.basename(imglist[i]))
dataset_result[im_name] = localize_obj_in_image(im_name,boxes_,clasies)
except(e):
import pdb
pdb.set_trace()
np.save('dictionary_answer.npy', dataset_result)
tmp = args.image_dir.split('/')
txt = str(tmp[-2:])+'_'+str(tmp[-1:])+'.csv'
save_localization_result(dataset_result,txt)
def main():
"""Main function"""
args = parse_args()
print('Called with args:')
print(args)
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
if args.cuda or cfg.NUM_GPUS > 0:
cfg.CUDA = True
else:
raise ValueError("Need Cuda device to run !")
if args.dataset == "coco2017":
cfg.TRAIN.DATASETS = ('coco_2017_train',)
elif args.dataset == "keypoints_coco2017":
cfg.TRAIN.DATASETS = ('keypoints_coco_2017_train',)
else:
raise ValueError("Unexpected args.dataset: {}".format(args.dataset))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
### Adaptively adjust some configs ###
original_batch_size = cfg.NUM_GPUS * cfg.TRAIN.IMS_PER_BATCH
if args.batch_size is None:
args.batch_size = original_batch_size
cfg.NUM_GPUS = torch.cuda.device_count()
assert (args.batch_size % cfg.NUM_GPUS) == 0, \
'batch_size: %d, NUM_GPUS: %d' % (args.batch_size, cfg.NUM_GPUS)
cfg.TRAIN.IMS_PER_BATCH = args.batch_size // cfg.NUM_GPUS
print('Batch size change from {} (in config file) to {}'.format(
original_batch_size, args.batch_size))
print('NUM_GPUs: %d, TRAIN.IMS_PER_BATCH: %d' % (cfg.NUM_GPUS, cfg.TRAIN.IMS_PER_BATCH))
if args.num_workers is not None:
cfg.DATA_LOADER.NUM_THREADS = args.num_workers
print('Number of data loading threads: %d' % cfg.DATA_LOADER.NUM_THREADS)
### Adjust learning based on batch size change linearly
old_base_lr = cfg.SOLVER.BASE_LR
cfg.SOLVER.BASE_LR *= args.batch_size / original_batch_size
print('Adjust BASE_LR linearly according to batch size change: {} --> {}'.format(
old_base_lr, cfg.SOLVER.BASE_LR))
### Overwrite some solver settings from command line arguments
if args.optimizer is not None:
cfg.SOLVER.TYPE = args.optimizer
if args.lr is not None:
cfg.SOLVER.BASE_LR = args.lr
if args.lr_decay_gamma is not None:
cfg.SOLVER.GAMMA = args.lr_decay_gamma
args.mGPUs = (cfg.NUM_GPUS > 1)
timers = defaultdict(Timer)
### Dataset ###
timers['roidb'].tic()
roidb, ratio_list, ratio_index = combined_roidb_for_training(
cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES)
timers['roidb'].toc()
train_size = len(roidb)
logger.info('{:d} roidb entries'.format(train_size))
logger.info('Takes %.2f sec(s) to construct roidb', timers['roidb'].average_time)
sampler = MinibatchSampler(ratio_list, ratio_index)
dataset = RoiDataLoader(
roidb,
cfg.MODEL.NUM_CLASSES,
training=True)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=args.batch_size,
sampler=sampler,
num_workers=cfg.DATA_LOADER.NUM_THREADS,
collate_fn=collate_minibatch)
assert_and_infer_cfg()
### Model ###
maskRCNN = Generalized_RCNN()
if cfg.CUDA:
maskRCNN.cuda()
### Optimizer ###
bias_params = []
nonbias_params = []
for key, value in dict(maskRCNN.named_parameters()).items():
if value.requires_grad:
if 'bias' in key:
bias_params.append(value)
else:
nonbias_params.append(value)
params = [
{'params': nonbias_params,
'lr': cfg.SOLVER.BASE_LR,
'weight_decay': cfg.SOLVER.WEIGHT_DECAY},
{'params': bias_params,
'lr': cfg.SOLVER.BASE_LR * (cfg.SOLVER.BIAS_DOUBLE_LR + 1),
'weight_decay': cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0}
]
if cfg.SOLVER.TYPE == "SGD":
optimizer = torch.optim.SGD(params, momentum=cfg.SOLVER.MOMENTUM)
elif cfg.SOLVER.TYPE == "Adam":
optimizer = torch.optim.Adam(params)
### Load checkpoint
if args.load_ckpt:
load_name = args.load_ckpt
logging.info("loading checkpoint %s", load_name)
checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.resume:
assert checkpoint['iters_per_epoch'] == train_size // args.batch_size, \
"iters_per_epoch should match for resume"
# There is a bug in optimizer.load_state_dict on Pytorch 0.3.1.
# However it's fixed on master.
# optimizer.load_state_dict(checkpoint['optimizer'])
misc_utils.load_optimizer_state_dict(optimizer, checkpoint['optimizer'])
if checkpoint['step'] == (checkpoint['iters_per_epoch'] - 1):
# Resume from end of an epoch
args.start_epoch = checkpoint['epoch'] + 1
args.start_iter = 0
else:
# Resume from the middle of an epoch.
# NOTE: dataloader is not synced with previous state
args.start_epoch = checkpoint['epoch']
args.start_iter = checkpoint['step'] + 1
del checkpoint
torch.cuda.empty_cache()
if args.load_detectron: #TODO resume for detectron weights (load sgd momentum values)
logging.info("loading Detectron weights %s", args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
lr = optimizer.param_groups[0]['lr'] # lr of non-bias parameters, for commmand line outputs.
maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'],
minibatch=True)
### Training Setups ###
run_name = misc_utils.get_run_name()
output_dir = misc_utils.get_output_dir(args, run_name)
if not args.no_save:
if not os.path.exists(output_dir):
os.makedirs(output_dir)
blob = {'cfg': yaml.dump(cfg), 'args': args}
with open(os.path.join(output_dir, 'config_and_args.pkl'), 'wb') as f:
pickle.dump(blob, f, pickle.HIGHEST_PROTOCOL)
if args.use_tfboard:
from tensorboardX import SummaryWriter
# Set the Tensorboard logger
tblogger = SummaryWriter(output_dir)
### Training Loop ###
maskRCNN.train()
iters_per_epoch = int(train_size / args.batch_size) # drop last
ckpt_interval_per_epoch = iters_per_epoch // args.ckpt_num_per_epoch
step = 0
try:
logger.info('Training starts !')
for epoch in range(args.start_epoch, args.start_epoch + args.num_epochs):
# ---- Start of epoch ----
loss_avg = 0
timers['train_loop'].tic()
# adjust learning rate
if args.lr_decay_epochs and epoch == args.lr_decay_epochs[0] and args.start_iter == 0:
args.lr_decay_epochs.pop(0)
net_utils.decay_learning_rate(optimizer, lr, cfg.SOLVER.GAMMA)
lr *= cfg.SOLVER.GAMMA
for step, input_data in zip(range(args.start_iter, iters_per_epoch), dataloader):
for key in input_data:
if key != 'roidb': # roidb is a list of ndarrays with inconsistent length
input_data[key] = list(map(Variable, input_data[key]))
outputs = maskRCNN(**input_data)
rois_label = outputs['rois_label']
cls_score = outputs['cls_score']
bbox_pred = outputs['bbox_pred']
loss_rpn_cls = outputs['loss_rpn_cls'].mean()
loss_rpn_bbox = outputs['loss_rpn_bbox'].mean()
loss_rcnn_cls = outputs['loss_rcnn_cls'].mean()
loss_rcnn_bbox = outputs['loss_rcnn_bbox'].mean()
loss = loss_rpn_cls + loss_rpn_bbox + loss_rcnn_cls + loss_rcnn_bbox
if cfg.MODEL.MASK_ON:
loss_rcnn_mask = outputs['loss_rcnn_mask'].mean()
loss += loss_rcnn_mask
if cfg.MODEL.KEYPOINTS_ON:
loss_rcnn_keypoints = outputs['loss_rcnn_keypoints'].mean()
loss += loss_rcnn_keypoints
loss_avg += loss.data.cpu().numpy()[0]
optimizer.zero_grad()
loss.backward()
optimizer.step()
if (step+1) % ckpt_interval_per_epoch == 0:
net_utils.save_ckpt(output_dir, args, epoch, step, maskRCNN, optimizer, iters_per_epoch)
if (step+1) % args.disp_interval == 0:
if (step + 1 - args.start_iter) >= args.disp_interval: # for the case of resume
diff = timers['train_loop'].toc(average=False)
loss_avg /= args.disp_interval
loss_rpn_cls = loss_rpn_cls.data[0]
loss_rpn_bbox = loss_rpn_bbox.data[0]
loss_rcnn_cls = loss_rcnn_cls.data[0]
loss_rcnn_bbox = loss_rcnn_bbox.data[0]
fg_cnt = torch.sum(rois_label.data.ne(0))
bg_cnt = rois_label.data.numel() - fg_cnt
print("[%s][epoch %2d][iter %4d / %4d]"
% (run_name, epoch, step, iters_per_epoch))
print("\t\tloss: %.4f, lr: %.2e" % (loss_avg, lr))
print("\t\tfg/bg=(%d/%d), time cost: %f" % (fg_cnt, bg_cnt, diff))
print("\t\trpn_cls: %.4f, rpn_bbox: %.4f, rcnn_cls: %.4f, rcnn_bbox %.4f"
% (loss_rpn_cls, loss_rpn_bbox, loss_rcnn_cls, loss_rcnn_bbox))
print_prefix = "\t\t"
if cfg.MODEL.MASK_ON:
loss_rcnn_mask = loss_rcnn_mask.data[0]
print("%srcnn_mask %.4f" % (print_prefix, loss_rcnn_mask))
print_prefix = ", "
if cfg.MODEL.KEYPOINTS_ON:
loss_rcnn_keypoints = loss_rcnn_keypoints.data[0]
print("%srcnn_keypoints %.4f" % (print_prefix, loss_rcnn_keypoints))
if args.use_tfboard:
info = {
'loss': loss_avg,
'loss_rpn_cls': loss_rpn_cls,
'loss_rpn_box': loss_rpn_bbox,
'loss_rcnn_cls': loss_rcnn_cls,
'loss_rcnn_box': loss_rcnn_bbox,
}
if cfg.MODEL.MASK_ON:
info['loss_rcnn_mask'] = loss_rcnn_mask
if cfg.MODEL.KEYPOINTS_ON:
info['loss_rcnn_keypoints'] = loss_rcnn_keypoints
for tag, value in info.items():
tblogger.add_scalar(tag, value, iters_per_epoch * epoch + step)
loss_avg = 0
timers['train_loop'].tic()
# ---- End of epoch ----
# save checkpoint
net_utils.save_ckpt(output_dir, args, epoch, step, maskRCNN, optimizer, iters_per_epoch)
# reset timer
timers['train_loop'].reset()
# reset starting iter number after first epoch
args.start_iter = 0
except (RuntimeError, KeyboardInterrupt) as e:
print('Save on exception:', e)
net_utils.save_ckpt(output_dir, args, epoch, step, maskRCNN, optimizer, iters_per_epoch)
stack_trace = traceback.format_exc()
print(stack_trace)
finally:
# ---- Training ends ----
if args.use_tfboard:
tblogger.close()
def main():
"""Main function"""
args = parse_args()
print('Called with args:')
print(args)
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
if args.cuda or cfg.NUM_GPUS > 0:
cfg.CUDA = True
else:
raise ValueError("Need Cuda device to run !")
if args.dataset == "coco2017":
cfg.TRAIN.DATASETS = ('coco_2017_train',)
cfg.MODEL.NUM_CLASSES = 81
elif args.dataset == "keypoints_coco2017":
cfg.TRAIN.DATASETS = ('keypoints_coco_2017_train',)
cfg.MODEL.NUM_CLASSES = 2
else:
raise ValueError("Unexpected args.dataset: {}".format(args.dataset))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
### Adaptively adjust some configs ###
original_batch_size = cfg.NUM_GPUS * cfg.TRAIN.IMS_PER_BATCH
original_ims_per_batch = cfg.TRAIN.IMS_PER_BATCH
original_num_gpus = cfg.NUM_GPUS
if args.batch_size is None:
args.batch_size = original_batch_size
cfg.NUM_GPUS = torch.cuda.device_count()
assert (args.batch_size % cfg.NUM_GPUS) == 0, \
'batch_size: %d, NUM_GPUS: %d' % (args.batch_size, cfg.NUM_GPUS)
cfg.TRAIN.IMS_PER_BATCH = args.batch_size // cfg.NUM_GPUS
effective_batch_size = args.iter_size * args.batch_size
print('effective_batch_size = batch_size * iter_size = %d * %d' % (args.batch_size, args.iter_size))
print('Adaptive config changes:')
print(' effective_batch_size: %d --> %d' % (original_batch_size, effective_batch_size))
print(' NUM_GPUS: %d --> %d' % (original_num_gpus, cfg.NUM_GPUS))
print(' IMS_PER_BATCH: %d --> %d' % (original_ims_per_batch, cfg.TRAIN.IMS_PER_BATCH))
### Adjust learning based on batch size change linearly
# For iter_size > 1, gradients are `accumulated`, so lr is scaled based
# on batch_size instead of effective_batch_size
old_base_lr = cfg.SOLVER.BASE_LR
cfg.SOLVER.BASE_LR *= args.batch_size / original_batch_size
print('Adjust BASE_LR linearly according to batch_size change:\n'
' BASE_LR: {} --> {}'.format(old_base_lr, cfg.SOLVER.BASE_LR))
### Adjust solver steps
step_scale = original_batch_size / effective_batch_size
old_solver_steps = cfg.SOLVER.STEPS
old_max_iter = cfg.SOLVER.MAX_ITER
cfg.SOLVER.STEPS = list(map(lambda x: int(x * step_scale + 0.5), cfg.SOLVER.STEPS))
cfg.SOLVER.MAX_ITER = int(cfg.SOLVER.MAX_ITER * step_scale + 0.5)
print('Adjust SOLVER.STEPS and SOLVER.MAX_ITER linearly based on effective_batch_size change:\n'
' SOLVER.STEPS: {} --> {}\n'
' SOLVER.MAX_ITER: {} --> {}'.format(old_solver_steps, cfg.SOLVER.STEPS,
old_max_iter, cfg.SOLVER.MAX_ITER))
# Scale FPN rpn_proposals collect size (post_nms_topN) in `collect` function
# of `collect_and_distribute_fpn_rpn_proposals.py`
#
# post_nms_topN = int(cfg[cfg_key].RPN_POST_NMS_TOP_N * cfg.FPN.RPN_COLLECT_SCALE + 0.5)
if cfg.FPN.FPN_ON and cfg.MODEL.FASTER_RCNN:
cfg.FPN.RPN_COLLECT_SCALE = cfg.TRAIN.IMS_PER_BATCH / original_ims_per_batch
print('Scale FPN rpn_proposals collect size directly propotional to the change of IMS_PER_BATCH:\n'
' cfg.FPN.RPN_COLLECT_SCALE: {}'.format(cfg.FPN.RPN_COLLECT_SCALE))
if args.num_workers is not None:
cfg.DATA_LOADER.NUM_THREADS = args.num_workers
print('Number of data loading threads: %d' % cfg.DATA_LOADER.NUM_THREADS)
### Overwrite some solver settings from command line arguments
if args.optimizer is not None:
cfg.SOLVER.TYPE = args.optimizer
if args.lr is not None:
cfg.SOLVER.BASE_LR = args.lr
if args.lr_decay_gamma is not None:
cfg.SOLVER.GAMMA = args.lr_decay_gamma
assert_and_infer_cfg()
timers = defaultdict(Timer)
### Dataset ###
timers['roidb'].tic()
roidb, ratio_list, ratio_index = combined_roidb_for_training(
cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES)
timers['roidb'].toc()
roidb_size = len(roidb)
logger.info('{:d} roidb entries'.format(roidb_size))
logger.info('Takes %.2f sec(s) to construct roidb', timers['roidb'].average_time)
# Effective training sample size for one epoch
train_size = roidb_size // args.batch_size * args.batch_size
batchSampler = BatchSampler(
sampler=MinibatchSampler(ratio_list, ratio_index),
batch_size=args.batch_size,
drop_last=True
)
dataset = RoiDataLoader(
roidb,
cfg.MODEL.NUM_CLASSES,
training=True)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_sampler=batchSampler,
num_workers=cfg.DATA_LOADER.NUM_THREADS,
collate_fn=collate_minibatch)
dataiterator = iter(dataloader)
### Model ###
maskRCNN = Generalized_RCNN()
if cfg.CUDA:
maskRCNN.cuda()
### Optimizer ###
gn_param_nameset = set()
for name, module in maskRCNN.named_modules():
if isinstance(module, nn.GroupNorm):
gn_param_nameset.add(name+'.weight')
gn_param_nameset.add(name+'.bias')
gn_params = []
gn_param_names = []
bias_params = []
bias_param_names = []
nonbias_params = []
nonbias_param_names = []
nograd_param_names = []
for key, value in dict(maskRCNN.named_parameters()).items():
if value.requires_grad:
if 'bias' in key:
bias_params.append(value)
bias_param_names.append(key)
elif key in gn_param_nameset:
gn_params.append(value)
gn_param_names.append(key)
else:
nonbias_params.append(value)
nonbias_param_names.append(key)
else:
nograd_param_names.append(key)
assert (gn_param_nameset - set(nograd_param_names) - set(bias_param_names)) == set(gn_param_names)
# Learning rate of 0 is a dummy value to be set properly at the start of training
params = [
{'params': nonbias_params,
'lr': 0,
'weight_decay': cfg.SOLVER.WEIGHT_DECAY},
{'params': bias_params,
'lr': 0 * (cfg.SOLVER.BIAS_DOUBLE_LR + 1),
'weight_decay': cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0},
{'params': gn_params,
'lr': 0,
'weight_decay': cfg.SOLVER.WEIGHT_DECAY_GN}
]
# names of paramerters for each paramter
param_names = [nonbias_param_names, bias_param_names, gn_param_names]
if cfg.SOLVER.TYPE == "SGD":
optimizer = torch.optim.SGD(params, momentum=cfg.SOLVER.MOMENTUM)
elif cfg.SOLVER.TYPE == "Adam":
optimizer = torch.optim.Adam(params)
### Load checkpoint
if args.load_ckpt:
load_name = args.load_ckpt
logging.info("loading checkpoint %s", load_name)
checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.resume:
args.start_step = checkpoint['step'] + 1
if 'train_size' in checkpoint: # For backward compatibility
if checkpoint['train_size'] != train_size:
print('train_size value: %d different from the one in checkpoint: %d'
% (train_size, checkpoint['train_size']))
# reorder the params in optimizer checkpoint's params_groups if needed
# misc_utils.ensure_optimizer_ckpt_params_order(param_names, checkpoint)
# There is a bug in optimizer.load_state_dict on Pytorch 0.3.1.
# However it's fixed on master.
# optimizer.load_state_dict(checkpoint['optimizer'])
misc_utils.load_optimizer_state_dict(optimizer, checkpoint['optimizer'])
del checkpoint
torch.cuda.empty_cache()
if args.load_detectron: #TODO resume for detectron weights (load sgd momentum values)
logging.info("loading Detectron weights %s", args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
lr = optimizer.param_groups[0]['lr'] # lr of non-bias parameters, for commmand line outputs.
maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'],
minibatch=True)
### Training Setups ###
args.run_name = misc_utils.get_run_name() + '_step'
output_dir = misc_utils.get_output_dir(args, args.run_name)
args.cfg_filename = os.path.basename(args.cfg_file)
if not args.no_save:
if not os.path.exists(output_dir):
os.makedirs(output_dir)
blob = {'cfg': yaml.dump(cfg), 'args': args}
with open(os.path.join(output_dir, 'config_and_args.pkl'), 'wb') as f:
pickle.dump(blob, f, pickle.HIGHEST_PROTOCOL)
if args.use_tfboard:
from tensorboardX import SummaryWriter
# Set the Tensorboard logger
tblogger = SummaryWriter(output_dir)
### Training Loop ###
maskRCNN.train()
CHECKPOINT_PERIOD = int(cfg.TRAIN.SNAPSHOT_ITERS / cfg.NUM_GPUS)
# Set index for decay steps
decay_steps_ind = None
for i in range(1, len(cfg.SOLVER.STEPS)):
if cfg.SOLVER.STEPS[i] >= args.start_step:
decay_steps_ind = i
break
if decay_steps_ind is None:
decay_steps_ind = len(cfg.SOLVER.STEPS)
training_stats = TrainingStats(
args,
args.disp_interval,
tblogger if args.use_tfboard and not args.no_save else None)
try:
logger.info('Training starts !')
step = args.start_step
for step in range(args.start_step, cfg.SOLVER.MAX_ITER):
# Warm up
if step < cfg.SOLVER.WARM_UP_ITERS:
method = cfg.SOLVER.WARM_UP_METHOD
if method == 'constant':
warmup_factor = cfg.SOLVER.WARM_UP_FACTOR
elif method == 'linear':
alpha = step / cfg.SOLVER.WARM_UP_ITERS
warmup_factor = cfg.SOLVER.WARM_UP_FACTOR * (1 - alpha) + alpha
else:
raise KeyError('Unknown SOLVER.WARM_UP_METHOD: {}'.format(method))
lr_new = cfg.SOLVER.BASE_LR * warmup_factor
net_utils.update_learning_rate(optimizer, lr, lr_new)
lr = optimizer.param_groups[0]['lr']
assert lr == lr_new
elif step == cfg.SOLVER.WARM_UP_ITERS:
net_utils.update_learning_rate(optimizer, lr, cfg.SOLVER.BASE_LR)
lr = optimizer.param_groups[0]['lr']
assert lr == cfg.SOLVER.BASE_LR
# Learning rate decay
if decay_steps_ind < len(cfg.SOLVER.STEPS) and \
step == cfg.SOLVER.STEPS[decay_steps_ind]:
logger.info('Decay the learning on step %d', step)
lr_new = lr * cfg.SOLVER.GAMMA
net_utils.update_learning_rate(optimizer, lr, lr_new)
lr = optimizer.param_groups[0]['lr']
assert lr == lr_new
decay_steps_ind += 1
training_stats.IterTic()
optimizer.zero_grad()
for inner_iter in range(args.iter_size):
try:
input_data = next(dataiterator)
except StopIteration:
dataiterator = iter(dataloader)
input_data = next(dataiterator)
for key in input_data:
if key != 'roidb': # roidb is a list of ndarrays with inconsistent length
input_data[key] = list(map(Variable, input_data[key]))
net_outputs = maskRCNN(**input_data)
training_stats.UpdateIterStats(net_outputs, inner_iter)
loss = net_outputs['total_loss']
loss.backward()
optimizer.step()
training_stats.IterToc()
training_stats.LogIterStats(step, lr)
if (step+1) % CHECKPOINT_PERIOD == 0:
save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer)
# ---- Training ends ----
# Save last checkpoint
save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer)
except (RuntimeError, KeyboardInterrupt):
del dataiterator
logger.info('Save ckpt on exception ...')
save_ckpt(output_dir, args, step, train_size, maskRCNN, optimizer)
logger.info('Save ckpt done.')
stack_trace = traceback.format_exc()
print(stack_trace)
finally:
if args.use_tfboard and not args.no_save:
tblogger.close()
import glob
import subprocess
from PIL import Image
from torchvision import transforms
import torch.nn.functional as F
# load config
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = len(dataset.classes)
cfg_from_file('configs/baselines/e2e_mask_rcnn_R-50-C4_1x.yaml')
cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights
assert_and_infer_cfg()
# load model
maskRCNN = Generalized_RCNN()
checkpoint = torch.load(
'/home/work/liupeng11/code/Detectron.pytorch/models/e2e_mask_rcnn_R-50-C4_1x.pth',
map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
maskRCNN.eval()
maskRCNN = mynn.DataParallel(maskRCNN,
cpu_keywords=['im_info', 'roidb'],
minibatch=True,
device_ids=[0])
# load image
img_path = "/home/work/liupeng11/code/Detectron.pytorch/demo/sample_images/img1.jpg"
im = cv2.imread(img_path)
# detect bouding boxes and segments
def main():
"""main function"""
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
args = parse_args()
print('Called with args:')
print(args)
assert args.image_dir or args.images
assert bool(args.image_dir) ^ bool(args.images)
if args.dataset.startswith("coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = len(dataset.classes)
elif args.dataset.startswith("keypoints_coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = 2
else:
raise ValueError('Unexpected dataset name: {}'.format(args.dataset))
print('load cfg from file: {}'.format(args.cfg_file))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
assert bool(args.load_ckpt) ^ bool(args.load_detectron), \
'Exactly one of --load_ckpt and --load_detectron should be specified.'
cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights
assert_and_infer_cfg()
maskRCNN = Generalized_RCNN()
if args.cuda:
maskRCNN.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
print("loading checkpoint %s" % (load_name))
checkpoint = torch.load(load_name,
map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.load_detectron:
print("loading detectron weights %s" % args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
maskRCNN = mynn.DataParallel(maskRCNN,
cpu_keywords=['im_info', 'roidb'],
minibatch=True,
device_ids=[0]) # only support single GPU
maskRCNN.eval()
params = list(maskRCNN.parameters())
k = 0
for i in params:
l = 1
for j in i.size():
l *= j
k = k + l
print('zonghe:' + str(k))
if args.image_dir:
imglist = misc_utils.get_imagelist_from_dir(args.image_dir)
else:
imglist = args.images
num_images = len(imglist)
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
for i in xrange(num_images):
print('img', i)
im = cv2.imread(imglist[i])
assert im is not None
timers = defaultdict(Timer)
start = time.time()
cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN,
im,
timers=timers)
class_result_boxes = []
for index, class_boxes in enumerate(cls_boxes):
if len(class_boxes) != 0:
class_boxes = class_boxes.tolist()
results_oneclass = threeD_detect(imglist[i], class_boxes,
index)
class_result_boxes.append(results_oneclass)
save_image = im
color_class = {
'Car': [0, 255, 255],
'Cyclist': [255, 0, 0],
'Pedestrian': [0, 0, 255]
}
for result_boxes in class_result_boxes:
for box in result_boxes:
cv2.rectangle(save_image, (box[0], box[1]), (box[2], box[3]),
color_class[box[-1]], 2)
height = round(box[-2][0], 2)
width = round(box[-2][1], 2)
length = round(box[-2][2], 2)
threeD_info = str(height) + ' ' + str(width) + ' ' + str(
length)
cv2.putText(save_image, threeD_info, (box[0], box[1] - 20),
cv2.FONT_HERSHEY_COMPLEX, 1, (255, 0, 0), 2)
_, imagename = os.path.split(imglist[i])
imagename2 = imagename.split('.')[0]
cv2.imwrite('../output1/%s.png' % imagename2, save_image)
end = time.time()
print(end - start)
im_name, _ = os.path.splitext(os.path.basename(imglist[i]))
vis_utils.vis_one_image(
im[:, :, ::-1], # BGR -> RGB for visualization
im_name,
args.output_dir,
cls_boxes,
cls_segms,
cls_keyps,
dataset=dataset,
box_alpha=0.3,
show_class=True,
thresh=0.7,
kp_thresh=2)
if args.merge_pdfs and num_images > 1:
merge_out_path = '{}/results.pdf'.format(args.output_dir)
if os.path.exists(merge_out_path):
os.remove(merge_out_path)
command = "pdfunite {}/*.pdf {}".format(args.output_dir,
merge_out_path)
subprocess.call(command, shell=True)
def main():
"""main function"""
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
args = parse_args()
print('Called with args:')
print(args)
assert args.image_dir or args.images
assert bool(args.image_dir) ^ bool(args.images)
if args.dataset.startswith("coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = len(dataset.classes)
elif args.dataset.startswith("keypoints_coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = 2
else:
raise ValueError('Unexpected dataset name: {}'.format(args.dataset))
print('load cfg from file: {}'.format(args.cfg_file))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
assert bool(args.load_ckpt) ^ bool(args.load_detectron), \
'Exactly one of --load_ckpt and --load_detectron should be specified.'
cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights
assert_and_infer_cfg()
maskRCNN = Generalized_RCNN()
if args.cuda:
maskRCNN.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
print("loading checkpoint %s" % (load_name))
checkpoint = torch.load(load_name, map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.load_detectron:
print("loading detectron weights %s" % args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
maskRCNN = mynn.DataParallel(maskRCNN, cpu_keywords=['im_info', 'roidb'],
minibatch=True, device_ids=[0]) # only support single GPU
maskRCNN.eval()
if args.image_dir:
imglist = misc_utils.get_imagelist_from_dir(args.image_dir)
else:
imglist = args.images
num_images = len(imglist)
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
for i in xrange(num_images):
print('img', i)
im = cv2.imread(imglist[i])
assert im is not None
timers = defaultdict(Timer)
cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN, im, timers=timers)
im_name, _ = os.path.splitext(os.path.basename(imglist[i]))
vis_utils.vis_one_image(
im[:, :, ::-1], # BGR -> RGB for visualization
im_name,
args.output_dir,
cls_boxes,
cls_segms,
cls_keyps,
dataset=dataset,
box_alpha=0.3,
show_class=True,
thresh=0.7,
kp_thresh=2
)
if args.merge_pdfs and num_images > 1:
merge_out_path = '{}/results.pdf'.format(args.output_dir)
if os.path.exists(merge_out_path):
os.remove(merge_out_path)
command = "pdfunite {}/*.pdf {}".format(args.output_dir,
merge_out_path)
subprocess.call(command, shell=True)
def test_net_on_dataset(args):
dataset = WAD_CVPR2018(args.dataset_dir)
cfg.MODEL.NUM_CLASSES = len(
dataset.eval_class) + 1 # with a background class
print('load cfg from file: {}'.format(args.cfg_file))
cfg_from_file(args.cfg_file)
if args.nms_soft:
cfg.TEST.SOFT_NMS.ENABLED = True
else:
cfg.TEST.NMS = args.nms
cfg.RESNETS.IMAGENET_PRETRAINED = False # Don't need to load imagenet pretrained weights
assert_and_infer_cfg()
maskRCNN = Generalized_RCNN()
maskRCNN.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
print("loading checkpoint %s" % (load_name))
checkpoint = torch.load(load_name,
map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
maskRCNN = mynn.DataParallel(maskRCNN,
cpu_keywords=['im_info', 'roidb'],
minibatch=True,
device_ids=[0]) # only support single GPU
maskRCNN.eval()
imglist_all = misc_utils.get_imagelist_from_dir(dataset.test_image_dir)
if args.nms_soft:
output_dir = os.path.join(
('/').join(args.load_ckpt.split('/')[:-2]),
'Images_' + str(cfg.TEST.SCALE) + '_SOFT_NMS')
elif args.nms:
output_dir = os.path.join(
('/').join(args.load_ckpt.split('/')[:-2]),
'Images_' + str(cfg.TEST.SCALE) + '_NMS_%.2f' % args.nms)
else:
output_dir = os.path.join(('/').join(args.load_ckpt.split('/')[:-2]),
'Images_' + str(cfg.TEST.SCALE))
if cfg.TEST.BBOX_AUG.ENABLED:
output_dir += '_TEST_AUG'
#if args.cls_boxes_confident_threshold < 0.5:
output_dir += '_cls_boxes_confident_threshold_%.1f' % args.cls_boxes_confident_threshold
if not os.path.exists(output_dir):
os.makedirs(output_dir)
output_vis_dir = os.path.join(output_dir, 'Image_Vis')
if not os.path.exists(output_vis_dir):
os.makedirs(output_vis_dir)
args.output_img_dir = os.path.join(output_dir, 'Image_Masks')
if not os.path.exists(args.output_img_dir):
os.makedirs(args.output_img_dir)
output_list_dir = os.path.join(output_dir, 'List_Masks')
if not os.path.exists(output_list_dir):
os.makedirs(output_list_dir)
# A break point
# img_produced = os.listdir(output_list_dir)
# imglist = [x for x in imglist_all if x.split('/')[-1][:-4]+'.txt' not in img_produced]
imglist = imglist_all
if False:
args.vis = True
imglist = [
os.path.join(args.dataset_dir, 'test',
'4f38c1d630209cfb777a3dcbb613ba56.jpg'),
os.path.join(args.dataset_dir, 'test',
'3b08ff3b969200982d0283aedbcaae20.jpg')
]
if not args.range:
start = 0
end = len(imglist)
else:
start = args.range[0]
end = args.range[1]
for i in tqdm(xrange(start, end)):
im = cv2.imread(imglist[i])
assert im is not None
timers = defaultdict(Timer)
im_name, _ = os.path.splitext(os.path.basename(imglist[i]))
args.current_im_name = im_name
cls_boxes, cls_segms, prediction_row = im_detect_all(args,
maskRCNN,
im,
dataset,
timers=timers)
im_name, _ = os.path.splitext(os.path.basename(imglist[i]))
print(im_name)
if args.vis:
vis_utils.vis_one_image_cvpr2018_wad(
im[:, :, ::-1], # BGR -> RGB for visualization
im_name,
output_vis_dir,
cls_boxes,
cls_segms,
None,
dataset=dataset,
box_alpha=0.3,
show_class=True,
thresh=0.5,
kp_thresh=2)
thefile = open(os.path.join(output_list_dir, im_name + '.txt'), 'w')
for item in prediction_row:
thefile.write("%s\n" % item)
def test_net_on_dataset_multigpu(args):
dataset = WAD_CVPR2018(args.dataset_dir)
cfg.MODEL.NUM_CLASSES = len(
dataset.eval_class) + 1 # with a background class
print('load cfg from file: {}'.format(args.cfg_file))
cfg_from_file(args.cfg_file)
cfg.RESNETS.IMAGENET_PRETRAINED = False # Don't need to load imagenet pretrained weights
if args.nms_soft:
cfg.TEST.SOFT_NMS.ENABLED = True
else:
cfg.TEST.NMS = args.nms
assert_and_infer_cfg()
maskRCNN = Generalized_RCNN()
maskRCNN.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
print("loading checkpoint %s" % (load_name))
checkpoint = torch.load(load_name,
map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
maskRCNN = mynn.DataParallel(maskRCNN,
cpu_keywords=['im_info', 'roidb'],
minibatch=True,
device_ids=[0]) # only support single GPU
maskRCNN.eval()
#imglist = misc_utils.get_imagelist_from_dir(dataset.test_image_dir)
img_produced = os.listdir(args.output_list_dir)
imglist_all = misc_utils.get_imagelist_from_dir(dataset.test_image_dir)
imglist = [
x for x in imglist_all
if x.split('/')[-1][:-4] + '.txt' not in img_produced
]
#for i in tqdm(xrange(args.range[0], args.range[1])):
start_idx = len(imglist_all) - args.range[0]
end_idx = len(imglist_all) - args.range[0] + args.range[1]
for i in tqdm(xrange(start_idx, end_idx)):
im = cv2.imread(imglist_all[i])
assert im is not None
timers = defaultdict(Timer)
im_name, _ = os.path.splitext(os.path.basename(imglist[i]))
args.current_im_name = im_name
cls_boxes, cls_segms, prediction_row = im_detect_all(args,
maskRCNN,
im,
dataset,
timers=timers)
im_name, _ = os.path.splitext(os.path.basename(imglist[i]))
print(im_name)
if args.vis:
vis_utils.vis_one_image_cvpr2018_wad(
im[:, :, ::-1], # BGR -> RGB for visualization
im_name,
args.output_vis_dir,
cls_boxes,
cls_segms,
None,
dataset=dataset,
box_alpha=0.3,
show_class=True,
thresh=0.5,
kp_thresh=2)
thefile = open(os.path.join(args.output_list_dir, im_name + '.txt'),
'w')
for item in prediction_row:
thefile.write("%s\n" % item)
def main():
"""main function"""
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
args = parse_args()
print('Called with args:')
print(args)
assert args.image_dir or args.images
assert bool(args.image_dir) ^ bool(args.images)
if args.dataset.startswith("coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = len(dataset.classes)
elif args.dataset.startswith("keypoints_coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = 2
else:
raise ValueError('Unexpected dataset name: {}'.format(args.dataset))
print('load cfg from file: {}'.format(args.cfg_file))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
assert bool(args.load_ckpt) ^ bool(args.load_detectron), \
'Exactly one of --load_ckpt and --load_detectron should be specified.'
cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights
assert_and_infer_cfg()
maskRCNN = Generalized_RCNN()
if args.cuda:
maskRCNN.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
print("loading checkpoint %s" % (load_name))
checkpoint = torch.load(load_name,
map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.load_detectron:
print("loading detectron weights %s" % args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
video_path = args.video
video_save_dir = args.video_save_dir
if not os.path.exists(video_save_dir):
os.mkdir(video_save_dir)
maskRCNN = mynn.DataParallel(maskRCNN,
cpu_keywords=['im_info', 'roidb'],
minibatch=True,
device_ids=[0]) # only support single GPU
maskRCNN.eval()
capture = cv2.VideoCapture(video_path)
frame_count = 0
batch_size = 1
frame_list = []
while True:
ret, frame = capture.read()
# Bail out when the video file ends
if not ret:
break
# Save each frame of the video to a list
frame_count += 1
frame_list.append(frame)
if len(frame_list) == batch_size:
for i, frame in enumerate(frame_list):
timers = defaultdict(Timer)
cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN,
frame,
timers=timers)
name = '{0}'.format(frame_count + i - batch_size)
print('\n' + name)
vis_utils.vis_one_image(
frame[:, :, ::-1], # BGR -> RGB for visualization
name,
video_save_dir,
cls_boxes,
cls_segms,
cls_keyps,
dataset=dataset,
box_alpha=0.3,
show_class=True,
thresh=0.8,
kp_thresh=2,
ext='jpg')
# Clear the frames array to start the next batch
frame_list = []
images = list(glob.iglob(os.path.join(video_save_dir, '*.jpg')))
# Sort the images by integer index
images = sorted(images, key=lambda x: float(os.path.split(x)[1][:-3]))
outvid = os.path.join(video_save_dir, "out.mp4")
make_video(outvid, images, fps=30)
extract_audio = 'ffmpeg -i %s -vn -acodec copy %s/out.aac' % (
video_path, video_save_dir)
subprocess.call(extract_audio, shell=True)
merge_audio_video = "ffmpeg -i %s/out.aac -i %s/out.mp4 -codec copy -shortest %s/final.mp4" % (
video_save_dir, video_save_dir, video_save_dir)
subprocess.call(merge_audio_video, shell=True)
p_tensor.copy_(torch.Tensor(src_blobs[d_name]))
def resnet_weights_name_pattern():
pattern = re.compile(r"conv1_w|conv1_gn_[sb]|res_conv1_.+|res\d+_\d+_.+")
return pattern
if __name__ == '__main__':
"""Testing"""
from pprint import pprint
import sys
sys.path.insert(0, '..')
from modeling.model_builder import Generalized_RCNN
from core.config import cfg, cfg_from_file
cfg.MODEL.NUM_CLASSES = 81
cfg_from_file('../../cfgs/res50_mask.yml')
net = Generalized_RCNN()
# pprint(list(net.state_dict().keys()), width=1)
mapping, orphans = net.detectron_weight_mapping
state_dict = net.state_dict()
for k in mapping.keys():
assert k in state_dict, '%s' % k
rest = set(state_dict.keys()) - set(mapping.keys())
assert len(rest) == 0
def main():
"""main function"""
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
args = parse_args()
print('Called with args:')
print(args)
assert args.image_dir or args.images
assert bool(args.image_dir) ^ bool(args.images)
dataset = datasets.get_hospital_dataset()
cfg.MODEL.NUM_CLASSES = 20 # with bg
num_class = cfg.MODEL.NUM_CLASSES
sents = dataset.sents
th_cls = dataset.th_cls
cls2eng = dataset.cls2eng
eng2type = dataset.eng2type
print('load cfg from file: {}'.format(args.cfg_file))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
assert bool(args.load_ckpt) ^ bool(args.load_detectron), \
'Exactly one of --load_ckpt and --load_detectron should be specified.'
cfg.MODEL.LOAD_IMAGENET_PRETRAINED_WEIGHTS = False # Don't need to load imagenet pretrained weights
assert_and_infer_cfg()
maskRCNN = Generalized_RCNN()
if args.cuda:
maskRCNN.cuda()
if args.load_ckpt:
load_name = args.load_ckpt
print("loading checkpoint %s" % (load_name))
checkpoint = torch.load(load_name,
map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.load_detectron:
print("loading detectron weights %s" % args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
maskRCNN = mynn.DataParallel(maskRCNN,
cpu_keywords=['im_info', 'roidb'],
minibatch=True,
device_ids=[0]) # only support single GPU
maskRCNN.eval()
if args.image_dir:
imglist = misc_utils.get_imagelist_from_dir(args.image_dir)
else:
imglist = args.images
num_images = len(imglist)
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
for i in xrange(num_images): # for each image
print('img', i)
im = cv2.imread(imglist[i])
assert im is not None
# segmentation
# d = segment(im)
# pdb.set_trace()
timers = defaultdict(Timer)
# detection
cls_boxes, cls_segms, cls_keyps = im_detect_all(maskRCNN,
im,
timers=timers)
# first we collect boxes from all classes
dets_total = np.empty([0, 6], dtype=np.float32)
for cls in range(1, num_class): # for each cls
dets = cls_boxes[cls]
if dets.shape[0] == 0:
continue
dets_extend = np.pad(
dets,
((0, 0),
(0, 1)), # add 0 rows above, below and left, but 1 row right
mode='constant',
constant_values=cls) # append cls to dets
dets_total = np.vstack((dets_total, dets_extend))
# then use a loose NMS to make each region has only one symptom
keep = box_utils.nms(dets_total, 0.7)
nms_dets = dets_total[keep, :]
# iterate through remained boxes
report, healthy = '', True
have_sym_of_cls = [False for _ in range(num_class)]
n = nms_dets.shape[0]
final_results = [] # return to the web
for idx in range(n): # for each region
th, cls = nms_dets[idx, -2], int(nms_dets[idx, -1])
if th > th_cls[cls]: # diagnosed to have the sym
report += sents[cls][1]
have_sym_of_cls[cls] = True
healthy = False
ename = cls2eng[int(cls)]
_type = eng2type[ename]
final_results.append({
'name': ename,
'type': _type,
'box': list(nms_dets[idx, 0:4])
})
for cls in range(1, num_class): # for each cls
if not have_sym_of_cls[cls]: # if have no sym of this cls
report += sents[cls][0]
if healthy:
report = sents[0][0]
print(report)
pdb.set_trace()
# healthy = True # flag indicating healthy or not
# for cls in range(1, num_class): # for each cls
# dets = cls_boxes[cls]
# if dets.shape[0] == 0:
# report += sents[cls][0]
# continue
# n = dets.shape[0]
# flag = False # indicates if the sym exists
# for k in range(n): # for each region
# if dets[k, -1] > th_cls[cls]: # above threshold for this cls, means have this cls of symptom
# report += sents[cls][1]
# flag = True
# healthy = False
# if not flag: # don't have this symptom
# report += sents[cls][0]
#
# if healthy: # use the report for healthy people
# report = sents[0][0]
im_name, _ = os.path.splitext(os.path.basename(imglist[i]))
# vis_utils.vis_one_image(
# im[:, :, ::-1], # BGR -> RGB for visualization
# im_name,
# args.output_dir,
# cls_boxes,
# cls_segms,
# cls_keyps,
# dataset=dataset,
# box_alpha=0.3,
# show_class=True,
# thresh=0.05,
# kp_thresh=2
# )
if args.merge_pdfs and num_images > 1:
merge_out_path = '{}/results.pdf'.format(args.output_dir)
if os.path.exists(merge_out_path):
os.remove(merge_out_path)
command = "pdfunite {}/*.pdf {}".format(args.output_dir,
merge_out_path)
subprocess.call(command, shell=True)
def main():
"""Main function"""
args = parse_args()
print('Called with args:')
print(args)
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
if args.cuda or cfg.NUM_GPUS > 0:
cfg.CUDA = True
else:
raise ValueError("Need Cuda device to run !")
if args.dataset == "coco2017":
cfg.TRAIN.DATASETS = ('coco_2017_train', )
cfg.MODEL.NUM_CLASSES = 81
elif args.dataset == "keypoints_coco2017":
cfg.TRAIN.DATASETS = ('keypoints_coco_2017_train', )
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset == "common":
cfg.TRAIN.DATASETS = ('common_train', )
cfg.MODEL.NUM_CLASSES = 81
else:
raise ValueError("Unexpected args.dataset: {}".format(args.dataset))
cfg_from_file(args.cfg_file)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
### Adaptively adjust some configs ###
original_batch_size = cfg.NUM_GPUS * cfg.TRAIN.IMS_PER_BATCH
if args.batch_size is None:
args.batch_size = original_batch_size
cfg.NUM_GPUS = torch.cuda.device_count()
assert (args.batch_size % cfg.NUM_GPUS) == 0, \
'batch_size: %d, NUM_GPUS: %d' % (args.batch_size, cfg.NUM_GPUS)
cfg.TRAIN.IMS_PER_BATCH = args.batch_size // cfg.NUM_GPUS
print('Batch size change from {} (in config file) to {}'.format(
original_batch_size, args.batch_size))
print('NUM_GPUs: %d, TRAIN.IMS_PER_BATCH: %d' %
(cfg.NUM_GPUS, cfg.TRAIN.IMS_PER_BATCH))
if args.num_workers is not None:
cfg.DATA_LOADER.NUM_THREADS = args.num_workers
print('Number of data loading threads: %d' % cfg.DATA_LOADER.NUM_THREADS)
### Adjust learning based on batch size change linearly
old_base_lr = cfg.SOLVER.BASE_LR
cfg.SOLVER.BASE_LR *= args.batch_size / original_batch_size
print('Adjust BASE_LR linearly according to batch size change: {} --> {}'.
format(old_base_lr, cfg.SOLVER.BASE_LR))
### Overwrite some solver settings from command line arguments
if args.optimizer is not None:
cfg.SOLVER.TYPE = args.optimizer
if args.lr is not None:
cfg.SOLVER.BASE_LR = args.lr
if args.lr_decay_gamma is not None:
cfg.SOLVER.GAMMA = args.lr_decay_gamma
timers = defaultdict(Timer)
### Dataset ###
timers['roidb'].tic()
roidb, ratio_list, ratio_index = combined_roidb_for_training(
cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES)
timers['roidb'].toc()
train_size = len(roidb)
logger.info('{:d} roidb entries'.format(train_size))
logger.info('Takes %.2f sec(s) to construct roidb',
timers['roidb'].average_time)
sampler = MinibatchSampler(ratio_list, ratio_index)
dataset = RoiDataLoader(roidb, cfg.MODEL.NUM_CLASSES, training=True)
dataloader = torch.utils.data.DataLoader(
dataset,
batch_size=args.batch_size,
sampler=sampler,
num_workers=cfg.DATA_LOADER.NUM_THREADS,
collate_fn=collate_minibatch)
assert_and_infer_cfg()
### Model ###
maskRCNN = Generalized_RCNN()
if cfg.CUDA:
maskRCNN.cuda()
### Optimizer ###
bias_params = []
nonbias_params = []
for key, value in dict(maskRCNN.named_parameters()).items():
if value.requires_grad:
if 'bias' in key:
bias_params.append(value)
else:
nonbias_params.append(value)
params = [{
'params': nonbias_params,
'lr': cfg.SOLVER.BASE_LR,
'weight_decay': cfg.SOLVER.WEIGHT_DECAY
}, {
'params':
bias_params,
'lr':
cfg.SOLVER.BASE_LR * (cfg.SOLVER.BIAS_DOUBLE_LR + 1),
'weight_decay':
cfg.SOLVER.WEIGHT_DECAY if cfg.SOLVER.BIAS_WEIGHT_DECAY else 0
}]
if cfg.SOLVER.TYPE == "SGD":
optimizer = torch.optim.SGD(params, momentum=cfg.SOLVER.MOMENTUM)
elif cfg.SOLVER.TYPE == "Adam":
optimizer = torch.optim.Adam(params)
### Load checkpoint
if args.load_ckpt:
load_name = args.load_ckpt
logging.info("loading checkpoint %s", load_name)
checkpoint = torch.load(load_name,
map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN, checkpoint['model'])
if args.resume:
assert checkpoint['iters_per_epoch'] == train_size // args.batch_size, \
"iters_per_epoch should match for resume"
# There is a bug in optimizer.load_state_dict on Pytorch 0.3.1.
# However it's fixed on master.
# optimizer.load_state_dict(checkpoint['optimizer'])
misc_utils.load_optimizer_state_dict(optimizer,
checkpoint['optimizer'])
if checkpoint['step'] == (checkpoint['iters_per_epoch'] - 1):
# Resume from end of an epoch
args.start_epoch = checkpoint['epoch'] + 1
args.start_iter = 0
else:
# Resume from the middle of an epoch.
# NOTE: dataloader is not synced with previous state
args.start_epoch = checkpoint['epoch']
args.start_iter = checkpoint['step'] + 1
del checkpoint
torch.cuda.empty_cache()
if args.load_detectron: #TODO resume for detectron weights (load sgd momentum values)
logging.info("loading Detectron weights %s", args.load_detectron)
load_detectron_weight(maskRCNN, args.load_detectron)
lr = optimizer.param_groups[0][
'lr'] # lr of non-bias parameters, for commmand line outputs.
maskRCNN = mynn.DataParallel(maskRCNN,
cpu_keywords=['im_info', 'roidb'],
minibatch=True)
### Training Setups ###
args.run_name = misc_utils.get_run_name()
output_dir = misc_utils.get_output_dir(args, args.run_name)
args.cfg_filename = os.path.basename(args.cfg_file)
if not args.no_save:
if not os.path.exists(output_dir):
os.makedirs(output_dir)
blob = {'cfg': yaml.dump(cfg), 'args': args}
with open(os.path.join(output_dir, 'config_and_args.pkl'), 'wb') as f:
pickle.dump(blob, f, pickle.HIGHEST_PROTOCOL)
if args.use_tfboard:
from tensorboardX import SummaryWriter
# Set the Tensorboard logger
tblogger = SummaryWriter(output_dir)
### Training Loop ###
maskRCNN.train()
training_stats = TrainingStats(
args, args.disp_interval,
tblogger if args.use_tfboard and not args.no_save else None)
iters_per_epoch = int(train_size / args.batch_size) # drop last
args.iters_per_epoch = iters_per_epoch
ckpt_interval_per_epoch = iters_per_epoch // args.ckpt_num_per_epoch
try:
logger.info('Training starts !')
args.step = args.start_iter
global_step = iters_per_epoch * args.start_epoch + args.step
for args.epoch in range(args.start_epoch,
args.start_epoch + args.num_epochs):
# ---- Start of epoch ----
# adjust learning rate
if args.lr_decay_epochs and args.epoch == args.lr_decay_epochs[
0] and args.start_iter == 0:
args.lr_decay_epochs.pop(0)
net_utils.decay_learning_rate(optimizer, lr, cfg.SOLVER.GAMMA)
lr *= cfg.SOLVER.GAMMA
for args.step, input_data in zip(
range(args.start_iter, iters_per_epoch), dataloader):
for key in input_data:
if key != 'roidb': # roidb is a list of ndarrays with inconsistent length
input_data[key] = list(map(Variable, input_data[key]))
training_stats.IterTic()
net_outputs = maskRCNN(**input_data)
training_stats.UpdateIterStats(net_outputs)
loss = net_outputs['total_loss']
optimizer.zero_grad()
loss.backward()
optimizer.step()
training_stats.IterToc()
if (args.step + 1) % ckpt_interval_per_epoch == 0:
net_utils.save_ckpt(output_dir, args, maskRCNN, optimizer)
if args.step % args.disp_interval == 0:
log_training_stats(training_stats, global_step, lr)
global_step += 1
# ---- End of epoch ----
# save checkpoint
net_utils.save_ckpt(output_dir, args, maskRCNN, optimizer)
# reset starting iter number after first epoch
args.start_iter = 0
# ---- Training ends ----
if iters_per_epoch % args.disp_interval != 0:
# log last stats at the end
log_training_stats(training_stats, global_step, lr)
except (RuntimeError, KeyboardInterrupt):
logger.info('Save ckpt on exception ...')
net_utils.save_ckpt(output_dir, args, maskRCNN, optimizer)
logger.info('Save ckpt done.')
stack_trace = traceback.format_exc()
print(stack_trace)
finally:
if args.use_tfboard and not args.no_save:
tblogger.close()
def main():
"""main function"""
if not torch.cuda.is_available():
sys.exit("Need a CUDA device to run the code.")
args = parse_args()
print('Called with args:')
print(args)
assert args.image_dir or args.images
assert bool(args.image_dir) ^ bool(args.images)
if args.dataset.startswith("coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = len(dataset.classes)
elif args.dataset.startswith("keypoints_coco"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = 2
elif args.dataset.startswith("keypoints_carfusion"):
dataset = datasets.get_coco_dataset()
cfg.MODEL.NUM_CLASSES = 2
else:
raise ValueError('Unexpected dataset name: {}'.format(args.dataset))
assert bool(args.load_ckpt_car) ^ bool(args.load_detectron), \
'Exactly one of --load_ckpt and --load_detectron should be specified.'
print('load cfg from file: {}'.format(args.cfg_file_person))
cfg_from_file(args.cfg_file_person)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
cfg.RESNETS.IMAGENET_PRETRAINED = False # Don't need to load imagenet pretrained weights
assert_and_infer_cfg()
maskRCNN_person = Generalized_RCNN()
#print('load cfg from file: {}'.format(args.cfg_file_person))
#cfg_from_file(args.cfg_file_person)
#assert_and_infer_cfg()
#maskRCNN_person = Generalized_RCNN()
if args.visualize:
save_image = True
if args.cuda:
maskRCNN_person.cuda()
if args.load_ckpt_person:
load_name = args.load_ckpt_person
print("loading checkpoint for person %s" % (load_name))
checkpoint_person = torch.load(
load_name, map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN_person, checkpoint_person['model'])
if args.load_detectron:
print("loading detectron weights %s" % args.load_detectron)
load_detectron_weight(maskRCNN_car, args.load_detectron)
maskRCNN_person = mynn.DataParallel(
maskRCNN_person,
cpu_keywords=['im_info', 'roidb'],
minibatch=True,
device_ids=[0]) # only support single GPU
maskRCNN_person.eval()
print('load cfg from file: {}'.format(args.cfg_file_car))
cfg_from_file(args.cfg_file_car)
if args.set_cfgs is not None:
cfg_from_list(args.set_cfgs)
#cfg.RESNETS.IMAGENET_PRETRAINED = False # Don't need to load imagenet pretrained weights
#assert_and_infer_cfg()
maskRCNN_car = Generalized_RCNN()
if args.cuda:
maskRCNN_car.cuda()
if args.load_ckpt_car:
load_name = args.load_ckpt_car
print("loading checkpoint for car %s" % (load_name))
checkpoint_car = torch.load(load_name,
map_location=lambda storage, loc: storage)
net_utils.load_ckpt(maskRCNN_car, checkpoint_car['model'])
maskRCNN_car = mynn.DataParallel(maskRCNN_car,
cpu_keywords=['im_info', 'roidb'],
minibatch=True,
device_ids=[0]) # only support single GPU
maskRCNN_car.eval()
if args.image_dir:
imglist = misc_utils.get_imagelist_from_dir(args.image_dir)
else:
imglist = args.images
num_images = len(imglist)
if not os.path.exists(args.output_dir):
os.makedirs(args.output_dir)
#imglist.sort(key=lambda f: int(filter(str.isdigit, f)))
l = imglist
try:
imglist = sorted(
l, key=lambda x: int(os.path.splitext(x)[0].split('/')[-1]))
except:
print('images couldnot be sorted')
for i in xrange(num_images):
print('img', i, ' out of ', num_images, ' filename: ',
imglist[i].split('/')[-1], ' in camera',
imglist[i].split('/')[-2])
im = cv2.imread(imglist[i])
try:
assert im is not None
except:
continue
timers = defaultdict(Timer)
im_name, _ = os.path.splitext(os.path.basename(imglist[i]))
#print(im_name)
output_name = os.path.basename(im_name) + '.txt'
output_file = os.path.join(args.output_dir, '{}'.format(output_name))
text_file = open(output_file, "w")
cfg_from_file(args.cfg_file_car)
cls_boxes_car, cls_segms_car, cls_keyps_car, features_car = im_detect_all(
maskRCNN_car, im, timers=timers)
if len(cls_boxes_car[1]) > 0:
features_car = features_car.data.cpu().numpy()
#print(loop,loop2)
#print(distance_matrix)
#plt.figure()
#plot_confusion_matrix(distance_matrix, classes=[0,1,2,3,4,5,6],
#title='Confusion matrix, without normalization')
#fig = plt.figure()
#ax = fig.add_subplot(1,1,1)
#ax.set_aspect('equal')
#plt.imshow(distance_matrix, interpolation='nearest', cmap=plt.cm.ocean)
#plt.colorbar()
#plt.show()
#fig.savefig('1.png')
count = 0
filename = 'finalized_model.txt'
loaded_model = pickle.load(open(filename, 'rb'))
for ind, bb in enumerate(cls_boxes_car[1]):
string = str(count)
keyps = [k for klist in cls_keyps_car for k in klist]
bb_new = [bb[0], bb[1], bb[2] - bb[0], bb[3] - bb[1]]
features = features_car[ind, :, :, :].flatten()
pca_feature = []
pca_feature.append(np.transpose(features.astype(np.float)))
#print(pca_feature)
features = loaded_model.transform(pca_feature)
features = features[0] #loaded_model.transform(pca_feature)
if bb[4] < 0.5:
continue
#for bb_ind,val in enumerate(bb_new):
# string = string + ',' + str(val)
for kp_ind, kp in enumerate(keyps[ind][0]):
string = string + ',' + str(kp) + ',' + str(
keyps[ind][1][kp_ind]) + ',' + str(
int(keyps[ind][2][kp_ind]))
for feature_ind, feature in enumerate(features):
string = string + ',' + str(feature)
string = string + ',car'
text_file.write(string)
text_file.write('\n')
#print(string)
count = count + 1
cfg_from_file(args.cfg_file_person)
cls_boxes_person, cls_segms_person, cls_keyps_person, features_person = im_detect_all(
maskRCNN_person, im, timers=timers)
if len(cls_boxes_person[1]) > 0:
features_person = features_person.data.cpu().numpy()
for ind, bb in enumerate(cls_boxes_person[1]):
string = str(count)
keyps = [k for klist in cls_keyps_person for k in klist]
bb_new = [bb[0], bb[1], bb[2] - bb[0], bb[3] - bb[1]]
features = features_person[ind, :, :, :].flatten()
pca_feature = []
pca_feature.append(np.transpose(features.astype(np.float)))
#print(pca_feature)
features = loaded_model.transform(pca_feature)
features = features[0] #loaded_model.transform(pca_feature)
#features = loaded_model.transform(np.transpose(features.astype(np.float)))
# print(features)
if bb[4] < 0.5:
continue
for bb_ind, val in enumerate(bb_new):
string = string + ',' + str(val)
for kp_ind, kp in enumerate(keyps[ind][0]):
string = string + ',' + str(kp) + ',' + str(
keyps[ind][1][kp_ind]) + ',' + str(
int(keyps[ind][2][kp_ind]))
for feature_ind, feature in enumerate(features):
string = string + ',' + str(feature)
string = string + ',person'
text_file.write(string)
text_file.write('\n')
#print(string)
count = count + 1
if save_image == True:
im_name, _ = os.path.splitext(os.path.basename(imglist[i]))
image_car = vis_utils.vis_one_image_car(
im[:, :, ::-1], # BGR -> RGB for visualization
im_name,
args.output_dir,
cls_boxes_car,
cls_segms_car,
cls_keyps_car,
dataset=dataset,
box_alpha=0.3,
show_class=True,
thresh=0.5,
kp_thresh=0.1)
output_name = os.path.basename(im_name) + '.png'
im = cv2.imread(
os.path.join(args.output_dir, '{}'.format(output_name)))
if im is None:
continue
continue
vis_utils.vis_one_image(
im[:, :, ::-1], # BGR -> RGB for visualization
im_name,
args.output_dir,
cls_boxes_person,
cls_segms_person,
cls_keyps_person,
dataset=dataset,
box_alpha=0.3,
show_class=True,
thresh=0.5,
kp_thresh=10)
if args.merge_pdfs and num_images > 1 and save_image == True:
merge_out_path = '{}/results.pdf'.format(args.output_dir)
if os.path.exists(merge_out_path):
os.remove(merge_out_path)
command = "pdfunite {}/*.pdf {}".format(args.output_dir,
merge_out_path)
subprocess.call(command, shell=True)
