mmclassification/demo/image_demo.py
Ezra-Yu 922a31a68e
[Feature] ADD RepMLP Backbone and checkpoints. (#709)
* add repmlp

* refactor backbone code

* add checkpoint, readme, unit tests

* unchange demo

* improve docstring.

* update pic

* refactor patch embed

* refactor unit tests

* fix lint

* update tools

* update tools

* Update checkpoint path

Co-authored-by: mzr1996 <mzr1996@163.com>
2022-03-30 20:40:25 +08:00

29 lines
924 B
Python

# Copyright (c) OpenMMLab. All rights reserved.
from argparse import ArgumentParser
from mmcls.apis import inference_model, init_model, show_result_pyplot
def main():
parser = ArgumentParser()
parser.add_argument('img', help='Image file')
parser.add_argument('config', help='Config file')
parser.add_argument('checkpoint', help='Checkpoint file')
parser.add_argument(
'--device', default='cuda:0', help='Device used for inference')
args = parser.parse_args()
# build the model from a config file and a checkpoint file
model = init_model(args.config, args.checkpoint, device=args.device)
for i, stage in enumerate(model.backbone.stages):
print(i, len(stage))
print(i, stage[0])
# test a single image
result = inference_model(model, args.img)
# show the results
show_result_pyplot(model, args.img, result)
if __name__ == '__main__':
main()