image.clj

(ns scottdw.image
  (:import
   [java.awt Color Image Dimension]
   [java.awt.image BufferedImage]
   [java.io File]
   [javax.imageio ImageIO]
   [javax.swing JFrame JPanel])
  (require [scottdw.stats :as stats]))

(defn read-image [^String filename]
  (ImageIO/read (File. filename)))

(defn create-image [w h p-fn]
  (let [imgr (BufferedImage. w
                             h
                             BufferedImage/TYPE_INT_ARGB)]
    (doseq [x (range w) y (range h)]
      (.setRGB imgr x y (p-fn x y)))
    imgr))

(defn viewer-frame [^String name ^BufferedImage img]
  (let [i (atom img)
        d (Dimension. (.getWidth img) (.getHeight img))
        p (proxy [JPanel] []
            (getPreferredSize [] d)
            (getMaximumSize [] d)
            (getMinimumSize [] d)
            (paintComponent [^Graphics g]
              (proxy-super paintComponent g)
              (.drawImage g @i 0 0 nil)))
        f (JFrame. name)]
    (-> f .getContentPane (.add p))
    (doto f
      (.pack)
      (.setVisible true))
    [i #(.repaint p)]))

(defn update-img [[ia pm] img]
  (swap! ia (constantly img))
  (pm))

(defrecord color [^String str  ^int rgb ^double red ^double grn ^double blu ^double hue ^double sat ^double lum ^double y ^double cb ^double cr])

(let [no-alpha-mask (unchecked-int 0x00FFFFFF)
      alpha-mask (unchecked-int 0xFF000000)
      red-mask (unchecked-int 0x00FF0000)
      green-mask (unchecked-int 0x0000FF00)
      blue-mask (unchecked-int 0x000000FF)
      to-255 (comp #(Math/round ^float %) (partial * 255))]

  (defn to-agrey [v]
    (let [vi (to-255 v)]
      (unchecked-int (bit-or alpha-mask
                             (bit-shift-left vi 16)
                             (bit-shift-left vi 8)
                             vi))))

  (defn get-red ^double [rgb]
    (/ (bit-shift-right (bit-and red-mask rgb) 16) 255))

  (defn get-green ^double [rgb]
    (/ (bit-shift-right (bit-and green-mask rgb) 8) 255))

  (defn get-blue ^double [rgb]
    (/ (bit-and blue-mask rgb) 255))

  (defn get-hue ^double [^double r ^double g ^double b]
    (let [M (Math/max r (Math/max g b))
          m (Math/min r (Math/min g b))
          C (- M m)]
      (/ (cond
          (== C 0.0) Double/NaN
          (== M r) (mod (/ (- g b) C) 6.0)
          (== M g) (+ (/ (- b r) C) 2.0)
          (== M b) (+ (/ (- r g) C) 4.0))
         6.0)))

  (defn get-lum ^double [^double r ^double g ^double b]
    (let [M (Math/max r (Math/max g b))
          m (Math/min r (Math/min g b))]
      (* 0.5 (+ M m))))

  (defn get-sat ^double [^double r ^double g ^double b]
    (let [M (Math/max r (Math/max g b))
          m (Math/min r (Math/min g b))
          C (- M m)
          L (* 0.5 (+ M m))]
          (if (== C 0.0)
            0.0
            (/ C (- 1.0 (Math/abs (- (* 2.0 L) 1.0)))))))

  (defn get-y ^double [^double r ^double g ^double b]
    (+ (* 0.299 r)
       (* 0.587 g)
       (* 0.114 b)))

  (defn get-cb ^double [^double r ^double g ^double b]
    (+ 0.5
       (* r -0.168736)
       (* g -0.331264)
       (* b 0.5)))

  (defn get-cr ^double [^double r ^double g ^double b]
    (+ 0.5
       (* r 0.5)
       (* g -0.418688)
       (* b -0.081312)))

  (defn get-r-from-ycr ^double [^double y ^double cr]
    (+ y (* 1.402 (- cr 0.5))))

  (defn get-g-from-ycbcr ^double [^double y ^double cb ^double cr]
    (+ y
       (* -0.344136 (- cb 0.5))
       (* -0.714136 (- cr 0.5))))

  (defn get-b-from-ycb ^double [^double y ^double cb]
    (+ y (* 1.772 (- cb 0.5))))

  (defn to-hex-color-string [argb]
    (format "#%06X" (bit-and no-alpha-mask argb)))

  (defn to-rgb [r g b]
    (unchecked-int (bit-or alpha-mask
                           (bit-shift-left r 16)
                           (bit-shift-left g 8)
                           b)))

  (defn ycbcr-to-rgb [y cb cr]
    (to-rgb (to-255 (get-r-from-ycr y cr))
            (to-255 (get-g-from-ycbcr y cb cr))
            (to-255 (get-b-from-ycb y cb))))

  (defn create-color [rgb]
    (let [hcs (to-hex-color-string rgb)
          r (get-red rgb)
          g (get-green rgb)
          b (get-blue rgb)
          h (get-hue r g b)
          s (get-sat r g b)
          l (get-lum r g b)
          y (get-y r g b)
          cb (get-cb r g b)
          cr (get-cr r g b)]
      (->color hcs rgb r g b h s l y cb cr))))

(defn hsl-to-rgb [hue sat lum]
  (let [h  (* 360 hue)
        s  sat
        l  lum
        c  (* s (- 1 (Math/abs (double (- (* 2 l) 1)))))
        m  (- l (* c 0.5))
        h_ (if (Double/isNaN h) 0 (/ h 60))
        x  (* c (- 1 (Math/abs (double (dec (mod h_ 2))))))
        rgb_ (cond
              (Double/isNaN h) [0 0 0]
              (< h_ 1) [c x 0]
              (< h_ 2) [x c 0]
              (< h_ 3) [0 c x]
              (< h_ 4) [0 x c]
              (< h_ 5) [x 0 c]
              :else    [c 0 x])
        [r g b] (map (comp #(Math/round ^float %) (partial * 255) (partial + m)) rgb_)]
    (to-rgb r g b)))

(defn pixels [^BufferedImage img]
  (for [x (range (.getWidth img)) y (range (.getHeight img))] (.getRGB img x y)))

(defn create-color-map [^Image img]
  (let [ps (into #{} (pixels img))]
    (binding [*unchecked-math* true]
      (persistent! (reduce #(assoc! %1 %2 (create-color %2)) (transient {}) ps)))))

(defn hue-centre [img]
  (let [cm (create-color-map img)]
    (let [[x y] (reduce (fn [[x1 y1] [x2 y2]] [(+ x1 x2) (+ y1 y2)]) (map (juxt #(Math/cos %) #(Math/sin %)) (remove #(Double/isNaN %) (map (comp :hue cm) (pixels img)))))]
      (Math/atan2 y x))))

(defn rgb-mean [^BufferedImage img]
  (let [cm (create-color-map img)
        c (* (.getWidth img) (.getHeight img))]
    (map (comp double #(/ % c)) (reduce (fn [[r g b] {:keys [red grn blu]}] [(+ r red) (+ g grn) (+ b blu)]) [0 0 0] (map cm (pixels img))))))

(defn cbcr-mean [^BufferedImage img]
  (let [cm (create-color-map img)
        c (* (.getWidth img) (.getHeight img))]
    (map (comp double #(/ % c)) (reduce (fn [[cb1 cr1] {:keys [cb cr]}] [(+ cb1 cb) (+ cr1 cr)]) [0 0] (map cm (pixels img))))))

(defn bucket-fn [n]
  (fn [x] (double (/ (Math/round ^float (* n x)) n))))

#_(defn img-color-centroid [img]
  (let [tau (* 2 Math/PI)
        bucket-fn (bucket-fn 360)
        cfm (frequencies (map (comp (partial * tau) bucket-fn :hue (memoize create-color)) (pixels img)))
        cs (keys cfm)
        cd-fn (fn [cm cM]
                (let [x1 (Math/cos cm) y1 (Math/sin cm)
                      x2 (Math/cos cM) y2 (Math/sin cm)]
                  (Math/sqrt (+ (Math/pow (- x1 x2) 2)
                                (Math/pow (- y1 y2) 2)))))
        cdm (zipmap cs (map #(zipmap cs (map (partial cd-fn %) cs)) cs))
        total-score-fn (fn [c] (reduce #(+ %1 (* ((cdm c) %2) (cfm %2))) 0 cs))
        median-score-fn (fn [c] (stats/median (reduce into [] (map (fn [[k d]] (repeat d (cfm k))) (cdm c)))))
        tm (zipmap cs (map median-score-fn cs))]
    (double (/ (first (reduce #(min-key second %1 %2) tm)) tau))))

(defn huegram-frame [^String name img]
  (let [black (unchecked-int 0xFF000000)
        i (atom img)
        r 200
        d (* r 2)
        dim (Dimension. d d)
        p (proxy [JPanel] []
            (getPreferredSize [] dim)
            (getMaximumSize [] dim)
            (getMinimumSize [] dim)
            (paintComponent [^Graphics g]
              (let [img @i
                    hf (frequencies (map (comp #(Math/round ^float %)
                                               (partial * 360)
                                               :hue
                                               (memoize create-color))
                                         (pixels img)))
                    sf (/ 100 (reduce max (map second hf)))]
                (proxy-super paintComponent g)
                (.drawImage g
                            (create-image d d
                                          (fn [x y]
                                            (let [black (unchecked-int 0xFF000000)
                                                  tx (- x r)
                                                  ty (- r y)
                                                  d (Math/sqrt (+ (* tx tx) (* ty ty)))
                                                  th (Math/toDegrees (Math/atan2 ty tx))
                                                  h (Math/round (float (if (neg? th) (+ 360 th) th)))]
                                              (cond
                                               (< d 50) black
                                               (and (>= d 50) (< d 60)) (hsl-to-rgb (/ h 360) 1 0.5)
                                               (and (>= d 80) (< d (+ 80 (* sf (or (hf h) 0))))) (hsl-to-rgb (/ h 360) 1 0.5)
                                               :else black))))
                            0 0 nil))))
        f (JFrame. name)]
    (-> f .getContentPane (.add p))
    (doto f
      (.pack)
      (.setVisible true))
    [i #(.repaint p)]))

(defn new-color-mapped-image [^BufferedImage img color-map-fn]
  (let [cm (create-color-map img)
        ncm (zipmap (keys cm) (map color-map-fn (keys cm)))]
    (create-image (.getWidth img) (.getHeight img) #(ncm (.getRGB img %1 %2)))))

(defn dither-to-bw [^BufferedImage img]
  (let [black (int -16777216)
        white (int -1)
        ef (double (/ 7 16))
        swf (double (/ 3 16))
        sf (double (/ 5 16))
        sef (double (/ 1 16))
        w (.getWidth img)
        bw (inc w)
        h (.getHeight img)
        buffer (make-array Double/TYPE (* 2 bw))
        ni (BufferedImage. w h BufferedImage/TYPE_INT_ARGB)
        cm (create-color-map img)
        lum-map (zipmap (keys cm) (map :y (vals cm)))]
    (loop [y 0 x 0 cr-offset 0 nr-offset bw]
      (cond
       (= y h) ni
       (= x w) (do
                 (java.util.Arrays/fill ^doubles buffer cr-offset (+ cr-offset bw) 0.0)
                 (recur (inc y) 0 nr-offset cr-offset))
       :else (let [xi (+ x cr-offset)
                   ei (inc xi)
                   si (+ x nr-offset)
                   swi (dec si)
                   sei (inc si)
                   l (+ (aget ^doubles buffer xi) (lum-map (.getRGB img x y)))
                   e (- l (if (> l 0.5) 1 0))]
               (.setRGB ni x y (if (> l 0.5) white black))
               (aset ^doubles buffer ei (+ (aget ^doubles buffer ei) (* e ef)))
               (aset ^doubles buffer sei (+ (aget ^doubles buffer sei) (* e sef)))
               (aset ^doubles buffer si (+ (aget ^doubles buffer si) (* e sf)))
               (when (pos? x)
                 (aset ^doubles buffer swi (+ (aget ^doubles buffer swi) (* e swf))))
               (recur y (inc x) cr-offset nr-offset))))))

(defn mean-diff [^BufferedImage img]
  (let [w (.getWidth img)
        h (.getHeight img)
        ni (BufferedImage. w h BufferedImage/TYPE_INT_ARGB)
        cm (create-color-map img)
        lum-map (zipmap (keys cm) (map :y (vals cm)))]
    (loop [y (int 0) x (int 0)]
      (cond
       (= y h) ni
       (= x w) (recur (inc y) 0)
       :else (let [px (dec x)
                   nx (inc x)
                   py (dec y)
                   ny (inc y)
                   p (.getRGB img x y)
                   ps (conj (cond
                             (and (zero? x) (zero? y)) [(.getRGB img nx y) (.getRGB img x ny) (.getRGB img nx ny)]
                             (and (= w nx) (zero? y)) [(.getRGB img px y) (.getRGB img px ny) (.getRGB img x ny)]
                             (and (zero? x) (= h ny)) [(.getRGB img x py) (.getRGB img nx py) (.getRGB img nx y)]
                             (and (= w nx) (= h ny)) [(.getRGB img px py) (.getRGB img x py) (.getRGB img px y)]
                             (zero? x) [(.getRGB img x py) (.getRGB img nx py) (.getRGB img nx y) (.getRGB img x ny) (.getRGB img nx ny)]
                             (= w nx) [(.getRGB img px py) (.getRGB img x py) (.getRGB img px y) (.getRGB img px ny) (.getRGB img x ny)]
                             (zero? y) [(.getRGB img px y) (.getRGB img nx y) (.getRGB img px ny) (.getRGB img x ny) (.getRGB img nx ny)]
                             (= h ny) [(.getRGB img px py) (.getRGB img x py) (.getRGB img nx py) (.getRGB img px y) (.getRGB img nx y)]
                             :else [(.getRGB img px py) (.getRGB img x py) (.getRGB img nx py)
                                    (.getRGB img px y) (.getRGB img nx y)
                                    (.getRGB img px ny) (.getRGB img x ny) (.getRGB img nx ny)])
                            p)
                   u (/ (reduce + (map lum-map ps)) (count ps))
                   ]
               (do
                 (.setRGB ni x y (to-agrey (- 1 (Math/abs ^double (- (lum-map p) u)))))
                 (recur y (inc x))))))))

(defn xy-clamp-to-index-fn [w h]
  (let [dw (dec w)
        dh (dec h)]
    (fn [x y]
      (let [x (cond (neg? x) 0
                    (> x dw) dw
                    :else x)
            y (cond (neg? y) 0
                    (> y dh) dh
                    :else y)]
        (+ x (* y w))))))

(defn img-to-component-array [^BufferedImage img component-key]
  (let [get-component (comp component-key (create-color-map img))
        w (.getWidth img)
        h (.getHeight img)
        index (xy-clamp-to-index-fn w h)]
    (into-array Double/TYPE (for [y (range h) x (range w)] (get-component (.getRGB img x y))))))

(defn lum-array-to-img [w h ^doubles lum-array]
  (let [index (xy-clamp-to-index-fn w h)]
    (create-image w h (fn [x y] (to-agrey (aget ^doubles lum-array (index x y)))))))

(defn convolve-extend [input iw ih kernel kw kh]
  (let [iindex (xy-clamp-to-index-fn iw ih)
        kindex (xy-clamp-to-index-fn kw kh)]
    (into-array Double/TYPE
                (for [y (range ih) x (range iw)]
                  (reduce + (for [ky (range kh) kx (range kw)]
                              (* (aget ^doubles input (iindex (dec (+ x kx)) (dec (+ y ky))))
                                 (aget ^doubles kernel (kindex kx kh)))))))))

(let [gauss-kernel (into-array Double/TYPE
                                 (map (comp double #(/ % 273))
                                      [1  4  7  4 1
                                       4 16 26 16 4
                                       7 26 41 26 7
                                       4 16 26 16 4
                                       1  4  7  4 1]))
        kw 5 kh 5]
    (defn gauss-blur [^BufferedImage img]
      (let [w (.getWidth img)
            h (.getHeight img)
            lum-array (img-to-component-array img :y)
            out-array (convolve-extend lum-array w h gauss-kernel 5 5)]
        out-array
       #_(lum-array-to-img w h out-array))))