1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98
| import os import warnings import numpy as np import cv2 as cv import hmmlearn.hmm as hl warnings.filterwarnings('ignore', category=DeprecationWarning) np.seterr(all='ignore') def search_objects(directory): directory = os.path.normpath(directory) if not os.path.isdir(directory): raise IOError("The directory '" + directory + "' doesn't exist!") objects = {} for curdir, subdirs, files in os.walk( directory): for jpeg in (file for file in files if file.endswith('.jpg')): path = os.path.join(curdir, jpeg) label = path.split(os.path.sep)[-2] if label not in objects: objects[label] = [] objects[label].append(path) return objects train_objects = search_objects( '../data/objects/training') train_x, train_y = [], [] for label, filenames in train_objects.items(): descs = np.array([]) for filename in filenames: image = cv.imread(filename) gray = cv.cvtColor(image, cv.COLOR_BGR2GRAY) h, w = gray.shape[:2] f = 200 / min(h, w) gray = cv.resize(gray, None, fx=f, fy=f) star = cv.xfeatures2d.StarDetector_create() keypoints = star.detect(gray) sift = cv.xfeatures2d.SIFT_create() _, desc = sift.compute(gray, keypoints) if len(descs) == 0: descs = desc else: descs = np.append(descs, desc, axis=0) train_x.append(descs) train_y.append(label) models = {} for descs, label in zip(train_x, train_y): model = hl.GaussianHMM( n_components=4, covariance_type='diag', n_iter=1000) models[label] = model.fit(descs) test_objects = search_objects( '../data/objects/testing') test_x, test_y, test_z = [], [], [] for label, filenames in test_objects.items(): test_z.append([]) descs = np.array([]) for filename in filenames: image = cv.imread(filename) test_z[-1].append(image) gray = cv.cvtColor(image, cv.COLOR_BGR2GRAY) h, w = gray.shape[:2] f = 200 / min(h, w) gray = cv.resize(gray, None, fx=f, fy=f) star = cv.xfeatures2d.StarDetector_create() keypoints = star.detect(gray) sift = cv.xfeatures2d.SIFT_create() _, desc = sift.compute(gray, keypoints) if len(descs) == 0: descs = desc else: descs = np.append(descs, desc, axis=0) test_x.append(descs) test_y.append(label) pred_test_y = [] for descs in test_x: best_score, best_label = None, None for label, model in models.items(): score = model.score(descs) if (best_score is None) or \ (best_score < score): best_score, best_label = \ score, label pred_test_y.append(best_label) i = 0 for label, pred_label, images in zip( test_y, pred_test_y, test_z): for image in images: i += 1 cv.imshow('{} - {} {} {}'.format( i, label, '==' if label == pred_label else '!=', pred_label), image) cv.waitKey()
|