1# Tencent is pleased to support the open source community by making ncnn available.
2#
3# Copyright (C) 2020 THL A29 Limited, a Tencent company. All rights reserved.
4#
5# Licensed under the BSD 3-Clause License (the "License"); you may not use this file except
6# in compliance with the License. You may obtain a copy of the License at
7#
8# https://opensource.org/licenses/BSD-3-Clause
9#
10# Unless required by applicable law or agreed to in writing, software distributed
11# under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR
12# CONDITIONS OF ANY KIND, either express or implied. See the License for the
13# specific language governing permissions and limitations under the License.
14
15import numpy as np
16import ncnn
17from .model_store import get_model_file
18from ..utils.objects import Detect_Object
19
20
21def clamp(v, lo, hi):
22    if v < lo:
23        return lo
24    elif hi < v:
25        return hi
26    else:
27        return v
28
29
30class MobileNetV3_SSDLite:
31    def __init__(self, target_size=300, num_threads=1, use_gpu=False):
32        self.target_size = target_size
33        self.num_threads = num_threads
34        self.use_gpu = use_gpu
35
36        self.mean_vals = [123.675, 116.28, 103.53]
37        self.norm_vals = [1.0, 1.0, 1.0]
38
39        self.net = ncnn.Net()
40        self.net.opt.use_vulkan_compute = self.use_gpu
41
42        # converted ncnn model from https://github.com/ujsyehao/mobilenetv3-ssd
43        # the ncnn model https://github.com/nihui/ncnn-assets/tree/master/models
44        self.net.load_param(get_model_file("mobilenetv3_ssdlite_voc.param"))
45        self.net.load_model(get_model_file("mobilenetv3_ssdlite_voc.bin"))
46
47        self.class_names = [
48            "background",
49            "aeroplane",
50            "bicycle",
51            "bird",
52            "boat",
53            "bottle",
54            "bus",
55            "car",
56            "cat",
57            "chair",
58            "cow",
59            "diningtable",
60            "dog",
61            "horse",
62            "motorbike",
63            "person",
64            "pottedplant",
65            "sheep",
66            "sofa",
67            "train",
68            "tvmonitor",
69        ]
70
71    def __del__(self):
72        self.net = None
73
74    def __call__(self, img):
75        img_h = img.shape[0]
76        img_w = img.shape[1]
77
78        mat_in = ncnn.Mat.from_pixels_resize(
79            img,
80            ncnn.Mat.PixelType.PIXEL_BGR2RGB,
81            img.shape[1],
82            img.shape[0],
83            self.target_size,
84            self.target_size,
85        )
86        mat_in.substract_mean_normalize([], self.norm_vals)
87        mat_in.substract_mean_normalize(self.mean_vals, [])
88
89        ex = self.net.create_extractor()
90        ex.set_light_mode(True)
91        ex.set_num_threads(self.num_threads)
92
93        ex.input("input", mat_in)
94
95        ret, mat_out = ex.extract("detection_out")
96
97        objects = []
98
99        # printf("%d %d %d\n", mat_out.w, mat_out.h, mat_out.c)
100
101        # method 1, use ncnn.Mat.row to get the result, no memory copy
102        for i in range(mat_out.h):
103            values = mat_out.row(i)
104
105            obj = Detect_Object()
106            obj.label = values[0]
107            obj.prob = values[1]
108
109            x1 = (
110                clamp(values[2] * self.target_size, 0.0, float(self.target_size - 1))
111                / self.target_size
112                * img_w
113            )
114            y1 = (
115                clamp(values[3] * self.target_size, 0.0, float(self.target_size - 1))
116                / self.target_size
117                * img_h
118            )
119            x2 = (
120                clamp(values[4] * self.target_size, 0.0, float(self.target_size - 1))
121                / self.target_size
122                * img_w
123            )
124            y2 = (
125                clamp(values[5] * self.target_size, 0.0, float(self.target_size - 1))
126                / self.target_size
127                * img_h
128            )
129
130            if np.isnan(x1) or np.isnan(y1) or np.isnan(x2) or np.isnan(y2):
131                continue
132
133            obj.rect.x = x1
134            obj.rect.y = y1
135            obj.rect.w = x2 - x1
136            obj.rect.h = y2 - y1
137
138            objects.append(obj)
139
140        """
141        #method 2, use ncnn.Mat->numpy.array to get the result, no memory copy too
142        out = np.array(mat_out)
143        for i in range(len(out)):
144            values = out[i]
145            obj = Detect_Object()
146            obj.label = values[0]
147            obj.prob = values[1]
148
149            x1 = clamp(values[2] * self.img_width, 0.0, float(self.img_width - 1)) / self.img_width * img_w
150            y1 = clamp(values[3] * self.img_height, 0.0, float(self.img_height - 1)) / self.img_height * img_h
151            x2 = clamp(values[4] * self.img_width, 0.0, float(self.img_width - 1)) / self.img_width * img_w
152            y2 = clamp(values[5] * self.img_height, 0.0, float(self.img_height - 1)) / self.img_height * img_h
153
154            obj.rect.x = x1
155            obj.rect.y = y1
156            obj.rect.w = x2 - x1
157            obj.rect.h = y2 - y1
158
159            objects.append(obj)
160        """
161
162        return objects
163