diff --git a/edgesTrain.m b/edgesTrain.m
index 9c5a4c3..19a24d4 100644
--- a/edgesTrain.m
+++ b/edgesTrain.m
@@ -164,13 +164,29 @@
 function trainTree( opts, stream, treeInd )
 % Train a single tree in forest model.
 
+% allowed image extensions
+imgExtWhitelist = {'jpg', 'jpeg', 'tiff', 'png'};
+
 % location of ground truth
 trnImgDir = [opts.bsdsDir '/images/train/'];
 trnDepDir = [opts.bsdsDir '/depth/train/'];
 trnGtDir = [opts.bsdsDir '/groundTruth/train/'];
-imgIds=dir(trnImgDir); imgIds=imgIds([imgIds.bytes]>0);
-imgIds={imgIds.name}; ext=imgIds{1}(end-2:end);
-nImgs=length(imgIds); for i=1:nImgs, imgIds{i}=imgIds{i}(1:end-4); end
+
+imgExt = cell(0); imgNames = cell(0);
+allFiles = dir(trnImgDir);
+for i=1:length(allFiles)
+    fullName = allFiles(i).name;
+    bytes = allFiles(i).bytes;
+    if bytes > 0 && ~ismember(fullName,{'.','..'})
+        [pathstr, name, ext] = fileparts([trnImgDir fullName]);
+        ext = ext(2:end);
+        if ismember(ext, imgExtWhitelist)
+            imgExt = [imgExt, {ext}];
+            imgNames = [imgNames, {name}];
+        end
+    end
+end
+nImgs = length(imgNames);
 
 % extract commonly used options
 imWidth=opts.imWidth; imRadius=imWidth/2;
@@ -199,9 +215,9 @@ function trainTree( opts, stream, treeInd )
 tid = ticStatus('Collecting data',30,1);
 for i = 1:nImgs
   % get image and compute channels
-  gt=load([trnGtDir imgIds{i} '.mat']); gt=gt.groundTruth;
-  I=imread([trnImgDir imgIds{i} '.' ext]); siz=size(I);
-  if(rgbd), D=single(imread([trnDepDir imgIds{i} '.png']))/1e4; end
+  gt=load([trnGtDir imgNames{i} '.mat']); gt=gt.groundTruth;
+  I=imread([trnImgDir imgNames{i} '.' imgExt{i}]); siz=size(I);
+  if(rgbd), D=single(imread([trnDepDir imgNames{i} '.png']))/1e4; end
   if(rgbd==1), I=D; elseif(rgbd==2), I=cat(3,single(I)/255,D); end
   p=zeros(1,4); p([2 4])=mod(4-mod(siz(1:2),4),4);
   if(any(p)), I=imPad(I,p,'symmetric'); end
@@ -298,3 +314,4 @@ function trainTree( opts, stream, treeInd )
 % optionally display different types of hs
 for i=1:0, figure(i); montage2(cell2array(segs(hs==i))); end
 end
+