scikit-learn/examples/ensemble/plot_forest_multioutput.py

70 lines
2.0 KiB
Python
Raw Normal View History

2012-07-04 18:02:22 +08:00
"""
=========================================
Face completion with multi-output forests
=========================================
This example shows the use of multi-output forests to complete images.
The goal is to predict the lower half of a face given its upper half.
2012-07-04 18:04:42 +08:00
The first row of images shows true faces. The second half illustrates
how the forest completes the lower half of those faces.
2012-07-04 18:02:22 +08:00
"""
print __doc__
import numpy as np
import pylab as pl
from sklearn.datasets import fetch_olivetti_faces
from sklearn.ensemble import ExtraTreesRegressor
2012-07-04 19:10:28 +08:00
# Load the faces datasets
2012-07-04 18:02:22 +08:00
data = fetch_olivetti_faces()
targets = data.target
data = data.images.reshape((len(data.images), -1))
train = data[targets < 30]
2012-07-04 19:02:48 +08:00
test = data[targets >= 30] # Test on independent people
2012-07-04 18:02:22 +08:00
n_pixels = data.shape[1]
2012-07-04 19:02:48 +08:00
X_train = train[:, :int(0.5 * n_pixels)] # Upper half of the faces
Y_train = train[:, int(0.5 * n_pixels):] # Lower half of the faces
X_test = test[:, :int(0.5 * n_pixels)]
Y_test = test[:, int(0.5 * n_pixels):]
2012-07-04 18:02:22 +08:00
2012-07-04 19:10:28 +08:00
# Build a multi-output forest
2012-07-04 18:02:22 +08:00
forest = ExtraTreesRegressor(n_estimators=10,
max_features=32,
random_state=0)
forest.fit(X_train, Y_train)
Y_test_predict = forest.predict(X_test)
# Plot the completed faces
n_faces = 5
image_shape = (64, 64)
pl.figure(figsize=(2. * n_faces, 2.26 * 2))
pl.suptitle("Face completion with multi-output forests", size=16)
for i in xrange(1, 1 + n_faces):
face_id = np.random.randint(X_test.shape[0])
true_face = np.hstack((X_test[face_id], Y_test[face_id]))
completed_face = np.hstack((X_test[face_id], Y_test_predict[face_id]))
pl.subplot(2, n_faces, i)
pl.axis("off")
pl.imshow(true_face.reshape(image_shape),
cmap=pl.cm.gray,
interpolation="nearest")
pl.subplot(2, n_faces, n_faces + i)
pl.axis("off")
pl.imshow(completed_face.reshape(image_shape),
cmap=pl.cm.gray,
interpolation="nearest")
pl.show()