add support

kyscg · Sep 10, 2018 · 45c27e8 · 45c27e8
1 parent b57a23c
commit 45c27e8
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/deepchem/data/data_loader.py b/deepchem/data/data_loader.py
@@ -349,7 +349,7 @@ def __init__(self, tasks=None):
       tasks = []
     self.tasks = tasks
 
-  def featurize(self, input_files, in_memory=True):
+  def featurize(self, input_files, labels=None, in_memory=True):
     """Featurizes image files.
 
     Parameters
@@ -409,8 +409,9 @@ def featurize(self, input_files, in_memory=True):
         raise ValueError("Unsupported image filetype for %s" % image_file)
     images = np.array(images)
     if in_memory:
-      return NumpyDataset(images)
+      return NumpyDataset(images, y=labels)
     else:
       # from_numpy currently requires labels. Make dummy labels
-      labels = np.zeros((len(images), 1))
+      if labels is None:
+        labels = np.zeros((len(images), 1))
       return DiskDataset.from_numpy(images, labels)