| |
| from argparse import ArgumentParser |
|
|
| from mmengine.config import Config |
| from mmengine.logging import print_log |
|
|
| from mmdet.apis import DetInferencer |
| from projects.XDecoder.xdecoder.inference import ( |
| ImageCaptionInferencer, RefImageCaptionInferencer, |
| TextToImageRegionRetrievalInferencer) |
|
|
| TASKINFOS = { |
| 'semseg': DetInferencer, |
| 'ref-seg': DetInferencer, |
| 'instance': DetInferencer, |
| 'panoptic': DetInferencer, |
| 'caption': ImageCaptionInferencer, |
| 'ref-caption': RefImageCaptionInferencer, |
| 'retrieval': TextToImageRegionRetrievalInferencer, |
| } |
|
|
|
|
| def parse_args(): |
| parser = ArgumentParser() |
| parser.add_argument( |
| 'inputs', type=str, help='Input image file or folder path.') |
| parser.add_argument('model', type=str, help='Config file name') |
| parser.add_argument('--weights', help='Checkpoint file') |
| parser.add_argument('--texts', help='text prompt') |
| parser.add_argument( |
| '--out-dir', |
| type=str, |
| default='outputs', |
| help='Output directory of images or prediction results.') |
| parser.add_argument( |
| '--device', default='cuda:0', help='Device used for inference') |
| parser.add_argument( |
| '--show', |
| action='store_true', |
| help='Display the image in a popup window.') |
| parser.add_argument( |
| '--no-save-vis', |
| action='store_true', |
| help='Do not save detection vis results') |
| parser.add_argument( |
| '--palette', |
| default='none', |
| choices=['ade20k', 'coco', 'voc', 'citys', 'random', 'none'], |
| help='Color palette used for visualization') |
|
|
| |
| parser.add_argument( |
| '--pred-score-thr', |
| type=float, |
| default=0.5, |
| help='bbox score threshold') |
| |
| parser.add_argument( |
| '--stuff-texts', |
| help='text prompt for stuff name in panoptic segmentation') |
|
|
| call_args = vars(parser.parse_args()) |
| if call_args['no_save_vis']: |
| call_args['out_dir'] = '' |
|
|
| init_kws = ['model', 'weights', 'device', 'palette'] |
| init_args = {} |
| for init_kw in init_kws: |
| init_args[init_kw] = call_args.pop(init_kw) |
|
|
| return init_args, call_args |
|
|
|
|
| def main(): |
| init_args, call_args = parse_args() |
|
|
| cfg = Config.fromfile(init_args['model']) |
| task = cfg.model.head.task |
| assert task in TASKINFOS |
|
|
| inferencer = TASKINFOS[task](**init_args) |
|
|
| if task != 'caption': |
| assert call_args[ |
| 'texts'] is not None, f'text prompts is required for {task}' |
| if task != 'panoptic': |
| call_args.pop('stuff_texts') |
| else: |
| call_args.pop('texts') |
| call_args.pop('stuff_texts') |
|
|
| inferencer(**call_args) |
|
|
| if call_args['out_dir'] != '' and not call_args['no_save_vis']: |
| print_log(f'results have been saved at {call_args["out_dir"]}') |
|
|
|
|
| if __name__ == '__main__': |
| main() |
|
|
