def initialize_model_from_cfg(gpu_id=0):
"""Initialize a model from the global cfg. Loads test-time weights and
creates the networks in the Caffe2 workspace.
"""
model = model_builder.create(cfg.MODEL.TYPE, train=False, gpu_id=gpu_id)
net_utils.initialize_gpu_from_weights_file(
model,
cfg.TEST.WEIGHTS,
gpu_id=gpu_id,
)
model_builder.add_inference_inputs(model)
workspace.CreateNet(model.net)
workspace.CreateNet(model.conv_body_net)
if cfg.MODEL.MASK_ON:
workspace.CreateNet(model.mask_net)
if cfg.MODEL.KEYPOINTS_ON:
workspace.CreateNet(model.keypoint_net)
graph = net_drawer.GetPydotGraphMinimal(model.net.Proto().op,
rankdir="LR",
minimal_dependency=True)
#display.Image(graph.create_png(), height=1800)
graph.write_png('graph.png')
with open(os.path.join("video/keypoint", "net.pbtxt"), 'w') as fid:
fid.write(str(model.net.Proto()))
return model
def initialize_model_from_cfg():
def create_input_blobs(net_def):
for op in net_def.op:
for blob_in in op.input:
if not workspace.HasBlob(blob_in):
workspace.CreateBlob(blob_in)
model = model_builder.create(
cfg.MODEL.TYPE, train=False,
init_params=cfg.TEST.INIT_RANDOM_VARS_BEFORE_LOADING)
model_builder.add_inputs(model)
if cfg.TEST.INIT_RANDOM_VARS_BEFORE_LOADING:
workspace.RunNetOnce(model.param_init_net)
net_utils.initialize_from_weights_file(
model, cfg.TEST.WEIGHTS, broadcast=False)
create_input_blobs(model.net.Proto())
workspace.CreateNet(model.net)
workspace.CreateNet(model.conv_body_net)
if cfg.MODEL.MASK_ON:
create_input_blobs(model.mask_net.Proto())
workspace.CreateNet(model.mask_net)
if cfg.MODEL.KEYPOINTS_ON:
create_input_blobs(model.keypoint_net.Proto())
workspace.CreateNet(model.keypoint_net)
return model
def create_model():
start_iter = 0
if cfg.CLUSTER.ON_CLUSTER and cfg.CLUSTER.AUTO_RESUME:
import re
output_dir = get_output_dir(training=True)
final_path = os.path.join(output_dir, 'model_final.pkl')
if os.path.exists(final_path):
logger.info('model_final.pkl exists; no need to train!')
return None, None, {'final': final_path}
files = os.listdir(output_dir)
for f in files:
iter_string = re.findall(r'(?<=model_iter)\d+(?=\.pkl)', f)
if len(iter_string) > 0:
checkpoint_iter = int(iter_string[0])
if checkpoint_iter > start_iter:
# Start one iteration immediately after the checkpoint iter
start_iter = checkpoint_iter + 1
resume_weights_file = f
if start_iter > 0:
cfg.TRAIN.WEIGHTS = os.path.join(output_dir, resume_weights_file)
logger.info(
'========> Resuming from checkpoint {} with start iter {}'.
format(cfg.TRAIN.WEIGHTS, start_iter))
logger.info('Building nework: {}'.format(cfg.MODEL.TYPE))
model = model_builder.create(cfg.MODEL.TYPE, train=True)
if cfg.MEMONGER:
optimize_memory(model)
workspace.RunNetOnce(model.param_init_net)
return model, start_iter, {}
def initialize_model_from_cfg():
def create_input_blobs(net_def):
for op in net_def.op:
for blob_in in op.input:
if not workspace.HasBlob(blob_in):
workspace.CreateBlob(blob_in)
model = model_builder.create(
cfg.MODEL.TYPE,
train=False,
init_params=cfg.TEST.INIT_RANDOM_VARS_BEFORE_LOADING)
model_builder.add_inputs(model)
if cfg.TEST.INIT_RANDOM_VARS_BEFORE_LOADING:
workspace.RunNetOnce(model.param_init_net)
net_utils.initialize_from_weights_file(model,
cfg.TEST.WEIGHTS,
broadcast=False)
create_input_blobs(model.net.Proto())
workspace.CreateNet(model.net)
workspace.CreateNet(model.conv_body_net)
if cfg.MODEL.MASK_ON:
create_input_blobs(model.mask_net.Proto())
workspace.CreateNet(model.mask_net)
if cfg.MODEL.KEYPOINTS_ON:
create_input_blobs(model.keypoint_net.Proto())
workspace.CreateNet(model.keypoint_net)
return model
def generate_rpn_on_range(ind_range=None):
assert cfg.TEST.WEIGHTS != '', \
'TEST.WEIGHTS must be set to the model file to test'
assert cfg.TEST.DATASET != '', \
'TEST.DATASET must be set to the dataset name to test'
assert cfg.MODEL.RPN_ONLY or cfg.MODEL.FASTER_RCNN
im_list, start_ind, end_ind, total_num_images = get_image_list(ind_range)
output_dir = get_output_dir(training=False)
logger.info(
'Output will be saved to: {:s}'.format(os.path.abspath(output_dir)))
model = model_builder.create(cfg.MODEL.TYPE, train=False)
model_builder.add_inputs(model)
nu.initialize_from_weights_file(model, cfg.TEST.WEIGHTS)
workspace.CreateNet(model.net)
boxes, scores, ids = im_list_proposals(
model,
im_list,
start_ind=start_ind,
end_ind=end_ind,
total_num_images=total_num_images)
cfg_yaml = yaml.dump(cfg)
if ind_range is not None:
rpn_name = 'rpn_proposals_range_%s_%s.pkl' % tuple(ind_range)
else:
rpn_name = 'rpn_proposals.pkl'
rpn_file = os.path.join(output_dir, rpn_name)
robust_pickle_dump(
dict(boxes=boxes, scores=scores, ids=ids, cfg=cfg_yaml), rpn_file)
logger.info('Wrote RPN proposals to {}'.format(os.path.abspath(rpn_file)))
return boxes, scores, ids, rpn_file
def test_restore_checkpoint():
# Create Model
model = model_builder.create(cfg.MODEL.TYPE, train=True)
add_momentum_init_ops(model)
init_weights(model)
# Fill input blobs
roidb = combined_roidb_for_training(
cfg.TRAIN.DATASETS, cfg.TRAIN.PROPOSAL_FILES
)
model_builder.add_training_inputs(model, roidb=roidb)
workspace.CreateNet(model.net)
# Bookkeeping for checkpoint creation
iter_num = 0
checkpoints = {}
output_dir = get_output_dir(cfg.TRAIN.DATASETS, training=True)
chk_file_path = os.path.join(output_dir, 'model_iter{}.pkl'.format(iter_num))
checkpoints[iter_num] = chk_file_path
# Save model weights
nu.save_model_to_weights_file(checkpoints[iter_num], model)
orig_gpu_0_params, orig_all_params = get_params(model)
# Change the model weights
init_weights(model)
# Reload the weights in the model
nu.initialize_gpu_from_weights_file(model, chk_file_path, gpu_id=0)
nu.broadcast_parameters(model)
shutil.rmtree(cfg.OUTPUT_DIR)
_, restored_all_params = get_params(model)
# Check if all params are loaded correctly
for scoped_name, blob in orig_all_params.items():
np.testing.assert_array_equal(blob, restored_all_params[scoped_name])
# Check if broadcast_parameters works
for scoped_name, blob in restored_all_params.items():
unscoped_name = c2_utils.UnscopeName(scoped_name)
np.testing.assert_array_equal(blob, orig_gpu_0_params[unscoped_name])
def _test_std(self):
current_dir = osp.dirname(osp.realpath(__file__))
cfg_file = osp.join(current_dir, '..', 'configs', 'R-50_1x.yaml')
merge_cfg_from_file(cfg_file)
cfg.TEST.WEIGHTS = osp.join(
current_dir, '..', 'outputs', 'train',
'coco_2014_train+coco_2014_valminusminival', 'R-50_1x', 'default',
'model_final.pkl')
cfg.RETINANET.INFERENCE_TH = 0.
dataset = JsonDataset('coco_2014_minival')
roidb = dataset.get_roidb()
model = model_builder.create(cfg.MODEL.TYPE, train=False, gpu_id=0)
utils.net.initialize_gpu_from_weights_file(model,
cfg.TEST.WEIGHTS,
gpu_id=0)
model_builder.add_inference_inputs(model)
workspace.CreateNet(model.net)
workspace.CreateNet(model.conv_body_net)
num_images = len(roidb)
num_classes = cfg.MODEL.NUM_CLASSES
entry = roidb[0]
im = cv2.imread(entry['image'])
with utils.c2.NamedCudaScope(0):
cls_boxes, cls_preds, cls_probs, box_preds, anchors, im_info = im_detect_bbox(
model, im, debug=True)
workspace.ResetWorkspace()
return cls_preds, cls_probs, box_preds, anchors, im_info
def test_retinanet(ind_range=None):
"""
Test RetinaNet model either on the entire dataset or the subset of dataset
specified by the index range
"""
assert cfg.RETINANET.RETINANET_ON, \
'RETINANET_ON must be set for testing RetinaNet model'
output_dir = get_output_dir(training=False)
dataset = JsonDataset(cfg.TEST.DATASET)
roidb = dataset.get_roidb()
if ind_range is not None:
start, end = ind_range
roidb = roidb[start:end]
# Create and load the model
model = model_builder.create(cfg.MODEL.TYPE, train=False)
if cfg.TEST.WEIGHTS:
nu.initialize_from_weights_file(
model, cfg.TEST.WEIGHTS, broadcast=False
)
model_builder.add_inference_inputs(model)
workspace.CreateNet(model.net)
# Compute the detections
all_boxes = im_list_detections(model, roidb)
# Save the detections
cfg_yaml = yaml.dump(cfg)
if ind_range is not None:
det_name = 'detection_range_%s_%s.pkl' % tuple(ind_range)
else:
det_name = 'detections.pkl'
det_file = os.path.join(output_dir, det_name)
save_object(
dict(all_boxes=all_boxes, cfg=cfg_yaml),
det_file)
logger.info('Wrote detections to: {}'.format(os.path.abspath(det_file)))
return all_boxes
def get_model(cfg_file, weights_file):
merge_cfg_from_file(cfg_file)
cfg.TRAIN.WEIGHTS = '' # NOTE: do not download pretrained model weights
cfg.TEST.WEIGHTS = weights_file
cfg.NUM_GPUS = 1
assert_and_infer_cfg()
#according the cfg to bulid model
model = model_builder.create(cfg.MODEL.TYPE,train=True)
return model
def initialize_model_from_cfg(weights_file, gpu_id=0):
"""Initialize a model from the global cfg. Loads test-time weights and
creates the networks in the Caffe2 workspace.
"""
model = model_builder.create(cfg.MODEL.TYPE, train=False, gpu_id=gpu_id)
net_utils.initialize_gpu_from_weights_file(
model, weights_file, gpu_id=gpu_id,
)
model_builder.add_inference_inputs(model)
workspace.CreateNet(model.net)
workspace.CreateNet(model.conv_body_net)
return model
def initialize_model_from_cfg(gpu_id=0):
"""Initialize a model from the global cfg. Loads test-time weights and
creates the networks in the Caffe2 workspace.
"""
model = model_builder.create(cfg.MODEL.TYPE, train=False, gpu_id=gpu_id)
net_utils.initialize_gpu_from_weights_file(
model, cfg.TEST.WEIGHTS, gpu_id=gpu_id,
)
model_builder.add_inference_inputs(model)
workspace.CreateNet(model.net)
workspace.CreateNet(model.conv_body_net)
if cfg.MODEL.MASK_ON:
workspace.CreateNet(model.mask_net)
if cfg.MODEL.KEYPOINTS_ON:
workspace.CreateNet(model.keypoint_net)
return model
def initialize_model_from_cfg():
"""Initialize a model from the global cfg. Loads test-time weights and
creates the networks in the Caffe2 workspace.
"""
model = model_builder.create(cfg.MODEL.TYPE, train=False)
net_utils.initialize_from_weights_file(
model, cfg.TEST.WEIGHTS, broadcast=False
)
model_builder.add_inference_inputs(model)
workspace.CreateNet(model.net)
workspace.CreateNet(model.conv_body_net)
if cfg.MODEL.MASK_ON:
workspace.CreateNet(model.mask_net)
if cfg.MODEL.KEYPOINTS_ON:
workspace.CreateNet(model.keypoint_net)
return model
def generate_rpn_on_range(
weights_file,
dataset_name,
_proposal_file_ignored,
output_dir,
ind_range=None,
gpu_id=0
):
"""Run inference on all images in a dataset or over an index range of images
in a dataset using a single GPU.
"""
assert cfg.MODEL.RPN_ONLY or cfg.MODEL.FASTER_RCNN
roidb, start_ind, end_ind, total_num_images = get_roidb(
dataset_name, ind_range
)
logger.info(
'Output will be saved to: {:s}'.format(os.path.abspath(output_dir))
)
model = model_builder.create(cfg.MODEL.TYPE, train=False, gpu_id=gpu_id)
nu.initialize_gpu_from_weights_file(
model, weights_file, gpu_id=gpu_id,
)
model_builder.add_inference_inputs(model)
workspace.CreateNet(model.net)
boxes, scores, ids = generate_proposals_on_roidb(
model,
roidb,
start_ind=start_ind,
end_ind=end_ind,
total_num_images=total_num_images,
gpu_id=gpu_id,
)
cfg_yaml = yaml.dump(cfg)
if ind_range is not None:
rpn_name = 'rpn_proposals_range_%s_%s.pkl' % tuple(ind_range)
else:
rpn_name = 'rpn_proposals.pkl'
rpn_file = os.path.join(output_dir, rpn_name)
save_object(
dict(boxes=boxes, scores=scores, ids=ids, cfg=cfg_yaml), rpn_file
)
logger.info('Wrote RPN proposals to {}'.format(os.path.abspath(rpn_file)))
return boxes, scores, ids, rpn_file
def create_model():
"""Build the model and look for saved model checkpoints in case we can
resume from one.
"""
logger = logging.getLogger(__name__)
start_iter = 0
checkpoints = {}
output_dir = get_output_dir(training=True)
weights_file = cfg.TRAIN.WEIGHTS
if cfg.TRAIN.AUTO_RESUME:
# Check for the final model (indicates training already finished)
final_path = os.path.join(output_dir, 'model_final.pkl')
if os.path.exists(final_path):
logger.info('model_final.pkl exists; no need to train!')
return None, None, None, {'final': final_path}, output_dir, None
# Find the most recent checkpoint (highest iteration number)
files = os.listdir(output_dir)
for f in files:
iter_string = re.findall(r'(?<=model_iter)\d+(?=\.pkl)', f)
if len(iter_string) > 0:
checkpoint_iter = int(iter_string[0])
if checkpoint_iter > start_iter:
# Start one iteration immediately after the checkpoint iter
start_iter = checkpoint_iter + 1
resume_weights_file = f
if start_iter > 0:
# Override the initialization weights with the found checkpoint
weights_file = os.path.join(output_dir, resume_weights_file)
logger.info(
'========> Resuming from checkpoint {} at start iter {}'.
format(weights_file, start_iter))
tb_dir = get_tb_dir(training=True)
writer = SummaryWriter(tb_dir, tag="detectron")
logger.info('Building model: {}'.format(cfg.MODEL.TYPE))
model = model_builder.create(cfg.MODEL.TYPE, train=True, writer=writer)
params = [blob._name for blob in model.TrainableParams(gpu_id=0)]
for param in params:
if 'mem' in param:
model.AddSummaryHistogram(param)
if cfg.MEMONGER:
optimize_memory(model)
# Performs random weight initialization as defined by the model
workspace.RunNetOnce(model.param_init_net)
return model, weights_file, start_iter, checkpoints, output_dir, writer
def generate_rpn_on_range(weights_file,
dataset_name,
_proposal_file_ignored,
output_dir,
ind_range=None,
gpu_id=0):
"""Run inference on all images in a dataset or over an index range of images
in a dataset using a single GPU.
"""
assert cfg.MODEL.RPN_ONLY or cfg.MODEL.FASTER_RCNN
roidb, start_ind, end_ind, total_num_images = get_roidb(
dataset_name, ind_range)
logger.info('Output will be saved to: {:s}'.format(
os.path.abspath(output_dir)))
model = model_builder.create(cfg.MODEL.TYPE, train=False, gpu_id=gpu_id)
nu.initialize_gpu_from_weights_file(
model,
weights_file,
gpu_id=gpu_id,
)
model_builder.add_inference_inputs(model)
workspace.CreateNet(model.net)
boxes, scores, ids = generate_proposals_on_roidb(
model,
roidb,
start_ind=start_ind,
end_ind=end_ind,
total_num_images=total_num_images,
gpu_id=gpu_id,
)
cfg_yaml = yaml.dump(cfg)
if ind_range is not None:
rpn_name = 'rpn_proposals_range_%s_%s.pkl' % tuple(ind_range)
else:
rpn_name = 'rpn_proposals.pkl'
rpn_file = os.path.join(output_dir, rpn_name)
save_object(dict(boxes=boxes, scores=scores, ids=ids, cfg=cfg_yaml),
rpn_file)
logger.info('Wrote RPN proposals to {}'.format(os.path.abspath(rpn_file)))
return boxes, scores, ids, rpn_file
def test_retinanet(ind_range=None):
"""
Test RetinaNet model either on the entire dataset or the subset of dataset
specified by the index range
"""
assert cfg.RETINANET.RETINANET_ON, \
'RETINANET_ON must be set for testing RetinaNet model'
output_dir = get_output_dir(training=False)
dataset = JsonDataset(cfg.TEST.DATASET)
im_list = dataset.get_roidb()
if ind_range is not None:
start, end = ind_range
im_list = im_list[start:end]
logger.info('Testing on roidb range: {}-{}'.format(start, end))
else:
# if testing over the whole dataset, use the NUM_TEST_IMAGES setting
# the NUM_TEST_IMAGES could be over a small set of images for quick
# debugging purposes
im_list = im_list[0:cfg.TEST.NUM_TEST_IMAGES]
model = model_builder.create(cfg.MODEL.TYPE, train=False)
if cfg.TEST.WEIGHTS:
nu.initialize_from_weights_file(model,
cfg.TEST.WEIGHTS,
broadcast=False)
model_builder.add_inference_inputs(model)
workspace.CreateNet(model.net)
boxes, scores, classes, image_ids = im_list_detections(
model, im_list[0:cfg.TEST.NUM_TEST_IMAGES])
cfg_yaml = yaml.dump(cfg)
if ind_range is not None:
det_name = 'retinanet_detections_range_%s_%s.pkl' % tuple(ind_range)
else:
det_name = 'retinanet_detections.pkl'
det_file = os.path.join(output_dir, det_name)
save_object(
dict(boxes=boxes,
scores=scores,
classes=classes,
ids=image_ids,
cfg=cfg_yaml), det_file)
logger.info('Wrote detections to: {}'.format(os.path.abspath(det_file)))
return boxes, scores, classes, image_ids
def _test_std(self):
current_dir = osp.dirname(osp.realpath(__file__))
cfg_file = osp.join(current_dir, '..', 'configs', 'R-50_1x.yaml')
merge_cfg_from_file(cfg_file)
cfg.TEST.WEIGHTS = osp.join(
current_dir, '..', 'outputs', 'train',
'coco_2014_train+coco_2014_valminusminival', 'R-50_1x', 'default',
'model_final.pkl')
cfg.RETINANET.INFERENCE_TH = 0.
dataset = JsonDataset('coco_2014_minival')
roidb = dataset.get_roidb()
model = model_builder.create(cfg.MODEL.TYPE, train=False, gpu_id=0)
utils.net.initialize_gpu_from_weights_file(model,
cfg.TEST.WEIGHTS,
gpu_id=0)
model_builder.add_inference_inputs(model)
workspace.CreateNet(model.net)
workspace.CreateNet(model.conv_body_net)
num_images = len(roidb)
num_classes = cfg.MODEL.NUM_CLASSES
entry = roidb[5]
im = cv2.imread(entry['image'])
with utils.c2.NamedCudaScope(0):
cls_boxes, cls_preds, cls_probs, box_preds, anchors, im_info = im_detect_bbox(
model, im, debug=True)
cls_boxes = cls_boxes[:, :5]
im_name = osp.splitext(osp.basename(entry['image']))[0]
# utils.vis.vis_one_image(im[:, :, ::-1],
# '{:s}-std-output'.format(im_name),
# current_dir,
# cls_boxes,
# segms=None,
# keypoints=None,
# thresh=0.,
# box_alpha=0.8,
# dataset=dataset,
# show_class=False)
workspace.ResetWorkspace()
return cls_preds, cls_probs, box_preds, anchors, im_info, im, im_name, current_dir, dataset
def create_model():
"""Build the model and look for saved model checkpoints in case we can
resume from one.
"""
logger = logging.getLogger(__name__)
start_iter = 0
checkpoints = {}
output_dir = get_output_dir(cfg.TRAIN.DATASETS, training=True)
weights_file = cfg.TRAIN.WEIGHTS
if cfg.TRAIN.AUTO_RESUME:
# Check for the final model (indicates training already finished)
final_path = os.path.join(output_dir, 'model_final.pkl')
if os.path.exists(final_path):
logger.info('model_final.pkl exists; no need to train!')
return None, None, None, {'final': final_path}, output_dir
# Find the most recent checkpoint (highest iteration number)
files = os.listdir(output_dir)
for f in files:
iter_string = re.findall(r'(?<=model_iter)\d+(?=\.pkl)', f)
if len(iter_string) > 0:
checkpoint_iter = int(iter_string[0])
if checkpoint_iter > start_iter:
# Start one iteration immediately after the checkpoint iter
start_iter = checkpoint_iter + 1
resume_weights_file = f
if start_iter > 0:
# Override the initialization weights with the found checkpoint
weights_file = os.path.join(output_dir, resume_weights_file)
logger.info(
'========> Resuming from checkpoint {} at start iter {}'.
format(weights_file, start_iter)
)
logger.info('Building model: {}'.format(cfg.MODEL.TYPE))
model = model_builder.create(cfg.MODEL.TYPE, train=True)
if cfg.MEMONGER:
optimize_memory(model)
# Performs random weight initialization as defined by the model
workspace.RunNetOnce(model.param_init_net)
return model, weights_file, start_iter, checkpoints, output_dir
def create_train_model():
"""Build the model and look for saved model checkpoints in case we can
resume from one.
"""
logger = logging.getLogger(__name__)
start_iter = 0
weights = torch.load(os.path.join('checkpoint', cfg.TRAIN.WEIGHTS))
if cfg.TRAIN.AUTO_RESUME:
checkpoints = torch.load(
os.path.join('checkpoint', cfg.TRAIN.RESUME_FILE))
start_iter = checkpoints['epoch']
if start_iter > 0:
# Override the initialization weights with the found checkpoint
weights = checkpoints['net']
logger.info(
'========> Resuming from checkpoint {} at start iter {}'.
format(cfg.TRAIN.RESUME_FILE, start_iter))
logger.info('Building model: {}'.format(cfg.MODEL.TYPE))
model = create(cfg.MODEL.TYPE, cfg.MODEL.CONV_BODY, cfg.MODEL.NUM_CLASSES)
return model, weights, start_iter
def test_retinanet(ind_range=None):
"""
Test RetinaNet model either on the entire dataset or the subset of dataset
specified by the index range
"""
assert cfg.RETINANET.RETINANET_ON, \
'RETINANET_ON must be set for testing RetinaNet model'
output_dir = get_output_dir(training=False)
dataset = JsonDataset(cfg.TEST.DATASET)
im_list = dataset.get_roidb()
if ind_range is not None:
start, end = ind_range
im_list = im_list[start:end]
logger.info('Testing on roidb range: {}-{}'.format(start, end))
else:
# if testing over the whole dataset, use the NUM_TEST_IMAGES setting
# the NUM_TEST_IMAGES could be over a small set of images for quick
# debugging purposes
im_list = im_list[0:cfg.TEST.NUM_TEST_IMAGES]
model = model_builder.create(cfg.MODEL.TYPE, train=False)
if cfg.TEST.WEIGHTS:
nu.initialize_from_weights_file(
model, cfg.TEST.WEIGHTS, broadcast=False
)
model_builder.add_inference_inputs(model)
workspace.CreateNet(model.net)
boxes, scores, classes, image_ids = im_list_detections(
model, im_list[0:cfg.TEST.NUM_TEST_IMAGES])
cfg_yaml = yaml.dump(cfg)
if ind_range is not None:
det_name = 'retinanet_detections_range_%s_%s.pkl' % tuple(ind_range)
else:
det_name = 'retinanet_detections.pkl'
det_file = os.path.join(output_dir, det_name)
save_object(
dict(boxes=boxes, scores=scores, classes=classes, ids=image_ids, cfg=cfg_yaml),
det_file)
logger.info('Wrote detections to: {}'.format(os.path.abspath(det_file)))
return boxes, scores, classes, image_ids
def test_model():
"""Model testing loop."""
logger = logging.getLogger(__name__)
model = create(cfg.MODEL.TYPE, cfg.MODEL.CONV_BODY, cfg.MODEL.NUM_CLASSES)
checkpoint = torch.load(os.path.join('checkpoint', cfg.TEST.WEIGHTS))
model.load_state_dict(checkpoint['net'])
if not torch.cuda.is_available():
logger.info('cuda not find')
sys.exit(1)
model = torch.nn.DataParallel(model,
device_ids=range(torch.cuda.device_count()))
model.cuda()
model.eval()
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
img_dir = os.path.join(cfg.TEST.DATASETS_DIR, cfg.TEST.DATASETS[0],
'JPEGImages')
img_list = os.path.join(cfg.TEST.DATASETS_DIR, cfg.TEST.DATASETS[0],
'ImageSets', 'Main', cfg.TEST.DATASETS[1])
with open(img_list, 'r') as lst:
img_list = lst.readlines()
img_nums = len(img_list)
test_scales = cfg.TEST.SCALES
dic = {}
for i in range(20):
dic[str(i)] = []
for im in range(img_nums):
if im % 100 == 0:
logger.info('{} imgs were processed, total {}'.format(
im, img_nums))
img = Image.open(os.path.join(img_dir, img_list[im].strip() + '.jpg'))
img_size = img.size
img = img.resize(test_scales)
# For visualization
filename = os.path.join("visual_results",
img_list[im].strip() + '.jpg')
x = transform(img)
x = x.unsqueeze(0)
x = torch.autograd.Variable(x)
loc_preds, cls_preds = model(x)
loc_preds = loc_preds.data.squeeze().type(torch.FloatTensor)
cls_preds = cls_preds.data.squeeze().type(torch.FloatTensor)
encoder = DataEncoder(test_scales)
boxes, labels, sco, is_found = encoder.decode(loc_preds, cls_preds,
test_scales)
if is_found:
img, boxes = resize(img, boxes, img_size)
boxes = boxes.ceil()
xmin = boxes[:, 0].clamp(min=1)
ymin = boxes[:, 1].clamp(min=1)
xmax = boxes[:, 2].clamp(max=img_size[0] - 1)
ymax = boxes[:, 3].clamp(max=img_size[1] - 1)
nums = len(boxes)
for i in range(nums):
dic[str(labels[i].item())].append([
img_list[im].strip(), sco[i].item(), xmin[i].item(),
ymin[i].item(), xmax[i].item(), ymax[i].item()
])
draw = ImageDraw.Draw(img)
font = ImageFont.truetype(
b'/usr/share/fonts/truetype/ancient-scripts/Symbola_hint.ttf',
20)
count = 0
for box in boxes:
draw.rectangle(list(box), outline='red', width=5)
draw.text((box[0], box[1] - 20),
str(category[labels[count]]),
font=font,
fill=(255, 0, 0, 255))
count = count + 1
img.save(filename)
for key in dic.keys():
logger.info('category id: {}, category name: {}'.format(
key, category[int(key)]))
file_name = cfg.TEST.OUTPUT_DIR + 'comp4_det_test_' + category[int(
key)] + '.txt'
with open(file_name, 'w') as comp4:
nums = len(dic[key])
for i in range(nums):
img, cls_preds, xmin, ymin, xmax, ymax = dic[key][i]
if cls_preds > 0.5:
cls_preds = '%.6f' % cls_preds
loc_preds = '%.6f %.6f %.6f %.6f' % (xmin, ymin, xmax,
ymax)
rlt = '{} {} {}\n'.format(img, cls_preds, loc_preds)
comp4.write(rlt)
def test_model():
"""Model testing loop."""
logger = logging.getLogger(__name__)
colors = np.random.randint(0, 255, size=(1, 3), dtype="uint8")
model = create(cfg.MODEL.TYPE, cfg.MODEL.CONV_BODY, cfg.MODEL.NUM_CLASSES)
checkpoint = torch.load(os.path.join('checkpoint', cfg.TEST.WEIGHTS))
model.load_state_dict(checkpoint['net'])
if not torch.cuda.is_available():
logger.info('cuda not find')
sys.exit(1)
#model = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count()))
model.cuda()
#model.cpu()
model.eval()
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
img_dir = os.path.join("/home/mia_dev/xeroblade2/dataset/train/img/")
g = os.walk(r"/home/mia_dev/xeroblade2/dataset/test/2021-03-22_21-49-34")
img_list = []
for path, dir_list, file_list in g:
for file_name in file_list:
img_list.append(file_name)
img_nums = len(img_list)
test_scales = cfg.TEST.SCALES
dic = {}
for i in range(20):
dic[str(i)] = []
frame_array = []
cap = cv2.VideoCapture(
"/home/mia_dev/xeroblade2/dataset/test/2021-03-22_21-47-43.mp4")
colors = np.random.randint(0, 255, size=(1, 3), dtype="uint8")
a = []
time_begin = time.time()
NUM = cap.get(cv2.CAP_PROP_FRAME_COUNT)
#NUM=0
frame_array = []
count = 0
test_scales = (1280, 960)
while cap.isOpened():
#img = Image.open(os.path.join("/home/mia_dev/xeroblade2/dataset/test/2021-03-22_21-49-34/", img_list[im].strip()))
ret, img = cap.read()
if ret is False:
break
#print(os.path.join("/home/mia_dev/xeroblade2/dataset/test/2021-03-22_21-49-34/", img_list[im].strip()))
img_size = img.shape
#img = cv2.resize(img, (1280, 960), interpolation=cv2.INTER_CUBIC)
img = cv2.resize(img, test_scales)
RGBimg = changeBGR2RGB(img)
x = transform(RGBimg)
x = x.cuda()
x = x.unsqueeze(0)
x = torch.autograd.Variable(x)
loc_preds, cls_preds = model(x)
loc_preds = loc_preds.data.squeeze().type(torch.FloatTensor)
cls_preds = cls_preds.data.squeeze().type(torch.FloatTensor)
encoder = DataEncoder(test_scales)
boxes, labels, sco, is_found = encoder.decode(loc_preds, cls_preds,
test_scales)
if is_found:
#img, boxes = resize(img, boxes, img_size)
boxes = boxes.ceil()
xmin = boxes[:, 0].clamp(min=1)
ymin = boxes[:, 1].clamp(min=1)
xmax = boxes[:, 2].clamp(max=img_size[0] - 1)
ymax = boxes[:, 3].clamp(max=img_size[1] - 1)
color = [int(c) for c in colors[0]]
nums = len(boxes)
print(nums)
#for i in range(nums) :
#dic[str(labels[i].item())].append([img_list[im].strip(), sco[i].item(), xmin[i].item(), ymin[i].item(), xmax[i].item(), ymax[i].item()])
for i in range(nums):
if float(sco[i]) > 0.5:
box_w = xmax[i] - xmin[i]
box_h = ymax[i] - ymin[i]
print(box_w, box_h)
box_w = int(box_w)
box_h = int(box_h)
# print(cls_conf)
x1 = int(xmin[i])
x2 = int(xmax[i])
y1 = int(ymin[i])
y2 = int(ymax[i])
img = cv2.rectangle(img, (x1, y1 + box_h), (x2, y1), color,
2)
cv2.putText(img, 'nopon', (x1, y1),
cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
cv2.putText(img, str("%.2f" % float(sco[i])),
(x2, y2 - box_h), cv2.FONT_HERSHEY_SIMPLEX,
0.5, color, 2)
vidframe = changeRGB2BGR(RGBimg)
#cv2.imshow('frame', vidframe)
frame_array.append(vidframe)
#if cv2.waitKey(25) & 0xFF == ord('q'):
#break
pathOut = '/home/mia_dev/xeroblade2/RetinaNet-Pytorch-master/2021-03-22_21-47-43.mp4'
fps = 60
out = cv2.VideoWriter(pathOut, cv2.VideoWriter_fourcc(*'DIVX'), fps,
test_scales)
print(len(frame_array))
for i in range(len(frame_array)):
# writing to a image array
out.write(frame_array[i])
out.release()
cap.release()
cv2.destroyAllWindows()
"""
def test_model():
"""Model testing loop."""
logger = logging.getLogger(__name__)
colors = np.random.randint(0, 255, size=(1, 3), dtype="uint8")
model = create(cfg.MODEL.TYPE, cfg.MODEL.CONV_BODY, cfg.MODEL.NUM_CLASSES)
checkpoint = torch.load(os.path.join('checkpoint', cfg.TEST.WEIGHTS))
model.load_state_dict(checkpoint['net'])
if not torch.cuda.is_available():
logger.info('cuda not find')
sys.exit(1)
#model = torch.nn.DataParallel(model, device_ids=range(torch.cuda.device_count()))
model.cuda()
#model.cpu()
model.eval()
transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.485,0.456,0.406), (0.229,0.224,0.225))])
img_dir = os.path.join("/home/mia_dev/xeroblade2/dataset/train/img/")
g = os.walk(r"/home/mia_dev/xeroblade2/dataset/test/2021-03-22_21-49-34")
img_list=[]
for path, dir_list, file_list in g:
for file_name in file_list:
img_list.append(file_name)
img_nums = len(img_list)
test_scales = cfg.TEST.SCALES
dic = {}
for i in range(20) :
dic[str(i)] = []
for im in range(img_nums):
if im % 100 == 0 : logger.info('{} imgs were processed, total {}'. format(im, img_nums))
img = Image.open(os.path.join("/home/mia_dev/xeroblade2/dataset/test/2021-03-22_21-49-34/", img_list[im].strip()))
print(os.path.join("/home/mia_dev/xeroblade2/dataset/test/2021-03-22_21-49-34/", img_list[im].strip()))
img_size = img.size
img = img.resize(test_scales)
x = transform(img)
x=x.cuda()
x = x.unsqueeze(0)
x = torch.autograd.Variable(x)
loc_preds, cls_preds = model(x)
loc_preds = loc_preds.data.squeeze().type(torch.FloatTensor)
cls_preds = cls_preds.data.squeeze().type(torch.FloatTensor)
encoder = DataEncoder(test_scales)
boxes, labels, sco, is_found = encoder.decode(loc_preds, cls_preds, test_scales)
if is_found :
img, boxes = resize(img, boxes, img_size)
boxes = boxes.ceil()
xmin = boxes[:, 0].clamp(min = 1)
ymin = boxes[:, 1].clamp(min = 1)
xmax = boxes[:, 2].clamp(max = img_size[0] - 1)
ymax = boxes[:, 3].clamp(max = img_size[1] - 1)
nums = len(boxes)
print(nums)
for i in range(nums) :
dic[str(labels[i].item())].append([img_list[im].strip(), sco[i].item(), xmin[i].item(), ymin[i].item(), xmax[i].item(), ymax[i].item()])
temp=''
for key in dic.keys() :
#logger.info('category id: {}, category name: {}'. format(key, category[int(key)]))
#file_name = cfg.TEST.OUTPUT_DIR + 'comp4_det_test_'+category[int(key)]+'.txt'
#with open(file_name, 'w') as comp4 :
nums = len(dic[key])
for i in range(nums) :
img, cls_preds, xmin, ymin, xmax, ymax = dic[key][i]
if temp!=img:
temp=img
imgs = cv2.imread("/home/mia_dev/xeroblade2/dataset/test/2021-03-22_21-49-34/" + img)
else:
imgs=imgs
print(cls_preds)
if cls_preds > 0 :
cls_preds = '%.6f' % cls_preds
loc_preds = '%.6f %.6f %.6f %.6f' % (xmin, ymin, xmax, ymax)
rlt = '{} {} {}\n'.format(img, cls_preds, loc_preds)
#comp4.write(rlt)
box_w = xmax - xmin
box_h = ymax - ymin
color = [int(c) for c in colors[0]]
print(box_w, box_h)
box_w=int(box_w)
box_h=int(box_h)
# print(cls_conf)
x1=int(xmin)
x2=int(xmax)
y1=int(ymin)
y2=int(ymax)
imgs = cv2.rectangle(imgs, (x1, y1 + box_h), (x2, y1), color, 2)
cv2.putText(imgs, 'nopon', (x1, y1), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
cv2.putText(imgs, str("%.2f" % float(cls_preds)), (x2, y2 - box_h), cv2.FONT_HERSHEY_SIMPLEX, 0.5,color, 2)
print("/home/mia_dev/xeroblade2/dataset/result/retinanet/"+img.split('/')[-1])
cv2.imwrite("/home/mia_dev/xeroblade2/dataset/result/retinanet/"+img.split('/')[-1],imgs)
def test_model():
"""Model testing loop."""
logger = logging.getLogger(__name__)
model = create(cfg.MODEL.TYPE, cfg.MODEL.CONV_BODY, cfg.MODEL.NUM_CLASSES)
checkpoint = torch.load(os.path.join('checkpoint', cfg.TEST.WEIGHTS))
model.load_state_dict(checkpoint['net'])
if not torch.cuda.is_available():
print("DUBUGGER that cuda is not found")
logger.info('cuda not find')
sys.exit(1)
print("Total cuda devices = ", torch.cuda.device_count())
model = torch.nn.DataParallel(model,
device_ids=range(torch.cuda.device_count()))
model.cuda()
model.eval()
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
print("DIRS = ", cfg.TEST.DATASETS_DIR)
print("Another DIRS = ", cfg.TEST.DATASETS[0])
img_dir = os.path.join(cfg.TEST.DATASETS_DIR, cfg.TEST.DATASETS[0],
'JPEGImages')
img_list = os.path.join(cfg.TEST.DATASETS_DIR, cfg.TEST.DATASETS[0],
'ImageSets', 'Main', cfg.TEST.DATASETS[1])
with open(img_list, 'r') as lst:
img_list = lst.readlines()
img_nums = len(img_list)
test_scales = cfg.TEST.SCALES
dic = {}
for i in range(20):
dic[str(i)] = []
for im in range(img_nums):
if im % 100 == 0:
logger.info('{} imgs were processed, total {}'.format(
im, img_nums))
img = Image.open(os.path.join(img_dir, img_list[im].strip() + '.jpg'))
img_size = img.size
img = img.resize(test_scales)
x = transform(img)
x = x.unsqueeze(0)
x = torch.autograd.Variable(x)
loc_preds, cls_preds = model(x)
loc_preds = loc_preds.data.squeeze().type(torch.FloatTensor)
cls_preds = cls_preds.data.squeeze().type(torch.FloatTensor)
encoder = DataEncoder(test_scales)
boxes, labels, sco, is_found = encoder.decode(loc_preds, cls_preds,
test_scales)
if is_found:
img, boxes = resize(img, boxes, img_size)
boxes = boxes.ceil()
xmin = boxes[:, 0].clamp(min=1)
ymin = boxes[:, 1].clamp(min=1)
xmax = boxes[:, 2].clamp(max=img_size[0] - 1)
ymax = boxes[:, 3].clamp(max=img_size[1] - 1)
nums = len(boxes)
for i in range(nums):
dic[str(labels[i].item())].append([
img_list[im].strip(), sco[i].item(), xmin[i].item(),
ymin[i].item(), xmax[i].item(), ymax[i].item()
])
for key in dic.keys():
logger.info('category id: {}, category name: {}'.format(
key, category[int(key)]))
file_name = cfg.TEST.OUTPUT_DIR + 'comp4_det_test_' + category[int(
key)] + '.txt'
with open(file_name, 'w') as comp4:
nums = len(dic[key])
for i in range(nums):
img, cls_preds, xmin, ymin, xmax, ymax = dic[key][i]
if cls_preds > 0.5:
cls_preds = '%.6f' % cls_preds
loc_preds = '%.6f %.6f %.6f %.6f' % (xmin, ymin, xmax,
ymax)
rlt = '{} {} {}\n'.format(img, cls_preds, loc_preds)
comp4.write(rlt)
