def main():
import sys, os
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../'))
from torchocr.models import build_model
from torchocr.utils.config_util import Config
from torchocr.utils.checkpoints import load_checkpoint, save_checkpoint
cfg_path = 'config/det/dbnet/61_hw_repb2_dbnet.py'
model_path = 'work_dirs/61_hw_repb2_dbnet/det.pth'
cfg = Config.fromfile(cfg_path)
cfg.model.pretrained = None
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
# build model
train_model = build_model(cfg.model)
load_checkpoint(train_model, model_path, map_location=device)
train_model = train_model.to(device)
cfg.model.backbone.is_deploy = True
deploy_model = build_model(cfg.model)
deploy_model = deploy_model.to(device)
deploy_weights = whole_model_convert(train_model, deploy_model)
save_checkpoint(deploy_weights, filepath='db_repvgg.pth')
def __init__(self, cfg, args):
self.postprocess = build_postprocess(cfg.postprocess)
self.mode = args.mode
if self.mode == 'torch':
model_path = args.model_path
# for rec cal head number
if hasattr(self.postprocess, 'character'):
char_num = len(getattr(self.postprocess, 'character'))
cfg.model.head.n_class = char_num
self.model = build_model(cfg.model)
self.device = torch.device(
'cuda:0' if torch.cuda.is_available() else 'cpu')
load_checkpoint(self.model, model_path, map_location=self.device)
self.model.to(self.device)
self.model.eval()
elif self.mode == 'onnx':
onnx_path = args.onnx_path
self.model = ONNXModel(onnx_path)
elif self.mode == 'engine':
engine_path = args.engine_path
self.model = TRTModel(engine_path)
# pre-heat
self.resize = RecResizeImg(image_shape=[3, 32, 800],
infer_mode=False,
character_type='ch')
def main():
args = parse_args()
cfg_path = args.config
cfg = Config.fromfile(cfg_path)
# 通用配置
global_config = cfg.options
# build model
model = build_model(cfg.model)
device, gpu_ids = select_device(global_config.gpu_ids)
load_checkpoint(model, args.model_path,map_location=device)
model = model.to(device)
model.device = device
eval_dataset = build_dataset(cfg.test_data.dataset)
eval_loader = build_dataloader(eval_dataset, loader_cfg=cfg.test_data.loader)
# build postprocess
postprocess = build_postprocess(cfg.postprocess)
# build metric
metric = build_metrics(cfg.metric)
result_metirc = eval(model, eval_loader, postprocess, metric)
print(result_metirc)
def __init__(self, cfg, model_path):
self.model = build_model(cfg.model)
load_checkpoint(self.model, model_path)
self.device = torch.device(
'cuda:0' if torch.cuda.is_available() else 'cpu')
self.model.to(self.device)
self.model.eval()
self.postprocess = build_postprocess(cfg.postprocess)
self.resize = RecResizeImg(image_shape=[3, 32, 100],
infer_mode=False,
character_type='ch')
def __init__(self, cfg, model_path):
self.model = build_model(cfg.model)
load_checkpoint(self.model, model_path)
self.device = torch.device(
'cuda:0' if torch.cuda.is_available() else 'cpu')
self.model.to(self.device)
self.model.eval()
self.postprocess = build_postprocess(cfg.postprocess)
self.normalize = NormalizeImage(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
self.to_chw = ToCHWImage()
self.resize = DetResizeForTest([736, 1280])
def main():
args = parse_args()
cfg_path = args.config
cfg = Config.fromfile(cfg_path)
# set pretrained model None
cfg.model.pretrained = None
# build postprocess
postprocess = build_postprocess(cfg.postprocess)
# for rec cal head number
if hasattr(postprocess, 'character'):
char_num = len(getattr(postprocess, 'character'))
cfg.model.head.n_class = char_num
eval_dataset = build_dataset(cfg.test_data.dataset)
eval_loader = build_dataloader(eval_dataset,
loader_cfg=cfg.test_data.loader)
# build metric
metric = build_metrics(cfg.metric)
mode = args.mode
if mode == 'torch':
# build model
model = build_model(cfg.model)
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
load_checkpoint(model, args.model_path, map_location=device)
if args.simple:
(filepath, filename) = os.path.split(args.model_path)
simple_model_path = os.path.join(filepath,
'sim_{}'.format(filename))
save_checkpoint(model, simple_model_path)
model = model.to(device)
model.device = device
result_metirc = eval(model, eval_loader, postprocess, metric)
elif mode == 'engine':
engine_path = args.engine_path
model = TRTModel(engine_path)
result_metirc = engine_eval(model, eval_loader, postprocess, metric)
print(result_metirc)
def __init__(self, cfg, args):
self.mode = args.mode
self.postprocess = build_postprocess(cfg.postprocess)
if self.mode == 'torch':
model_path = args.model_path
self.model = build_model(cfg.model)
self.device = torch.device(
'cuda:0' if torch.cuda.is_available() else 'cpu')
load_checkpoint(self.model, model_path, map_location=self.device)
self.model.to(self.device)
self.model.eval()
elif self.mode == 'onnx':
onnx_path = args.onnx_path
self.model = ONNXModel(onnx_path)
elif self.mode == 'engine':
engine_path = args.engine_path
self.model = TRTModel(engine_path)
self.normalize = NormalizeImage(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225])
self.to_chw = ToCHWImage()
self.resize = DetResizeForTest(short_size=736, mode='db')
def main():
args = parse_args()
cfg_path = args.config
cfg = Config.fromfile(cfg_path)
global_config = cfg.options # 通用配置
# local_rank = 0 is logger
global_config['local_rank'] = args.local_rank
# amp train
if args.amp:
global_config['is_amp'] = True
else:
global_config['is_amp'] = False
# ema train
if args.ema:
global_config['is_ema'] = True
else:
global_config['is_ema'] = False
# set cudnn_benchmark,如数据size一致能加快训练
if global_config.get('cudnn_benchmark', False):
torch.backends.cudnn.benchmark = True
if global_config.get('work_dir', None) is None:
# use config filename as default work_dir if cfg.work_dir is None
global_config.work_dir = osp.join(
'./work_dirs',
osp.splitext(osp.basename(args.config))[0])
# create work_dir
file_util.mkdir_or_exist(global_config.work_dir)
# init the logger before other steps
timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
log_file = osp.join(global_config.work_dir, '{}.log'.format(timestamp))
logger = get_logger(name='ocr', log_file=log_file)
# # log env info
if args.local_rank == 0:
env_info_dict = collect_env()
env_info = '\n'.join([('{}: {}'.format(k, v))
for k, v in env_info_dict.items()])
dash_line = '-' * 60 + '\n'
logger.info('Environment info:\n' + dash_line + env_info + '\n' +
dash_line)
## log some basic info
logger.info('Config:\n{}'.format(cfg.text))
# set random seeds
logger.info('Set random seed to {}, deterministic: {}'.format(
global_config.seed, args.deterministic))
# set random seed
set_random_seed(global_config.seed, deterministic=args.deterministic)
# select device
# dist init
if torch.cuda.device_count() > 1 and args.distributed:
device = init_dist(launcher='pytorch',
backend='nccl',
rank=args.local_rank)
global_config['distributed'] = True
else:
device, gpu_ids = select_device(global_config.gpu_ids)
global_config.gpu_ids = gpu_ids
global_config['distributed'] = False
# build train dataset
train_dataset = build_dataset(cfg.train_data.dataset)
train_loader = build_dataloader(train_dataset,
loader_cfg=cfg.train_data.loader,
distributed=global_config['distributed'])
# if is eval , build eval dataloader,postprocess,metric
# 移动到前面，由于rec-head的输出需要用postprocess计算
if global_config.is_eval:
eval_dataset = build_dataset(cfg.test_data.dataset)
eval_loader = build_dataloader(
eval_dataset,
loader_cfg=cfg.test_data.loader,
distributed=global_config['distributed'])
# build postprocess
postprocess = build_postprocess(cfg.postprocess)
# build metric
metric = build_metrics(cfg.metric)
else:
eval_loader = None
postprocess = None
metric = None
# for rec cal head number
if hasattr(postprocess, 'character'):
char_num = len(getattr(postprocess, 'character'))
cfg.model.head.n_class = char_num
# build model
model = build_model(cfg.model)
model = model.to(device)
# set model to device
if device.type != 'cpu' and torch.cuda.device_count(
) > 1 and global_config['distributed'] == True:
model = DDP(model,
device_ids=[args.local_rank],
output_device=args.local_rank)
device = torch.device('cuda', args.local_rank)
is_cuda = True
elif device.type != 'cpu' and global_config[
'distributed'] == False and len(gpu_ids) >= 1:
model = nn.DataParallel(model, device_ids=global_config.gpu_ids)
model.gpu_ids = gpu_ids
is_cuda = True
else:
is_cuda = False
global_config['is_cuda'] = is_cuda
model.device = device
# build optimizer
optimizer = build_optimizer(cfg.optimizer, model)
# build lr_scheduler
lr_scheduler = build_lr_scheduler(cfg.lr_scheduler, optimizer)
# build loss
criterion = build_loss(cfg.loss).to(device)
runner = TrainRunner(global_config, model, optimizer, lr_scheduler,
postprocess, criterion, train_loader, eval_loader,
metric, logger)
# # Resume
if global_config.resume_from is not None and args.resume:
runner.resume(global_config.resume_from, map_location=device)
if global_config.load_from is not None:
runner.load_checkpoint(global_config.load_from, map_location=device)
runner.run()
def main():
args = parse_args()
cfg_path = args.config
cfg = Config.fromfile(cfg_path)
# set pretrained model None
cfg.model.pretrained = None
# build postprocess
postprocess = build_postprocess(cfg.postprocess)
# for rec cal head number
if hasattr(postprocess, 'character'):
char_num = len(getattr(postprocess, 'character'))
cfg.model.head.n_class = char_num
# use config build model
model = build_model(cfg.model)
# set weights to model and set model to device/eval()
model_path = args.weights
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
load_checkpoint(model, model_path, map_location=device)
model = model.to(device)
model.eval()
onnx_output = args.onnx_output
input_shape = args.input_shape
input_shape = eval(input_shape)
mode = args.mode
# transform torch model to onnx model
input_names = ['input']
output_names = ['output']
# input shape
input_data = torch.randn(input_shape).to(device)
if args.is_dynamic:
if mode == 'rec':
# #rec
dynamic_axes = {"input": {0: "batch_size"}, "output": {0: "batch_size"}}
elif mode == 'det':
## det
dynamic_axes = {"input": {0: "batch_size", 2: 'height', 3: 'width'},
"output": {0: "batch_size", 2: 'height', 3: 'width'}}
else:
dynamic_axes = None
onnx_model_name = torch2onnx(
model=model,
dummy_input=input_data,
onnx_model_name=onnx_output,
input_names=input_names,
output_names=output_names,
opset_version=12,
is_dynamic=args.is_dynamic,
dynamic_axes=dynamic_axes
)
onnx_model = onnx.load(onnx_model_name)
# check that the model converted fine
onnx.checker.check_model(onnx_model)
onnx.helper.printable_graph(onnx_model.graph)
print("Model was successfully converted to ONNX format.")
print("It was saved to", onnx_model_name)
sys.path.append('./')
os.environ['CUDA_VISIBLE_DEVICES'] = '1'
from torchocr.models import build_model
from torchocr.models import build_head
device = torch.device('cuda:0')
det_data = torch.randn(1, 3, 640, 640).to(device)
db_model = dict(
type='DetectionModel',
transform=None,
backbone=dict(type='DetResNet', in_channels=3, depth=50),
neck=dict(type='EASTWithUnet',
in_channels=[256, 512, 1024, 2048],
out_channels=128),
# head=dict(
# type='C2TDHead',
# in_channels=3,
# ),
head=None)
det_model = build_model(cfg=db_model)
# print(det_model)
det_model = det_model.to(device)
y = det_model(det_data)
print(y.shape)
