参照博客基于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日)