def load_model(args, verbose=False):
if args.command != 'train' and not os.path.isfile(args.model):
raise RuntimeError('Model file {} does not exist!'.format(args.model))
model = None
state = {}
_, ext = os.path.splitext(args.model)
if args.command == 'train' and (not os.path.exists(args.model) or args.override):
if verbose: print('Initializing model...')
model = Model(args.backbone, args.classes)
model.initialize(args.fine_tune)
if verbose: print(model)
elif ext == '.pth' or ext == '.torch':
if verbose: print('Loading model from {}...'.format(os.path.basename(args.model)))
model, state = Model.load(args.model)
if verbose: print(model)
elif args.command == 'infer' and ext in ['.engine', '.plan']:
model = None
else:
raise RuntimeError('Invalid model format "{}"!'.format(args.ext))
state['path'] = args.model
return model, state
def load_model(args, verbose=False):
config = {}
if args.config:
with open(args.config, 'r') as config_file:
config = json.load(config_file)
if args.command != 'train' and not os.path.isfile(args.model):
raise RuntimeError('Model file {} does not exist!'.format(args.model))
model = None
state = {}
model_name, ext = os.path.splitext(args.model)
if args.command == 'train' and (not os.path.exists(args.model)
or args.override):
if verbose: print('Initializing model...')
model = Model(backbones=args.backbone,
classes=args.classes,
rotated_bbox=args.rotated_bbox,
config=config)
model.initialize(args.fine_tune)
if verbose: print(model)
elif ext == '.pth' or ext == '.torch':
if verbose:
print('Loading model from {}...'.format(
os.path.basename(args.model)))
exporting = False
if args.command == 'eval':
exporting = True
model, state = Model.load(filename=args.model,
rotated_bbox=args.rotated_bbox,
config=config,
exporting=exporting)
if verbose: print(model)
elif args.command == 'infer' and ext in ['.engine', '.plan']:
model = None
else:
raise RuntimeError('Invalid model format "{}"!'.format(args.ext))
state['path'] = model_name
return model, state
