diff --git a/python/classify.py b/python/classify.py
index 96a2a6b..22a7a55 100644
--- a/python/classify.py
+++ b/python/classify.py
@@ -7,6 +7,8 @@ Created on Sun Jul 23 16:04:33 2017
 
 import numpy
 import colorsys
+from scipy import misc
+from envi import envi
 
 #generate a 2D color class map using a stack of binary class images
 def classcolor2(C):
@@ -29,15 +31,72 @@ def classcolor2(C):
 #input: list of class image names
 #output: X x Y x C stack of binary mask images
 #example: image2class(("class_coll.bmp", "class_epith.bmp"))
+def image2class(*masks):
+    #get num of mask file names
+    num_masks = len(masks)
+
+    if num_masks == 0:
+        print("ERROR: mask filenames not provided")
+        print("Usage example: image2class(('class_coll.bmp', 'class_epith.bmp'))")
+        return
+
+    #load first mask to get dimensions -- assuming all masks have same dimensions
+    mask = misc.imread(masks[0], flatten=True).astype(numpy.bool)
+
+    mask_stack = numpy.zeros((mask.shape[0], mask.shape[1], num_masks), dtype=numpy.bool)
+    mask_stack[:,:,0] = mask
+
+    #load the rest of masks
+    for i in range(1, num_masks):
+        mask_stack[:,:,i] = misc.imread(masks[i], flatten=True).astype(numpy.bool)
+
+    return mask_stack
+
 
 #create a set of feature/target pairs for classification
 #input: envi file object, stack of class masks, list of class names
 #output: feature matrix (features x pixels), target matrix (1 x pixels)
 #example: generate_training(("class_coll.bmp", "class_epith.bmp"), (1, 2))
+def generate_training(filename, mask_stack):
+    #create envi file object
+    E = envi(filename)
+
+    # get number of classes
+    C = mask_stack.shape[2]
+
+    #get number of annotated pixels
+    num_pixels = 0
+    for i in range(0,C):
+        num_pixels += numpy.count_nonzero(mask_stack[:,:,i])
+
+    feature_matrix = numpy.zeros((E.header.bands, num_pixels), dtype=E.header.data_type)
+    target_matrix = numpy.zeros((1, num_pixels))
+
+    #load masks and append the corresponding pixels to feature matrix and labels to target_matrix
+    idx = 0
+    for i in range(0,C):
+        mask = E.loadmask(mask_stack[:,:,i])
+        feature_matrix[:, idx:idx + mask.shape[1]] = mask
+        target_matrix[idx:idx+mask.shape[1]] = (i+1)
+        idx += mask.shape[1]
+
+    return feature_matrix, target_matrix
 
 #create a class mask stack from an X x Y x C probability image
 #input: X x Y x C image giving the probability P(c |x,y)
 #output: X x Y x C binary class image
+def prob2class(prob_image):
+    class_image = numpy.zeros_like(prob_image)
+    #get nonzero indices
+    nnz_idx = numpy.transpose(numpy.nonzero(numpy.sum(prob_image, axis=0)))
+
+    #set pixel corresponding to max probability to 1
+    for idx in nnz_idx:
+        idx_max_prob = numpy.argmax(prob_image[:, idx[0], idx[1]])
+        class_image[idx_max_prob, idx[0], idx[1]] = 1
+
+    return class_image
+
 
 #create an ROC curve calculator
 #input: X x Y x C image giving the probability P(c | x,y)
diff --git a/python/envi.py b/python/envi.py
index f013b5b..c1fe886 100644
--- a/python/envi.py
+++ b/python/envi.py
@@ -139,9 +139,18 @@ class envi_header:
             #handle the complicated conditions
             elif l[li].startswith("description"):
                 desc = [l[li]]
-                while l[li].strip()[-1] != '}':
+                ''' 
+                while l[li].strip()[-1] != '}': #will fail if l[li].strip() is empty
+                    li += 1
+                    desc.append(l[li])
+                '''
+                while True:
+                    if l[li].strip():
+                       if  l[li].strip()[-1] == '}':
+                           break
                     li += 1
                     desc.append(l[li])
+
                 desc = ''.join(list(map(str.strip, desc)))           #strip all white space from the string list
                 i0 = desc.rindex('{')
                 i1 = desc.rindex('}')
diff --git a/python/example.py b/python/example.py
new file mode 100644
index 0000000..51ac6de
--- /dev/null
+++ b/python/example.py
@@ -0,0 +1,22 @@
+import numpy
+import classify
+import matplotlib.pyplot as plt
+from envi import envi
+
+mask_path = '/home/sberisha/data/masks/'
+mask_stack = classify.image2class(mask_path + "class_blood.png", mask_path + "class_coll.png", mask_path +  "class_epith.png",
+                         mask_path + "class_lymph.png", mask_path + "class_necrosis.png")
+
+color_image = classify.classcolor2(mask_stack)
+plt.imshow(color_image)
+
+data_path ='/home/sberisha/data/cnn/brc961-nfp8/envi/'
+
+feature_matrix, target_matrix = classify.generate_training(data_path + 'brc961-nfp8-project-br1003', mask_stack)
+
+prob_path = '/home/sberisha/data/'
+prob_envi= envi(prob_path + "cnn-response")
+prob_image = prob_envi.loadall()
+
+class_image = classify.prob2class(prob_image)
+plt.imshow(class_image[4,:,:])
\ No newline at end of file
--
libgit2 0.21.4