diff --git a/src/napari_annotatorj/_dock_widget.py b/src/napari_annotatorj/_dock_widget.py
index 79ffedd..6fbd4f4 100644
--- a/src/napari_annotatorj/_dock_widget.py
+++ b/src/napari_annotatorj/_dock_widget.py
@@ -9751,8 +9751,10 @@ def prepTraining(self):
             width,height=self.prepOrigFolder(self.originalFolder)
 
             ms=0
+            maskDir=os.path.join(self.trainDataFolder,'unet_masks')
             for mask in self.curExpList:
                 if os.path.isfile(os.path.join(self.annotationFolder,mask)):
+                    from shutil import copyfile
                     copyfile(os.path.join(self.annotationFolder,mask),os.path.join(maskDir,mask))
                     ms+=1
             if ms>0:
@@ -9867,6 +9869,7 @@ def prepOrigFolder(self,origFolder):
 
             ims=0
             img=None
+            imDir=os.path.join(self.trainDataFolder,'images')
             for im in self.curFileList:
                 if os.path.isfile(os.path.join(origFolder,im)):
                     from shutil import copyfile