Amazing-Python-Scripts

1
# *^_^* coding:utf-8 *^_^*
2
"""
3
SmokeDect-v0.2
4
检测方法：用MOG2提取运动区域，分小块提取hog特征，用svm分类
5
"""
6

7
from __future__ import print_function  # use print()
8

9
__author__ = 'stone'
10
__date__ = '15-11-30'
11

12
import cv2
13
import numpy as np
14
from time import time
15
import os
16

17
DEBUG_MOD = True
18
vName = '/home/stone/Code/FlameSmokeDetect/medias/videos/CTC.BG.055_11_320x240.avi'
19
train_directory = '/home/stone/Code/FlameSmokeDetect/tools/clips'
20
BLOCK_SIZE = 40
21

22

23
def hog_feature(img):
24
    """
25
    Extract HOG feature
26
    """
27
    win_size = (16, 16)
28
    block_size = (8, 8)
29
    block_stride = (8, 8)
30
    cell_size = (8, 8)
31
    nbins = 9
32
    descriptor = cv2.HOGDescriptor(
33
        win_size, block_size, block_stride, cell_size, nbins)
34
    hog_feature_hog = descriptor.compute(img)
35
    hog_feature_hog = np.float32(hog_feature_hog)
36
    return hog_feature_hog
37

38

39
def svm_create():
40
    """
41
    initial a SVM
42
    """
43
    svm_params = dict(kernel_type=cv2.ml.SVM_LINEAR,
44
                      svm_type=cv2.ml.SVM_C_SVC,
45
                      C=1,
46
                      gamma=0.5)
47
    svm = cv2.ml.SVM_create()
48
    svm.setKernel(svm_params['kernel_type'])
49
    svm.setType(svm_params['svm_type'])
50
    svm.setC(svm_params['C'])
51
    svm.setGamma(svm_params['gamma'])
52
    return svm
53

54

55
def file_path(directory):
56
    """
57
    generate a full directory path
58
    """
59
    files = os.listdir(directory)
60
    path = []
61

62
    for name in files:
63
        full_name = os.path.join(directory, name)
64
        path.append(full_name)
65
    print('%s 的文件数: %d\n' % (directory, len(path)))
66
    return path
67

68

69
# def load_hog(hog_path):
70
#     """
71
#     load hog files
72
#     """
73
#     print('loading hog files...')
74
#     hog = []
75
#     for fp in hog_path:
76
#         content = np.loadtxt(fp)
77
#         hog.append(content)
78
#     hog = np.float32(hog)
79
#     # print(hog)
80
#     print('finished!')
81
#     return hog
82

83

84
def split(img, cell_size, flatten=False):
85
    """
86
    Split into small blocks
87
    """
88
    split_h, split_w = img.shape[:2]
89
    sx, sy = cell_size
90
    split_cells = [np.hsplit(row, split_w // sx)
91
                   for row in np.vsplit(img, split_h // sy)]
92
    split_cells = np.array(split_cells)
93
    if flatten:
94
        split_cells = split_cells.reshape(-1, sy, sx)
95
    return split_cells
96

97

98
if __name__ == '__main__':
99
    print(__doc__)
100
    start_time = time()
101
    count = 0
102

103
    # create and train a SVM
104
    train_path = file_path(train_directory)
105
    train_hog = []
106
    response = []
107
    for img_path in train_path:
108
        print(img_path.split('-')[0])
109
        if img_path.split('-')[0] == '/home/stone/Code/FlameSmokeDetect/tools/clips/p':
110
            response.append(1)
111
        elif img_path.split('-')[0] == '/home/stone/Code/FlameSmokeDetect/tools/clips/n':
112
            response.append(0)
113
        img = cv2.imread(img_path)
114
        train_hog.append(hog_feature(img))
115
    print(len(train_hog[0]))
116
    train_hog = np.float32(train_hog)
117
    print(response)
118
    response = np.array(response)
119
    svm = svm_create()
120
    svm.train(train_hog, cv2.ml.ROW_SAMPLE, response)
121
    if svm.isTrained():
122
        print('SVM is trained!')
123

124
    cap = cv2.VideoCapture(vName)
125
    fgbg = cv2.createBackgroundSubtractorMOG2(detectShadows=False)
126

127
    while 1:
128
        result = []
129
        ret, frame = cap.read()
130
        if frame is None:
131
            print("The End!!!")
132
            break
133

134
        # if count < 500:
135
        #     count += 1
136
        #     continue
137

138
        frame_copy = frame.copy()
139
        frame = cv2.GaussianBlur(frame, (5, 5), 2)
140
        frame = cv2.medianBlur(frame, 5)
141

142
        gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
143
        fmask = fgbg.apply(gray_frame)
144

145
        kernel = np.ones((20, 20), np.uint8)
146
        fmask = cv2.medianBlur(fmask, 3)
147
        fmask = cv2.dilate(fmask, kernel)
148

149
        fmask_copy = fmask.copy()
150

151
        contour_img, contours, hierarchy = cv2.findContours(
152
            fmask_copy, cv2.RETR_LIST, cv2.CHAIN_APPROX_SIMPLE)
153

154
        if contours is None:
155
            continue
156

157
        for cnt in contours:
158
            x, y, w, h = cv2.boundingRect(cnt)
159

160
            if w < BLOCK_SIZE or h < BLOCK_SIZE:
161
                pass
162
            else:
163
                cells = split(frame, (BLOCK_SIZE, BLOCK_SIZE), flatten=False)
164
                cells_x_1 = x / BLOCK_SIZE
165
                cells_x_2 = (x + w) / BLOCK_SIZE
166
                cells_y_1 = y / BLOCK_SIZE
167
                cells_y_2 = (y + h) / BLOCK_SIZE
168
                for candidate in cells[cells_y_1:cells_y_2, cells_x_1:cells_x_2]:
169
                    for i in xrange(candidate.shape[0]):
170
                        hog = []
171
                        hog.append(hog_feature(candidate[i]))
172
                        hog = np.float32(hog)
173

174
                        result.append(svm.predict(hog)[0])
175
                        # print('HogFeature/hog%d-%d-%d' % (count, i, j))
176
                        # hog_file = 'HogFeature/hog%d-%d-%d' % (count, i, j)
177
                        # np.savetxt(hog_file, hog)
178

179
            if DEBUG_MOD is True:
180
                cv2.rectangle(frame_copy, (x, y),
181
                              (x + w, y + h), (0, 255, 0), 2)
182

183
        print(result)
184
        if 0 in result:
185
            print('yes')
186
        print('frame %d' % count)
187
        cv2.imshow('fmask', frame_copy)
188
        frame_and = cv2.bitwise_and(frame, frame, mask=fmask)
189

190
        count += 1
191
        k = cv2.waitKey(100) & 0xFF
192
        if k == 27:
193
            break
194

195
    end_time = time()
196
    total_time = end_time - start_time
197
    print('Total time: %d' % total_time)
198
    print(count)
199
    cap.release()
200
    cv2.destroyAllWindows()
201