pcd_demo.py 2.78 KB
Newer Older
dingchang's avatar
dingchang committed
1
# Copyright (c) OpenMMLab. All rights reserved.
2
3
import logging
import os
wuyuefeng's avatar
Demo  
wuyuefeng committed
4
5
from argparse import ArgumentParser

6
7
8
from mmengine.logging import print_log

from mmdet3d.apis import LidarDet3DInferencer
wuyuefeng's avatar
Demo  
wuyuefeng committed
9

ZCMax's avatar
ZCMax committed
10
11

def parse_args():
wuyuefeng's avatar
Demo  
wuyuefeng committed
12
13
    parser = ArgumentParser()
    parser.add_argument('pcd', help='Point cloud file')
14
15
    parser.add_argument('model', help='Config file')
    parser.add_argument('weights', help='Checkpoint file')
wuyuefeng's avatar
Demo  
wuyuefeng committed
16
17
18
    parser.add_argument(
        '--device', default='cuda:0', help='Device used for inference')
    parser.add_argument(
19
20
21
22
        '--pred-score-thr',
        type=float,
        default=0.3,
        help='bbox score threshold')
wuyuefeng's avatar
Demo  
wuyuefeng committed
23
    parser.add_argument(
24
25
26
27
        '--out-dir',
        type=str,
        default='outputs',
        help='Output directory of prediction and visualization results.')
28
    parser.add_argument(
Zongbao Feng's avatar
Zongbao Feng committed
29
30
        '--show',
        action='store_true',
31
32
33
34
35
36
37
38
39
40
41
        help='Show online visualization results')
    parser.add_argument(
        '--wait-time',
        type=float,
        default=-1,
        help='The interval of show (s). Demo will be blocked in showing'
        'results, if wait_time is -1. Defaults to -1.')
    parser.add_argument(
        '--no-save-vis',
        action='store_true',
        help='Do not save detection visualization results')
42
    parser.add_argument(
43
        '--no-save-pred',
44
        action='store_true',
45
46
47
48
49
50
51
52
        help='Do not save detection prediction results')
    parser.add_argument(
        '--print-result',
        action='store_true',
        help='Whether to print the results.')
    call_args = vars(parser.parse_args())

    call_args['inputs'] = dict(points=call_args.pop('pcd'))
ZCMax's avatar
ZCMax committed
53

54
55
    if call_args['no_save_vis'] and call_args['no_save_pred']:
        call_args['out_dir'] = ''
ZCMax's avatar
ZCMax committed
56

57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
    init_kws = ['model', 'weights', 'device']
    init_args = {}
    for init_kw in init_kws:
        init_args[init_kw] = call_args.pop(init_kw)

    # NOTE: If your operating environment does not have a display device,
    # (e.g. a remote server), you can save the predictions and visualize
    # them in local devices.
    if os.environ.get('DISPLAY') is None and call_args['show']:
        print_log(
            'Display device not found. `--show` is forced to False',
            logger='current',
            level=logging.WARNING)
        call_args['show'] = False

    return init_args, call_args


def main():
ZCMax's avatar
ZCMax committed
76
    # TODO: Support inference of point cloud numpy file.
77
78
79
80
81
82
83
84
85
86
    init_args, call_args = parse_args()

    inferencer = LidarDet3DInferencer(**init_args)
    inferencer(**call_args)

    if call_args['out_dir'] != '' and not (call_args['no_save_vis']
                                           and call_args['no_save_pred']):
        print_log(
            f'results have been saved at {call_args["out_dir"]}',
            logger='current')
wuyuefeng's avatar
Demo  
wuyuefeng committed
87
88
89


if __name__ == '__main__':
90
    main()