--- a/pylearn/datasets/shapeset1.py	Sat Jan 17 17:38:17 2009 -0500
+++ b/pylearn/datasets/shapeset1.py	Sat Jan 17 19:06:21 2009 -0500
@@ -9,6 +9,7 @@
 
 from ..io.amat import AMat
 from .config import data_root
+from .dataset import Dataset
 
 def _head(path, n):
     dat = AMat(path=path, head=n)
@@ -25,10 +26,11 @@
 def head_train(n=10000):
     """Load the first Shapeset1 training examples.
 
-    Returns two matrices: x, y.  x has N rows of 1024 columns.  Each row of x represents the
-    32x32 grey-scale pixels in raster order.  y is a vector of N integers.  Each element y[i]
-    is the label of the i'th row of x.
-    
+    Returns two matrices: x, y.
+    x has N rows of 1024 columns.
+    Each row of x represents the 32x32 grey-scale pixels in raster order.
+    y is a vector of N integers between 0 and 2.
+    Each element y[i] is the label of the i'th row of x.
     """
     path = os.path.join(data_root(), 'shapeset1','shapeset1_1cspo_2_3.10000.train.shape.amat')
     return _head(path, n)
@@ -36,10 +38,11 @@
 def head_valid(n=5000):
     """Load the first Shapeset1 validation examples.
 
-    Returns two matrices: x, y.  x has N rows of 1024 columns.  Each row of x represents the
-    32x32 grey-scale pixels in raster order.  y is a vector of N integers.  Each element y[i]
-    is the label of the i'th row of x.
-    
+    Returns two matrices: x, y.
+    x has N rows of 1024 columns.
+    Each row of x represents the 32x32 grey-scale pixels in raster order.
+    y is a vector of N integers between 0 and 2.
+    Each element y[i] is the label of the i'th row of x.
     """
     path = os.path.join(data_root(), 'shapeset1','shapeset1_1cspo_2_3.5000.valid.shape.amat')
     return _head(path, n)
@@ -47,15 +50,28 @@
 def head_test(n=5000):
     """Load the first Shapeset1 testing examples.
 
-    Returns two matrices: x, y.  x has N rows of 1024 columns.  Each row of x represents the
-    32x32 grey-scale pixels in raster order.  y is a vector of N integers.  Each element y[i]
-    is the label of the i'th row of x.
-    
+    Returns two matrices: x, y.
+    x has N rows of 1024 columns.
+    Each row of x represents the 32x32 grey-scale pixels in raster order.
+    y is a vector of N integers between 0 and 2.
+    Each element y[i] is the label of the i'th row of x.
     """
     path = os.path.join(data_root(), 'shapeset1','shapeset1_1cspo_2_3.5000.test.shape.amat')
     return _head(path, n)
 
 def train_valid_test(ntrain=10000, nvalid=5000, ntest=5000):
-    return head_train(n=ntrain), head_valid(n=nvalid), head_test(n=ntest)
+    train_x, train_y = head_train(n=ntrain)
+    valid_x, valid_y = head_valid(n=nvalid)
+    test_x,  test_y  = head_test(n=test)
+
+    rval = Dataset()
+    rval.train = Dataset.Obj(x = train_x, y = train_y)
+    rval.valid = Dataset.Obj(x = valid_x, y = valid_y)
+    rval.test  = Dataset.Obj(x = test_x,  y = test_y)
+
+    rval.n_classes = 3
+    rval.img_shape = (32, 32)
+
+    return rval