基于BOW和SVM的目标检测

数据集制作

参照博客基于PIL批量修改图片JPG格式为PGM格式制作数据集。

源程序

import cv2
import numpy as np

data_path = "train_data/"
SAMPLES = 10


def path(cls, i):
    return "%s/%s%d.pgm" % (data_path, cls, i + 1)  # pgm格式


pos, neg = "pos-", "neg-"   # 类别

detect = cv2.xfeatures2d.SIFT_create()      # 提取关键点
extract = cv2.xfeatures2d.SIFT_create()     # 提取特征

flann_params = dict(algorithm=1, tree=5)    # FLANN匹配器
flann = cv2.FlannBasedMatcher(flann_params, {})

bow_kmeans_trainer = cv2.BOWKMeansTrainer(40)       # BOW训练器/提取器
extract_bow = cv2.BOWImgDescriptorExtractor(extract, flann)


def extract_sift(fn):
    im = cv2.imread(fn, 0)
    return extract.compute(im, detect.detect(im))[1]


for i in range(1, SAMPLES):
    bow_kmeans_trainer.add(extract_sift(path(pos, i)))
    bow_kmeans_trainer.add(extract_sift(path(neg, i)))

voc = bow_kmeans_trainer.cluster()  # 执行k-means分类并返回词汇
extract_bow.setVocabulary(voc)      # 从测试图像中提取描述符


def bow_feature(fn):
    im = cv2.imread(fn, 0)
    return extract_bow.compute(im, detect.detect(im))       # 返回基于BOW的描述符提取器计算得到的描述符


train_data, train_labels = [], []

for i in range(1, SAMPLES):
    train_data.extend(bow_feature(path(pos, i)))
    train_labels.append(1)
    train_data.extend(bow_feature(path(neg, i)))
    train_labels.append(-1)

svm = cv2.ml.SVM_create()       # SVM实例
svm.train(np.array(train_data), cv2.ml.ROW_SAMPLE, np.array(train_labels))


def predict(fn):
    f = bow_feature(fn)
    p = svm.predict(f)
    print(fn, "\t", p[1][0][0])
    return p


bulldozer = "bulldozer.jpg"
image = cv2.imread(bulldozer)
image_predict = predict(bulldozer)

font = cv2.FONT_HERSHEY_COMPLEX

if (image_predict[1][0][0] == 1.0):
    cv2.putText(image, 'Bulldozer Detected', (10, 30), font, 1, (0, 255, 0), 2, cv2.LINE_AA)
else:
    cv2.putText(image, 'Bulldozer Not Detected', (10, 30), font, 1, (0, 0, 255), 2, cv2.LINE_AA)

cv2.imshow('Detect Success for %s' % bulldozer[:-4], image)
cv2.waitKey(0)
cv2.destroyAllWindows()

（最近更新：2020年01月10日）