import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import fetch_olivetti_faces
from sklearn.utils.validation import check_random_state
from sklearn.ensemble import RandomForestRegressor
from sklearn.tree import DecisionTreeRegressor
#加载脸部数据
data = fetch_olivetti_faces()
targets = data.target
#划分训练、测试集
data = data.images.reshape((len(data.images), -1))
train = data[targets < 30]
test = data[targets >= 30]
#从测试集中随机找出5个用来作预测并且对其可视化
n_faces = 5
rng = check_random_state(4)
face_ids = rng.randint(test.shape[0], size=(n_faces, ))
test = test[face_ids, :]
#数据处理
n_pixels = data.shape[1]
X_train = train[:, :(n_pixels + 1) // 2] #上半部分脸
y_train = train[:, n_pixels // 2:] #下半部分脸
X_test = test[:, :(n_pixels + 1) // 2]
y_test = test[:, n_pixels // 2:]
#训练模型并预测
ESTIMATORS = {
"Random Forest": RandomForestRegressor(n_estimators=100, max_features=32, random_state=0),
"Decision Tree Regressor": DecisionTreeRegressor()
}
y_test_predict = dict()
for name, estimator in ESTIMATORS.items():
estimator.fit(X_train, y_train)
y_test_predict[name] = estimator.predict(X_test)
# 可视化补全的脸
image_shape = (64, 64) #图片像素
n_cols = 1 + len(ESTIMATORS)
plt.figure(figsize=(2. * n_cols, 2.26 * n_faces))
plt.suptitle("Face completion with RandomForests and DecisionTreeReression", size=8)
for i in range(n_faces):
true_face = np.hstack((X_test[i], y_test[i]))
if i:
sub = plt.subplot(n_faces, n_cols, i * n_cols + 1)
else:
sub = plt.subplot(n_faces, n_cols, i * n_cols + 1, title="true faces")
sub.axis("off")
sub.imshow(true_face.reshape(image_shape), cmap=plt.cm.gray, interpolation="nearest")
for j, est in enumerate(sorted(ESTIMATORS)):
completed_face = np.hstack((X_test[i], y_test_predict[est][i]))
if i:
sub = plt.subplot(n_faces, n_cols, i * n_cols + 2 + j)
else:
sub = plt.subplot(n_faces, n_cols, i * n_cols + 2 + j, title=est)
sub.axis("off")
sub.imshow(completed_face.reshape(image_shape), cmap=plt.cm.gray, interpolation="nearest")
plt.show()
