inference.py 5.65 KB
Newer Older
1
2
import warnings

Kai Chen's avatar
Kai Chen committed
3
import matplotlib.pyplot as plt
myownskyW7's avatar
myownskyW7 committed
4
5
import mmcv
import numpy as np
6
import pycocotools.mask as maskUtils
myownskyW7's avatar
myownskyW7 committed
7
import torch
8
from mmcv.parallel import collate, scatter
9
from mmcv.runner import load_checkpoint
myownskyW7's avatar
myownskyW7 committed
10

11
from mmdet.core import get_classes
12
from mmdet.datasets.pipelines import Compose
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
from mmdet.models import build_detector


def init_detector(config, checkpoint=None, device='cuda:0'):
    """Initialize a detector from config file.

    Args:
        config (str or :obj:`mmcv.Config`): Config file path or the config
            object.
        checkpoint (str, optional): Checkpoint path. If left as None, the model
            will not load any weights.

    Returns:
        nn.Module: The constructed detector.
    """
    if isinstance(config, str):
        config = mmcv.Config.fromfile(config)
    elif not isinstance(config, mmcv.Config):
        raise TypeError('config must be a filename or Config object, '
                        'but got {}'.format(type(config)))
    config.model.pretrained = None
    model = build_detector(config.model, test_cfg=config.test_cfg)
    if checkpoint is not None:
        checkpoint = load_checkpoint(model, checkpoint)
        if 'CLASSES' in checkpoint['meta']:
38
            model.CLASSES = checkpoint['meta']['CLASSES']
39
40
41
42
43
44
45
46
47
48
        else:
            warnings.warn('Class names are not saved in the checkpoint\'s '
                          'meta data, use COCO classes by default.')
            model.CLASSES = get_classes('coco')
    model.cfg = config  # save the config in the model for convenience
    model.to(device)
    model.eval()
    return model


49
50
51
52
53
54
55
56
57
58
class LoadImage(object):

    def __call__(self, results):
        img = mmcv.imread(results['img'])
        results['img'] = img
        results['ori_shape'] = img.shape
        return results


def inference_detector(model, img):
59
60
61
62
63
64
65
66
67
68
69
70
71
    """Inference image(s) with the detector.

    Args:
        model (nn.Module): The loaded detector.
        imgs (str/ndarray or list[str/ndarray]): Either image files or loaded
            images.

    Returns:
        If imgs is a str, a generator will be returned, otherwise return the
        detection results directly.
    """
    cfg = model.cfg
    device = next(model.parameters()).device  # model device
72
73
74
75
76
77
78
79
    # build the data pipeline
    test_pipeline = [LoadImage()] + cfg.data.test.pipeline[1:]
    test_pipeline = Compose(test_pipeline)
    # prepare data
    data = dict(img=img)
    data = test_pipeline(data)
    data = scatter(collate([data], samples_per_gpu=1), [device])[0]
    # forward the model
Kai Chen's avatar
Kai Chen committed
80
81
82
    with torch.no_grad():
        result = model(return_loss=False, rescale=True, **data)

83
    return result
Kai Chen's avatar
Kai Chen committed
84

myownskyW7's avatar
myownskyW7 committed
85

86
# TODO: merge this method with the one in BaseDetector
simon wu's avatar
simon wu committed
87
88
89
90
91
def show_result(img,
                result,
                class_names,
                score_thr=0.3,
                wait_time=0,
Kai Chen's avatar
Kai Chen committed
92
                show=True,
simon wu's avatar
simon wu committed
93
                out_file=None):
94
    """Visualize the detection results on the image.
myownskyW7's avatar
myownskyW7 committed
95

96
97
98
99
100
101
    Args:
        img (str or np.ndarray): Image filename or loaded image.
        result (tuple[list] or list): The detection result, can be either
            (bbox, segm) or just bbox.
        class_names (list[str] or tuple[str]): A list of class names.
        score_thr (float): The threshold to visualize the bboxes and masks.
simon wu's avatar
simon wu committed
102
        wait_time (int): Value of waitKey param.
Kai Chen's avatar
Kai Chen committed
103
        show (bool, optional): Whether to show the image with opencv or not.
104
105
        out_file (str, optional): If specified, the visualization result will
            be written to the out file instead of shown in a window.
Kai Chen's avatar
Kai Chen committed
106
107
108
109

    Returns:
        np.ndarray or None: If neither `show` nor `out_file` is specified, the
            visualized image is returned, otherwise None is returned.
110
111
    """
    assert isinstance(class_names, (tuple, list))
112
    img = mmcv.imread(img)
Kai Chen's avatar
Kai Chen committed
113
    img = img.copy()
114
115
116
117
118
119
120
121
122
123
    if isinstance(result, tuple):
        bbox_result, segm_result = result
    else:
        bbox_result, segm_result = result, None
    bboxes = np.vstack(bbox_result)
    # draw segmentation masks
    if segm_result is not None:
        segms = mmcv.concat_list(segm_result)
        inds = np.where(bboxes[:, -1] > score_thr)[0]
        for i in inds:
124
            color_mask = np.random.randint(0, 256, (1, 3), dtype=np.uint8)
125
126
127
            mask = maskUtils.decode(segms[i]).astype(np.bool)
            img[mask] = img[mask] * 0.5 + color_mask * 0.5
    # draw bounding boxes
myownskyW7's avatar
myownskyW7 committed
128
129
    labels = [
        np.full(bbox.shape[0], i, dtype=np.int32)
130
        for i, bbox in enumerate(bbox_result)
myownskyW7's avatar
myownskyW7 committed
131
132
133
    ]
    labels = np.concatenate(labels)
    mmcv.imshow_det_bboxes(
Kai Chen's avatar
Kai Chen committed
134
        img,
myownskyW7's avatar
myownskyW7 committed
135
136
137
        bboxes,
        labels,
        class_names=class_names,
138
        score_thr=score_thr,
Kai Chen's avatar
Kai Chen committed
139
        show=show,
simon wu's avatar
simon wu committed
140
        wait_time=wait_time,
zhijl's avatar
zhijl committed
141
        out_file=out_file)
Kai Chen's avatar
Kai Chen committed
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
    if not (show or out_file):
        return img


def show_result_pyplot(img,
                       result,
                       class_names,
                       score_thr=0.3,
                       fig_size=(15, 10)):
    """Visualize the detection results on the image.

    Args:
        img (str or np.ndarray): Image filename or loaded image.
        result (tuple[list] or list): The detection result, can be either
            (bbox, segm) or just bbox.
        class_names (list[str] or tuple[str]): A list of class names.
        score_thr (float): The threshold to visualize the bboxes and masks.
        fig_size (tuple): Figure size of the pyplot figure.
        out_file (str, optional): If specified, the visualization result will
            be written to the out file instead of shown in a window.
    """
    img = show_result(
        img, result, class_names, score_thr=score_thr, show=False)
    plt.figure(figsize=fig_size)
    plt.imshow(mmcv.bgr2rgb(img))