1
# -*- coding: utf-8 -*-
2
"""
3
Created on Fri May  1 22:45:22 2020
4

5
@author: hp
6
"""
7

8
import tensorflow as tf
9
import numpy as np
10
import cv2
11

12
from tensorflow.keras import Model
13
from tensorflow.keras.layers import (
14
    Add,
15
    Concatenate,
16
    Conv2D,
17
    Input,
18
    Lambda,
19
    LeakyReLU,
20
    UpSampling2D,
21
    ZeroPadding2D,
22
    BatchNormalization
23
)
24
from tensorflow.keras.regularizers import l2
25
import wget
26

27
def load_darknet_weights(model, weights_file):
28
    '''
29
    Helper function used to load darknet weights.
30
    
31
    :param model: Object of the Yolo v3 model
32
    :param weights_file: Path to the file with Yolo V3 weights
33
    '''
34
    
35
    #Open the weights file
36
    wf = open(weights_file, 'rb')
37
    major, minor, revision, seen, _ = np.fromfile(wf, dtype=np.int32, count=5)
38

39
    #Define names of the Yolo layers (just for a reference)    
40
    layers = ['yolo_darknet',
41
            'yolo_conv_0',
42
            'yolo_output_0',
43
            'yolo_conv_1',
44
            'yolo_output_1',
45
            'yolo_conv_2',
46
            'yolo_output_2']
47

48
    for layer_name in layers:
49
        sub_model = model.get_layer(layer_name)
50
        for i, layer in enumerate(sub_model.layers):
51
          
52
            
53
            if not layer.name.startswith('conv2d'):
54
                continue
55
                
56
            #Handles the special, custom Batch normalization layer
57
            batch_norm = None
58
            if i + 1 < len(sub_model.layers) and \
59
                    sub_model.layers[i + 1].name.startswith('batch_norm'):
60
                batch_norm = sub_model.layers[i + 1]
61

62
            filters = layer.filters
63
            size = layer.kernel_size[0]
64
            in_dim = layer.input.shape[-1]
65

66
            if batch_norm is None:
67
                conv_bias = np.fromfile(wf, dtype=np.float32, count=filters)
68
            else:
69
                # darknet [beta, gamma, mean, variance]
70
                bn_weights = np.fromfile(
71
                    wf, dtype=np.float32, count=4 * filters)
72
                # tf [gamma, beta, mean, variance]
73
                bn_weights = bn_weights.reshape((4, filters))[[1, 0, 2, 3]]
74

75
            # darknet shape (out_dim, in_dim, height, width)
76
            conv_shape = (filters, in_dim, size, size)
77
            conv_weights = np.fromfile(
78
                wf, dtype=np.float32, count=np.product(conv_shape))
79
            # tf shape (height, width, in_dim, out_dim)
80
            conv_weights = conv_weights.reshape(
81
                conv_shape).transpose([2, 3, 1, 0])
82

83
            if batch_norm is None:
84
                layer.set_weights([conv_weights, conv_bias])
85
            else:
86
                layer.set_weights([conv_weights])
87
                batch_norm.set_weights(bn_weights)
88

89
    assert len(wf.read()) == 0, 'failed to read all data'
90
    wf.close()
91
    
92
def draw_outputs(img, outputs, class_names):
93
    '''
94
    Helper, util, function that draws predictons on the image.
95
    
96
    :param img: Loaded image
97
    :param outputs: YoloV3 predictions
98
    :param class_names: list of all class names found in the dataset
99
    '''
100
    boxes, objectness, classes, nums = outputs
101
    boxes, objectness, classes, nums = boxes[0], objectness[0], classes[0], nums[0]
102
    wh = np.flip(img.shape[0:2])
103
    for i in range(nums):
104
        x1y1 = tuple((np.array(boxes[i][0:2]) * wh).astype(np.int32))
105
        x2y2 = tuple((np.array(boxes[i][2:4]) * wh).astype(np.int32))
106
        img = cv2.rectangle(img, x1y1, x2y2, (255, 0, 0), 2)
107
        img = cv2.putText(img, '{} {:.4f}'.format(
108
            class_names[int(classes[i])], objectness[i]),
109
            x1y1, cv2.FONT_HERSHEY_COMPLEX_SMALL, 1, (0, 0, 255), 2)
110
    return img
111

112
yolo_anchors = np.array([(10, 13), (16, 30), (33, 23), (30, 61), (62, 45),
113
                         (59, 119), (116, 90), (156, 198), (373, 326)],
114
                        np.float32) / 416
115

116
yolo_anchor_masks = np.array([[6, 7, 8], [3, 4, 5], [0, 1, 2]])
117
    
118
def DarknetConv(x, filters, kernel_size, strides=1, batch_norm=True):
119
    '''
120
    Call this function to define a single Darknet convolutional layer
121
    
122
    :param x: inputs
123
    :param filters: number of filters in the convolutional layer
124
    :param kernel_size: Size of kernel in the Conv layer
125
    :param strides: Conv layer strides
126
    :param batch_norm: Whether or not to use the custom batch norm layer.
127
    '''
128
    #Image padding
129
    if strides == 1:
130
        padding = 'same'
131
    else:
132
        x = ZeroPadding2D(((1, 0), (1, 0)))(x)  # top left half-padding
133
        padding = 'valid'
134
        
135
    #Defining the Conv layer
136
    x = Conv2D(filters=filters, kernel_size=kernel_size,
137
               strides=strides, padding=padding,
138
               use_bias=not batch_norm, kernel_regularizer=l2(0.0005))(x)
139
    
140
    if batch_norm:
141
        x = BatchNormalization()(x)
142
        x = LeakyReLU(alpha=0.1)(x)
143
    return x
144

145
def DarknetResidual(x, filters):
146
    '''
147
    Call this function to define a single DarkNet Residual layer
148
    
149
    :param x: inputs
150
    :param filters: number of filters in each Conv layer.
151
    '''
152
    prev = x
153
    x = DarknetConv(x, filters // 2, 1)
154
    x = DarknetConv(x, filters, 3)
155
    x = Add()([prev, x])
156
    return x
157
  
158
  
159
def DarknetBlock(x, filters, blocks):
160
    '''
161
    Call this function to define a single DarkNet Block (made of multiple Residual layers)
162
    
163
    :param x: inputs
164
    :param filters: number of filters in each Residual layer
165
    :param blocks: number of Residual layers in the block
166
    '''
167
    x = DarknetConv(x, filters, 3, strides=2)
168
    for _ in range(blocks):
169
        x = DarknetResidual(x, filters)
170
    return x
171

172
def Darknet(name=None):
173
    '''
174
    The main function that creates the whole DarkNet.
175
    '''
176
    x = inputs = Input([None, None, 3])
177
    x = DarknetConv(x, 32, 3)
178
    x = DarknetBlock(x, 64, 1)
179
    x = DarknetBlock(x, 128, 2)  # skip connection
180
    x = x_36 = DarknetBlock(x, 256, 8)  # skip connection
181
    x = x_61 = DarknetBlock(x, 512, 8)
182
    x = DarknetBlock(x, 1024, 4)
183
    return tf.keras.Model(inputs, (x_36, x_61, x), name=name)
184

185
def YoloConv(filters, name=None):
186
    '''
187
    Call this function to define the Yolo Conv layer.
188
    
189
    :param flters: number of filters for the conv layer
190
    :param name: name of the layer
191
    '''
192
    def yolo_conv(x_in):
193
        if isinstance(x_in, tuple):
194
            inputs = Input(x_in[0].shape[1:]), Input(x_in[1].shape[1:])
195
            x, x_skip = inputs
196

197
            # concat with skip connection
198
            x = DarknetConv(x, filters, 1)
199
            x = UpSampling2D(2)(x)
200
            x = Concatenate()([x, x_skip])
201
        else:
202
            x = inputs = Input(x_in.shape[1:])
203

204
        x = DarknetConv(x, filters, 1)
205
        x = DarknetConv(x, filters * 2, 3)
206
        x = DarknetConv(x, filters, 1)
207
        x = DarknetConv(x, filters * 2, 3)
208
        x = DarknetConv(x, filters, 1)
209
        return Model(inputs, x, name=name)(x_in)
210
    return yolo_conv
211

212
def YoloOutput(filters, anchors, classes, name=None):
213
    '''
214
    This function defines outputs for the Yolo V3. (Creates output projections)
215
     
216
    :param filters: number of filters for the conv layer
217
    :param anchors: anchors
218
    :param classes: list of classes in a dataset
219
    :param name: name of the layer
220
    '''
221
    def yolo_output(x_in):
222
        x = inputs = Input(x_in.shape[1:])
223
        x = DarknetConv(x, filters * 2, 3)
224
        x = DarknetConv(x, anchors * (classes + 5), 1, batch_norm=False)
225
        x = Lambda(lambda x: tf.reshape(x, (-1, tf.shape(x)[1], tf.shape(x)[2],
226
                                            anchors, classes + 5)))(x)
227
        return tf.keras.Model(inputs, x, name=name)(x_in)
228
    return yolo_output
229

230
def yolo_boxes(pred, anchors, classes):
231
    '''
232
    Call this function to get bounding boxes from network predictions
233
    
234
    :param pred: Yolo predictions
235
    :param anchors: anchors
236
    :param classes: List of classes from the dataset
237
    '''
238
    
239
    # pred: (batch_size, grid, grid, anchors, (x, y, w, h, obj, ...classes))
240
    grid_size = tf.shape(pred)[1]
241
    #Extract box coortinates from prediction vectors
242
    box_xy, box_wh, objectness, class_probs = tf.split(
243
        pred, (2, 2, 1, classes), axis=-1)
244

245
    #Normalize coortinates
246
    box_xy = tf.sigmoid(box_xy)
247
    objectness = tf.sigmoid(objectness)
248
    class_probs = tf.sigmoid(class_probs)
249
    pred_box = tf.concat((box_xy, box_wh), axis=-1)  # original xywh for loss
250

251
    # !!! grid[x][y] == (y, x)
252
    grid = tf.meshgrid(tf.range(grid_size), tf.range(grid_size))
253
    grid = tf.expand_dims(tf.stack(grid, axis=-1), axis=2)  # [gx, gy, 1, 2]
254

255
    box_xy = (box_xy + tf.cast(grid, tf.float32)) / \
256
        tf.cast(grid_size, tf.float32)
257
    box_wh = tf.exp(box_wh) * anchors
258

259
    box_x1y1 = box_xy - box_wh / 2
260
    box_x2y2 = box_xy + box_wh / 2
261
    bbox = tf.concat([box_x1y1, box_x2y2], axis=-1)
262

263
    return bbox, objectness, class_probs, pred_box
264

265
def yolo_nms(outputs, anchors, masks, classes):
266
    # boxes, conf, type
267
    b, c, t = [], [], []
268

269
    for o in outputs:
270
        b.append(tf.reshape(o[0], (tf.shape(o[0])[0], -1, tf.shape(o[0])[-1])))
271
        c.append(tf.reshape(o[1], (tf.shape(o[1])[0], -1, tf.shape(o[1])[-1])))
272
        t.append(tf.reshape(o[2], (tf.shape(o[2])[0], -1, tf.shape(o[2])[-1])))
273

274
    bbox = tf.concat(b, axis=1)
275
    confidence = tf.concat(c, axis=1)
276
    class_probs = tf.concat(t, axis=1)
277

278
    scores = confidence * class_probs
279
    boxes, scores, classes, valid_detections = tf.image.combined_non_max_suppression(
280
        boxes=tf.reshape(bbox, (tf.shape(bbox)[0], -1, 1, 4)),
281
        scores=tf.reshape(
282
        scores, (tf.shape(scores)[0], -1, tf.shape(scores)[-1])),
283
        max_output_size_per_class=100,
284
        max_total_size=100,
285
        iou_threshold=0.5,
286
        score_threshold=0.6
287
    )
288

289
    return boxes, scores, classes, valid_detections
290

291

292
def YoloV3(size=None, channels=3, anchors=yolo_anchors,
293
           masks=yolo_anchor_masks, classes=80):
294
  
295
    x = inputs = Input([size, size, channels], name='input')
296

297
    x_36, x_61, x = Darknet(name='yolo_darknet')(x)
298

299
    x = YoloConv(512, name='yolo_conv_0')(x)
300
    output_0 = YoloOutput(512, len(masks[0]), classes, name='yolo_output_0')(x)
301

302
    x = YoloConv(256, name='yolo_conv_1')((x, x_61))
303
    output_1 = YoloOutput(256, len(masks[1]), classes, name='yolo_output_1')(x)
304

305
    x = YoloConv(128, name='yolo_conv_2')((x, x_36))
306
    output_2 = YoloOutput(128, len(masks[2]), classes, name='yolo_output_2')(x)
307

308
    boxes_0 = Lambda(lambda x: yolo_boxes(x, anchors[masks[0]], classes),
309
                     name='yolo_boxes_0')(output_0)
310
    boxes_1 = Lambda(lambda x: yolo_boxes(x, anchors[masks[1]], classes),
311
                     name='yolo_boxes_1')(output_1)
312
    boxes_2 = Lambda(lambda x: yolo_boxes(x, anchors[masks[2]], classes),
313
                     name='yolo_boxes_2')(output_2)
314

315
    outputs = Lambda(lambda x: yolo_nms(x, anchors, masks, classes),
316
                     name='yolo_nms')((boxes_0[:3], boxes_1[:3], boxes_2[:3]))
317

318
    return Model(inputs, outputs, name='yolov3')
319

320
def weights_download(out='models/yolov3.weights'):
321
    _ = wget.download('https://pjreddie.com/media/files/yolov3.weights', out='models/yolov3.weights')
322
    
323
# weights_download() # to download weights
324
yolo = YoloV3()
325
load_darknet_weights(yolo, 'models/yolov3.weights') 
326

327

328
def detect_phone_and_person(video_path):
329
    cap = cv2.VideoCapture(video_path)
330

331
    while(True):
332
        ret, image = cap.read()
333
        if ret == False:
334
            break
335
        img = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
336
        img = cv2.resize(img, (320, 320))
337
        img = img.astype(np.float32)
338
        img = np.expand_dims(img, 0)
339
        img = img / 255
340
        class_names = [c.strip() for c in open("models/classes.TXT").readlines()]
341
        boxes, scores, classes, nums = yolo(img)
342
        count=0
343
        for i in range(nums[0]):
344
            if int(classes[0][i] == 0):
345
                count +=1
346
            if int(classes[0][i] == 67):
347
                print('Mobile Phone detected')
348
        if count == 0:
349
            print('No person detected')
350
        elif count > 1: 
351
            print('More than one person detected')
352
            
353
        image = draw_outputs(image, (boxes, scores, classes, nums), class_names)
354

355
        cv2.imshow('Prediction', image)
356
        if cv2.waitKey(1) & 0xFF == ord('q'):
357
            break
358

359

360
    cap.release()
361
    cv2.destroyAllWindows()
362

363

364