Testing gradients with custom gradient fns and tensors.

src/tech/v2/tensor/gradients.clj

@@ -11,7 +11,9 @@
             [tech.parallel.for :as pfor]
             [tech.libs.buffered-image :as bufimg]
             [clojure.tools.logging :as log])
-  (:import [java.awt.image BufferedImage]))
+  (:import [java.awt.image BufferedImage]
+           [tech.v2.tensor ByteTensorReader]
+           [clojure.lang IFn]))
 
 
 (def gradient-map (delay (-> (io/resource "gradients.edn")
@@ -30,14 +32,32 @@
 
 
 (defn apply-gradient
-  "Apply a gradient to a tensor returning a new tensor.  If data-min, data-max aren't
+  "Apply a gradient to a tensor returning an image.  If data-min, data-max aren't
   provided they are found in the data.  A buffered image is returned.
-  gradient-names are (keys gradient-map).  Gradients themselves are taken from
-  https://reference.wolfram.com/language/guide/ColorSchemes.html"
+  src-tens - Source tensor whose shape determines the shape of the final image.
+  gradient-name -  may be a keyword, in which it must be a key in @gradient-map and
+      these gradients come from:
+      https://reference.wolfram.com/language/guide/ColorSchemes.html.
+    gradient-name may be a tensor of dimensions [n 3].
+    gradient-name may be a function that takes a value from 0-1 and returns a tuple
+    of length 3.
+
+  Additional arguments:
+  :data-min :data-max - If provided then the data isn't scanned for min and max.  If min
+    is equal to 0 and max is equal to 1.0 then the data doesn't need to be normalized.
+    data ranges are clamped to min and max.
+  :alpha? - If true, an image with an alpha channel is returned.  This is useful for
+    when your data has NAN or INFs as in that case the returned image is transparent
+    in those sections.
+  :check-invalid? - If true then the data is scanned for NAN or INF's.  Used in
+    conjunction with :alpha?
+  :invert-gradient? - When true, reverses the provided gradient."
   [src-tens gradient-name & {:keys [data-min data-max
                                     alpha?
                                     check-invalid?
-                                    invert-gradient?]}]
+                                    invert-gradient?
+                                    gradient-default-n]
+                             :or {gradient-default-n 200}}]
   (let [img-shape (dtype/shape src-tens)
         n-pixels (dtype/ecount src-tens)
         valid-indexes (when check-invalid?
@@ -91,17 +111,39 @@ This leads to ambiguous results as pixels not written to will be black but not t
                    res-tens)
         res-tens (tens-typecast/datatype->tensor-writer
                   :uint8 res-tens)
-        src-gradient-info (get @gradient-map gradient-name)
-        _ (when-not src-gradient-info
-            (throw (Exception. (format "Failed to find gradient %s"
-                                       gradient-name))))
-        gradient-tens @gradient-tens
+        n-gradient-increments (long gradient-default-n)
+        gradient-line
+        (cond
+          (keyword? gradient-name)
+          (let [src-gradient-info (get @gradient-map gradient-name)
+                _ (when-not src-gradient-info
+                    (throw (Exception.
+                            (format "Failed to find gradient %s"
+                                    gradient-name))))
+                gradient-tens @gradient-tens]
+            (dtt/select gradient-tens
+                        (:tensor-index src-gradient-info)
+                        :all :all))
+          (dtt/tensor? gradient-name)
+          gradient-name
+          (instance? IFn gradient-name)
+          (dtt/->tensor
+           (->> (range n-gradient-increments)
+                (map (fn [idx]
+                       (let [p-val (/ (double idx)
+                                      (double n-gradient-increments))
+                             grad-val (gradient-name p-val)]
+                         (when-not (= 3 (count grad-val))
+                           (throw (Exception. (format
+                                               "Gradient fns must return bgr tuples:
+function returned: %s"
+                                               grad-val))))
+                         grad-val))))))
+        n-gradient-increments (long (first (dtype/shape gradient-line)))
         gradient-line (tens-typecast/datatype->tensor-reader
                        :uint8
-                       (dtt/select gradient-tens
-                                   (:tensor-index src-gradient-info)
-                                   :all :all))
-        line-last-idx (double (dec (long (second (dtype/shape gradient-tens)))))
+                       gradient-line)
+        line-last-idx (double (dec n-gradient-increments))
         n-pixels (long (if valid-indexes
                          (dtype/ecount valid-indexes)
                          n-pixels))]
@@ -147,4 +189,26 @@ This leads to ambiguous results as pixels not written to will be black but not t
                                   :alpha? true
                                   :check-invalid? true)
                   "PNG"
-                  (format "gradient-demo/%s-nan.png" (name grad-name)))))
+                  (format "gradient-demo/%s-nan.png" (name grad-name))))
+
+  (def custom-gradient-tens (dtt/->tensor
+                             (->> (range 100)
+                                  (map (fn [idx]
+                                         (let [p-value (/ (double idx)
+                                                          (double 100))]
+                                           [(* 255 p-value) 0 (* (- 1.0 p-value)
+                                                                 255)]))))))
+
+  (bufimg/save! (apply-gradient test-src-tens custom-gradient-tens)
+                "PNG"
+                "gradient-demo/custom-tensor-gradient.png")
+
+  (defn custom-gradient-fn
+    [^double p-value]
+    (let [one-m-p (- 1.0 p-value)]
+      [(* 255 one-m-p) (* 255 p-value) (* 255 one-m-p)]))
+
+  (bufimg/save! (apply-gradient test-src-tens custom-gradient-fn)
+                "PNG"
+                "gradient-demo/custom-ifn-gradient.png")
+  )