def main(args):
if args.dataset:
dataset = get_dataset(args.dataset, False, False, None)
classes_num = dataset.classes_num
else:
classes_num = args.classes_num
net = locate(args.model)(classes_num, fc_detection_head=False)
load_checkpoint(net, args.ckpt)
if args.show_flops:
net = add_flops_counting_methods(net)
net.reset_flops_count()
net.start_flops_count()
printable_graph = onnx_export(net, args.input_size, args.output_file, check=args.check, verbose=args.verbose)
if args.verbose:
logging.info(printable_graph)
if args.show_flops:
net.stop_flops_count()
logging.info('Computational complexity: {}'.format(flops_to_string(net.compute_average_flops_cost())))
if args.verbose:
print_model_with_flops(net)
def main(args):
args.output_dir = osp.abspath(args.output_dir)
downloaders = dict(web=download_file_from_web,
google_drive=download_file_from_google_drive)
with open(args.model_zoo, 'rt') as f:
model_zoo = yaml.load(f)['models']
logging.info('Models to be fetched:')
for target in model_zoo:
logging.info('\t{} ({}, {})'.format(colored(target['name'], 'green'),
target['dataset'],
target['framework']))
logging.info('')
for target in model_zoo:
output_file = osp.join(args.output_dir, 'raw', target['dst_file'])
mkdir_for_file(output_file)
logging.info('Fetching {} ({}, {})'.format(
colored(target['name'], 'green'), target['dataset'],
target['framework']))
try:
# Download weights.
if target['storage_type'] in downloaders:
downloaders[target['storage_type']](target['url'], output_file)
else:
logging.warning(
'No downloaders available for storage {}'.format(
target['storage_type']))
continue
# Convert weights.
logging.info('Downloaded to {}'.format(output_file))
if target.get('weights_converter', None):
logging.info('Converting weights...')
converter = locate(target['weights_converter'])
if converter is None:
logging.warning('Invalid weights converter {}'.format(
target['weights_converter']))
continue
output_converted_file = osp.join(args.output_dir, 'converted',
target['dst_file'])
output_converted_file = osp.splitext(
output_converted_file)[0] + '.pth'
mkdir_for_file(output_converted_file)
try:
converter(output_file, output_converted_file)
logging.info('Converted weights file saved to {}'.format(
output_converted_file))
except Exception as ex:
logging.warning('Failed to convert weights.')
logging.warning(ex)
continue
if target.get('convert_to_ir', False):
# Convert to ONNX.
logging.info('Exporting to ONNX...')
output_onnx_file = osp.join(args.output_dir, 'onnx',
target['dst_file'])
output_onnx_file = osp.splitext(
output_onnx_file)[0] + '.onnx'
mkdir_for_file(output_onnx_file)
net = locate(target['model'])(81, fc_detection_head=False)
load_checkpoint(net, output_converted_file, verbose=False)
onnx_export(net, target['input_size'], output_onnx_file)
logging.info(
'ONNX file is saved to {}'.format(output_onnx_file))
# Convert to IR.
logging.info('Converting to IR...')
output_ir_dir = osp.join(args.output_dir, 'ir',
target['dst_file'])
mkdir_for_file(output_ir_dir)
output_ir_dir = osp.dirname(output_ir_dir)
status = call([
args.model_optimizer, '--framework', 'onnx',
'--input_model', output_onnx_file, '--output_dir',
output_ir_dir, '--input', 'im_data,im_info',
'--output', 'boxes,scores,classes,batch_ids,raw_masks',
'--mean_values', 'im_data{},im_info[0,0,0]'.format(
str(target['mean_pixel']).replace(' ', ''))
])
if status:
logging.warning('Failed to convert model to IR.')
else:
logging.info(
'IR files saved to {}'.format(output_ir_dir))
except Exception as ex:
logging.warning(repr(ex))
def main(args):
transforms = Compose([
Resize(max_size=args.fit_max_image_size,
window_size=args.fit_window_size,
size=args.size),
ToTensor(),
Normalize(mean=args.mean_pixel, std=[1., 1., 1.], rgb=args.rgb),
])
dataset = get_dataset(args.dataset, False, False, transforms)
logging.info(dataset)
batch_size = 1
logging.info('Using {} backend'.format(args.backend))
logging.info('Loading network...')
if args.backend == 'pytorch':
net = locate(args.pytorch_model_class)(dataset.classes_num)
net.eval()
load_checkpoint(net, args.checkpoint_file_path)
if torch.cuda.is_available():
net = net.cuda()
net = add_flops_counting_methods(net)
net.reset_flops_count()
net.start_flops_count()
elif args.backend == 'openvino':
net = MaskRCNNOpenVINO(
args.openvino_model_path,
args.checkpoint_file_path,
device=args.device,
plugin_dir=args.plugin_dir,
cpu_extension_lib_path=args.cpu_extension,
collect_perf_counters=args.show_performance_counters)
else:
raise ValueError('Unknown backend "{}"'.format(args.backend))
viz = Visualizer(dataset.classes,
confidence_threshold=args.prob_threshold,
show_boxes=args.show_boxes,
show_scores=args.show_scores)
inference_timer = Timer(cuda_sync=True, warmup=1)
timer = Timer(cuda_sync=False, warmup=1)
timer.tic()
logging.info('Configuring data source...')
if args.video:
try:
args.video = int(args.video)
except ValueError:
pass
demo_dataset = VideoDataset(args.video,
labels=dataset.classes,
transforms=transforms)
num_workers = 0
tracker = StaticIOUTracker()
else:
demo_dataset = ImagesDataset(args.images,
labels=dataset.classes,
transforms=transforms)
num_workers = 1
tracker = None
data_loader = torch.utils.data.DataLoader(demo_dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=False,
collate_fn=collate)
logging.info('Processing data...')
frames_num = len(demo_dataset)
for data_batch in tqdm(
iter(data_loader),
total=frames_num if frames_num != sys.maxsize else 0):
im_data = data_batch['im_data']
im_info = data_batch['im_info']
if torch.cuda.is_available():
im_data = [i.cuda() for i in im_data]
im_info = [i.cuda() for i in im_info]
with torch.no_grad(), inference_timer:
boxes, classes, scores, _, masks = net(im_data, im_info)
meta = data_batch['meta'][0]
scores, classes, boxes, masks = postprocess(
scores,
classes,
boxes,
masks,
im_h=meta['original_size'][0],
im_w=meta['original_size'][1],
im_scale_y=meta['processed_size'][0] / meta['original_size'][0],
im_scale_x=meta['processed_size'][1] / meta['original_size'][1],
full_image_masks=True,
encode_masks=False,
confidence_threshold=args.prob_threshold)
masks_ids = tracker(masks, classes) if tracker is not None else None
image = data_batch['original_image'][0]
visualization = viz(image,
boxes,
classes,
scores,
segms=masks,
ids=masks_ids)
fps = 1 / timer.toc()
if args.show_fps:
visualization = cv2.putText(visualization,
'FPS: {:>2.2f}'.format(fps), (30, 30),
cv2.FONT_HERSHEY_SIMPLEX, 1,
(0, 0, 255), 2)
cv2.imshow('result', visualization)
key = cv2.waitKey(args.delay)
if key == 27:
break
timer.tic()
if inference_timer.average_time > 0:
logging.info('Average inference FPS: {:3.2f}'.format(
1 / inference_timer.average_time))
if args.backend == 'pytorch':
net.stop_flops_count()
if args.show_flops:
logging.info('Average FLOPs: {}'.format(
flops_to_string(net.compute_average_flops_cost())))
if args.show_layers_flops:
logging.info('Thorough computational complexity statistics:')
print_model_with_flops(net)
if torch.cuda.is_available():
logging.info('GPU memory footprint:')
logging.info('\tMax allocated: {:.2f} MiB'.format(
torch.cuda.max_memory_allocated() / 1024**2))
logging.info('\tMax cached: {:.2f} MiB'.format(
torch.cuda.max_memory_cached() / 1024**2))
else:
if args.show_performance_counters:
net.print_performance_counters()
cv2.destroyAllWindows()
del net
def main():
""" Does export to onnx. """
args = parse_args()
with open(args.model) as file:
config = json.load(file)
text_spotter = make_text_detector(**config['model'])(
2, fc_detection_head=False, shape=args.input_size)
load_checkpoint(text_spotter, args.ckpt)
text_spotter.export_mode = True
net = text_spotter
if args.show_flops:
net = add_flops_counting_methods(net)
net.reset_flops_count()
net.start_flops_count()
# Export of text detection part (Mask-RCNN subgraph).
printable_graph = onnx_export(net,
args.input_size,
args.output_folder,
check=args.check,
verbose=args.verbose)
if args.verbose:
logging.info(printable_graph)
if args.show_flops:
net.stop_flops_count()
logging.info(
'Computational complexity of text detection part: {}'.format(
flops_to_string(net.compute_average_flops_cost())))
if args.verbose:
print_model_with_flops(net)
# Export of text recognition encoder
net = text_spotter.text_recogn_head.encoder
if args.show_flops:
net = add_flops_counting_methods(net)
net.reset_flops_count()
net.start_flops_count()
printable_graph = export_to_onnx_text_recognition_encoder(
net, text_spotter.text_recogn_head.input_feature_size,
args.output_folder)
if args.verbose:
logging.info(printable_graph)
if args.show_flops:
net.stop_flops_count()
logging.info(
'Computational complexity of text recognition encoder part: {}'.
format(flops_to_string(net.compute_average_flops_cost())))
if args.verbose:
print_model_with_flops(net)
# Export of text recognition decoder
net = text_spotter.text_recogn_head.decoder
if args.show_flops:
net = add_flops_counting_methods(net)
net.reset_flops_count()
net.start_flops_count()
printable_graph = export_to_onnx_text_recognition_decoder(
net, text_spotter.text_recogn_head.input_feature_size,
args.output_folder)
if args.verbose:
logging.info(printable_graph)
if args.show_flops:
net.stop_flops_count()
logging.info(
'Computational complexity of text recognition decoder part: {}'.
format(flops_to_string(net.compute_average_flops_cost())))
if args.verbose:
print_model_with_flops(net)
def __init__(self):
super().__init__()
self.identifier = 'instance-segmentation-security-0050'
self.description = 'Training of instance-segmentation-security-0050'
self.root_directory = osp.join(osp.dirname(osp.abspath(__file__)), '..')
self.run_directory = self.create_run_directory(osp.join(self.root_directory, 'outputs'))
setup_logging(file_path=osp.join(self.run_directory, 'log.txt'))
logger.info('Running {}'.format(self.identifier))
logger.info(self.description)
logger.info('Working directory "{}"'.format(self.run_directory))
self.batch_size = 32
self.virtual_iter_size = 1
# Training dataset.
training_transforms = Compose(
[
RandomResize(mode='size', heights=(416, 448, 480, 512, 544), widths=(416, 448, 480, 512, 544)),
RandomHorizontalFlip(prob=0.5),
ToTensor(),
Normalize(mean=[102.9801, 115.9465, 122.7717], std=[1., 1., 1.], rgb=False),
],
)
training_dataset_name = 'coco_2017_train'
logger.info('Training dataset {}'.format(training_dataset_name))
training_dataset = get_dataset(training_dataset_name, True, True, training_transforms)
logger.info(training_dataset)
self.training_data_loader = torch.utils.data.DataLoader(
training_dataset, batch_size=self.batch_size, num_workers=0,
shuffle=True, drop_last=True, collate_fn=collate
)
# Validation datasets.
validation_transforms = Compose(
[
Resize(size=[480, 480]),
ToTensor(),
Normalize(mean=[102.9801, 115.9465, 122.7717], std=[1., 1., 1.], rgb=False),
]
)
validation_datasets = []
validation_dataset_name = 'coco_2017_val'
logger.info('Validation dataset #{}: {}'.format(len(validation_datasets) + 1, validation_dataset_name))
validation_datasets.append(get_dataset(validation_dataset_name, False, False, validation_transforms))
logger.info(validation_datasets[-1])
self.validation_data_loaders = []
for validation_dataset in validation_datasets:
self.validation_data_loaders.append(torch.utils.data.DataLoader(
validation_dataset,
batch_size=1, num_workers=8,
shuffle=False, drop_last=False, collate_fn=collate)
)
self.validate_every = 10000
for validation_dataset in validation_datasets:
assert training_dataset.classes_num == validation_dataset.classes_num
# Model and optimizer.
logger.info('Model:')
self.model = Model(training_dataset.classes_num)
logger.info(self.model)
self.training_iterations_num = 270000
lr_scheduler_milestones = [220000, 250000]
base_lr = 0.02
weight_decay = 0.0001
logger.info('Optimizer:')
self.optimizer = torch.optim.SGD(self.setup_optimizer(self.model, base_lr, weight_decay),
lr=base_lr, weight_decay=weight_decay, momentum=0.9)
logger.info(self.optimizer)
logger.info('Learning Rate scheduler:')
self.lr_scheduler = MultiStepLRWithWarmUp(
self.optimizer,
milestones=lr_scheduler_milestones,
warmup_iters=1000,
warmup_method='linear',
warmup_factor_base=0.333,
gamma=0.1,
last_epoch=0
)
logger.info(self.lr_scheduler)
self.start_step = 0
checkpoint_file_path = osp.join(self.root_directory, 'data', 'pretrained_models',
'converted', 'imagenet', 'detectron', 'resnet50.pth')
if not osp.exists(checkpoint_file_path):
raise IOError('Initial checkpoint file "{}" does not exist. '
'Please fetch pre-trained backbone networks using '
'tools/download_pretrained_weights.py script first.'.format(checkpoint_file_path))
logger.info('Loading weights from "{}"'.format(checkpoint_file_path))
load_checkpoint(self.model.backbone, checkpoint_file_path)
# Loggers and misc. stuff.
self.loggers = [TextLogger(logger),
TensorboardLogger(self.run_directory)]
self.log_every = 50
self.checkpoint_every = 10000
def main(args):
""" Tests text spotter. """
transforms = Compose(
[
Resize(size=args.size),
ToTensor(),
Normalize(mean=args.mean_pixel, std=args.std_pixel, rgb=args.rgb),
]
)
dataset = get_dataset(args.dataset, False, False, transforms,
alphabet_decoder=AlphabetDecoder())
logging.info(dataset)
num_workers = args.num_workers
inference_timer = Timer()
logging.info('Using {} backend'.format(args.backend))
logging.info('Loading network...')
batch_size = 1
if args.backend == 'pytorch':
with open(args.pytorch_model_class) as file:
config = json.load(file)
net = make_text_detector(**config['model'])(dataset.classes_num,
force_max_output_size=False, shape=args.size)
net.eval()
load_checkpoint(net, args.checkpoint_file_path)
if torch.cuda.is_available():
net = net.cuda()
net = add_flops_counting_methods(net)
net.reset_flops_count()
net.start_flops_count()
if torch.cuda.is_available():
torch.backends.cudnn.deterministic = True
net = net.cuda()
net = ShallowDataParallel(net)
elif args.backend == 'openvino':
net = TextMaskRCNNOpenVINO(args.openvino_detector_model_path,
args.openvino_encoder_model_path,
args.openvino_decoder_model_path,
collect_perf_counters=args.show_performance_counters)
else:
raise ValueError('Unknown backend "{}"'.format(args.backend))
logging.info('Using batch size {}'.format(batch_size))
logging.info('Number of prefetching processes {}'.format(num_workers))
data_loader = torch.utils.data.DataLoader(
dataset,
batch_size=batch_size,
num_workers=num_workers,
shuffle=False,
collate_fn=collate
)
logging.info('Processing dataset...')
boxes_all = []
masks_all = []
classes_all = []
scores_all = []
text_probs_all = []
for data_batch in tqdm(iter(data_loader)):
batch_meta = data_batch['meta']
actual_batch_size = len(batch_meta)
with torch.no_grad(), inference_timer:
boxes, classes, scores, batch_ids, masks, text_probs = net(**data_batch)
im_heights = [meta['original_size'][0] for meta in batch_meta]
im_widths = [meta['original_size'][1] for meta in batch_meta]
im_scale_y = [meta['processed_size'][0] / meta['original_size'][0] for meta in batch_meta]
im_scale_x = [meta['processed_size'][1] / meta['original_size'][1] for meta in batch_meta]
scores, classes, boxes, masks, text_probs = postprocess_batch(
batch_ids, scores, classes, boxes, masks, text_probs, actual_batch_size,
im_h=im_heights,
im_w=im_widths,
im_scale_y=im_scale_y,
im_scale_x=im_scale_x,
full_image_masks=True, encode_masks=True,
confidence_threshold=args.prob_threshold)
boxes_all.extend(boxes)
masks_all.extend(masks)
classes_all.extend(classes)
scores_all.extend(scores)
text_probs_all.extend(text_probs)
try:
del data_loader
except ConnectionResetError:
pass
logging.info('Evaluating results...')
evaluation_results = dataset.evaluate(scores_all, classes_all, boxes_all, masks_all,
text_probs_all, dump='dump', visualize=args.visualize)
logging.info(evaluation_results)
logging.info('Average inference time {}'.format(inference_timer.average_time))
if args.backend == 'pytorch':
if torch.cuda.is_available():
net = net.module
net.stop_flops_count()
if args.show_flops:
logging.info(
'Average FLOPs: {}'.format(flops_to_string(net.compute_average_flops_cost())))
if args.show_layers_flops:
logging.info('Thorough computational complexity statistics:')
print_model_with_flops(net)
else:
if args.show_performance_counters:
net.print_performance_counters()
del net
def load_checkpoint(net,
optimizer,
load_ckpt=None,
load_backbone=None,
resume=False):
start_step = 0
if load_ckpt:
logging.info('loading checkpoint "{}"'.format(load_ckpt))
checkpoint = torch.load(load_ckpt,
map_location=lambda storage, loc: storage)
weight_utils.load_rcnn_ckpt(net, checkpoint['model'])
if resume:
start_step = checkpoint['step']
optimizer.load_state_dict(checkpoint['optimizer'])
corrected_matcher = checkpoint.get('optimizer_corrected', {})
if len(corrected_matcher) > 0:
for group in optimizer.param_groups:
for param_original in group['params']:
if 'momentum_buffer' in optimizer.state[
param_original]:
for param_loaded, buffer in corrected_matcher.items(
):
shapes_are_equal = np.array_equal(
list(param_original.shape),
list(param_loaded.shape))
if shapes_are_equal and torch.all(
torch.eq(
param_original.data,
param_loaded.data.cuda())):
optimizer.state[param_original][
'momentum_buffer'] = buffer.cuda()
break
else:
# If a checkpoint does not have additional dictionary with matched
# parameters and its buffers, just match them by shapes
if len(
optimizer.state
) > 0: # It means that a checkpoint has momentum_buffer
used_buffers = {}
copy_buf = None
for p in optimizer.state.keys():
used_buffers[p] = False
for group in optimizer.param_groups:
for param in group['params']:
for p, buffer in optimizer.state.items():
if 'momentum_buffer' not in buffer:
continue
if np.array_equal(list(param.shape), list(buffer['momentum_buffer'].shape)) and not \
used_buffers[p]:
copy_buf = optimizer.state[param][
'momentum_buffer'].cuda()
optimizer.state[param][
'momentum_buffer'] = buffer[
'momentum_buffer'].cuda()
optimizer.state[p][
'momentum_buffer'] = copy_buf.cuda(
)
used_buffers[param] = True
del used_buffers
del copy_buf
logging.info('Resume training from {} step'.format(start_step))
del checkpoint
torch.cuda.empty_cache()
if load_backbone:
logging.info(
'loading backbone weights from "{}"'.format(load_backbone))
assert hasattr(net, 'backbone')
weight_utils.load_checkpoint(net.backbone, load_backbone)
return start_step
def __init__(self, work_dir, config):
super().__init__()
self.identifier = config['identifier']
self.description = config['description']
self.root_directory = work_dir if work_dir else osp.join(
osp.dirname(osp.abspath(__file__)), '..')
self.run_directory = self.create_run_directory(
osp.join(self.root_directory, 'models'))
setup_logging(file_path=osp.join(self.run_directory, 'log.txt'))
logger.info('Running {}'.format(self.identifier))
logger.info(self.description)
logger.info('Working directory "{}"'.format(self.run_directory))
self.batch_size = config['training_details']['batch_size']
self.virtual_iter_size = config['training_details'][
'virtual_iter_size']
model_class = make_text_detector(**config['model'])
alphabet_decoder = AlphabetDecoder()
# Training dataset.
training_transforms = Compose([
getattr(sys.modules[__name__], k)(**v)
for k, v in config['training_transforms'].items()
] + [AlphabetDecodeTransform(alphabet_decoder)])
training_dataset_name = config['training_dataset_name']
logger.info('Training dataset {}'.format(training_dataset_name))
training_dataset = get_dataset(training_dataset_name,
True,
True,
training_transforms,
alphabet_decoder=alphabet_decoder,
remove_images_without_text=True)
logger.info(training_dataset)
self.training_data_loader = torch.utils.data.DataLoader(
training_dataset,
batch_size=self.batch_size,
num_workers=0,
shuffle=True,
drop_last=True,
collate_fn=collate)
# Validation datasets.
validation_transforms = Compose([
getattr(sys.modules[__name__], k)(**v)
for k, v in config['validation_transforms'].items()
])
self.confidence_threshold = config['validation_confidence_threshold']
validation_datasets = []
validation_dataset_name = config['validation_dataset_name']
logger.info('Validation dataset #{}: {}'.format(
len(validation_datasets) + 1, validation_dataset_name))
validation_datasets.append(
get_dataset(validation_dataset_name,
False,
False,
validation_transforms,
alphabet_decoder=alphabet_decoder))
logger.info(validation_datasets[-1])
self.validation_data_loaders = []
for validation_dataset in validation_datasets:
self.validation_data_loaders.append(
torch.utils.data.DataLoader(validation_dataset,
batch_size=1,
num_workers=8,
shuffle=False,
drop_last=False,
collate_fn=collate))
self.validate_every = config['training_details']['validate_every']
for validation_dataset in validation_datasets:
assert training_dataset.classes_num == validation_dataset.classes_num
# Model and optimizer.
logger.info('Model:')
self.model = model_class(cls_num=training_dataset.classes_num,
shape=config['shape'],
num_chars=len(alphabet_decoder.alphabet))
logger.info(self.model)
self.training_iterations_num = config['training_details'][
'training_iterations_num']
lr_scheduler_milestones = config['training_details'][
'lr_scheduler_milestones']
base_lr = config['training_details']['base_lr']
weight_decay = config['training_details']['weight_decay']
logger.info('Optimizer:')
self.optimizer = torch.optim.SGD(self.setup_optimizer(
self.model, base_lr, weight_decay),
lr=base_lr,
weight_decay=weight_decay,
momentum=0.9)
logger.info(self.optimizer)
logger.info('Learning Rate scheduler:')
self.lr_scheduler = MultiStepLRWithWarmUp(
self.optimizer,
milestones=lr_scheduler_milestones,
warmup_iters=1000,
warmup_method='linear',
warmup_factor_base=0.333,
gamma=0.1,
last_epoch=0)
logger.info(self.lr_scheduler)
self.start_step = 0
if 'backbone_checkpoint' in config and config['backbone_checkpoint']:
checkpoint_file_path = osp.join(self.root_directory,
config['backbone_checkpoint'])
if not osp.exists(checkpoint_file_path):
raise IOError(
'Initial checkpoint file "{}" does not exist. '
'Please fetch pre-trained backbone networks using '
'tools/download_pretrained_weights.py script first.'.
format(checkpoint_file_path))
logger.info(
'Loading weights from "{}"'.format(checkpoint_file_path))
load_checkpoint(self.model.backbone,
checkpoint_file_path,
verbose=True,
skip_prefix='text_recogn')
if 'checkpoint' in config and config['checkpoint']:
checkpoint_file_path = osp.join(self.root_directory,
config['checkpoint'])
if not osp.exists(checkpoint_file_path):
raise IOError('Checkpoint file "{}" does not exist. '.format(
checkpoint_file_path))
logger.info(
'Loading weights from "{}"'.format(checkpoint_file_path))
load_checkpoint(self.model, checkpoint_file_path, verbose=True)
# Loggers and misc. stuff.
self.loggers = [
TextLogger(logger),
TensorboardLogger(self.run_directory)
]
self.log_every = 50
self.checkpoint_every = config['training_details']['checkpoint_every']