 
 (defn- image->observation-data
   "Create an observation from input."
-  [image datatype image-aug-fn]
+  [image datatype colorspace image-aug-fn]
   (patch/image->patch (if image-aug-fn
                         (image-aug-fn image)
                         image)
                       :datatype datatype
-                      :colorspace :gray))
+                      :colorspace colorspace))
 
 
 (defn- file->observation
   "Given a file, returns an observation map (an element of a dataset)."
-  [{:keys [class-name->index]} image-aug-fn datatype ^File file]
+  [{:keys [class-name->index]} image-aug-fn datatype colorspace ^File file]
   (try
-    {:data (image->observation-data (i/load-image file) datatype image-aug-fn)
+    {:data (image->observation-data (i/load-image file) datatype colorspace image-aug-fn)
      :labels (util/idx->one-hot (class-name->index (.. file getParentFile getName))
                                 (count (keys class-name->index)))}
     (catch Throwable _


(defn create-dataset-from-folder
  "Turns a folder of folders of png images into a dataset (a sequence of maps)."
  [folder-name class-mapping & {:keys [image-aug-fn datatype]
                                :or {datatype :float}}]
  (println "Building dataset from folder:" folder-name)
  (->> folder-name
       (io/as-file)

       (map (partial file->observation
                     class-mapping
                     image-aug-fn
                     datatype
                     colorspace))
       (remove nil?)))
