;; Facial Expression Recognition Demo
;; Convolutional Network on the JAFFE dataset
;; 426 images, 213 and 213 reversed horizontally
;; Patrick Winters
;;
;; 7 labels, 0-6
;; HA - 0
;; SU - 1
;; NE - 2
;; AN - 3
;; SA - 4
;; DI - 5
;; FE - 6
;;
;; 426 images with any distribution for train and test
;; (build-databases-jaffe 319 107) ex
;;

(libload "gblearn2/net-lenet5")
(libload "gblearn2/gb-trainers")
(libload "gblearn2/gb-meters")
(libload "gblearn2/demos/dsource-mnist")
(libload "libimage/image-io")
(libload "libimage/ubimage")
(libload "libimage/rgbaimage")
(libload "libidx/idx-convol")
(libload "show-image")

(when (not *jaffe-dir*)
  (setq *jaffe-dir* "jaffedata/"))
(when (not *friends-dir*)
  (setq *friends-dir* "jaffetest/"))
(when (not *net-size*)
  (setq *net-size* 64))

;; size of the database is 426 images
(when (not *train-size*)
  (setq *train-size* 426))



;; open a new window and show the image and classification
;; returns the classification result
(de show-classify (file)
 (printf "label, output") 
)

;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;
;;; The following is for reading and setting up the
;;; datasource.  The network should be able to use the
;;; globally defined database variables after running
;;; (build-databases-jaffe (trsize tesize))
;;;
;;; trainingdb - dsource-idx31-narrow
;;;   and
;;; testingdb  - dsource-idx31-narrow
;;;
;;; trimages - matrix trsizex64x64
;;; trlabels - matrix trsize
;;; teimages - matrix tesizex64x64
;;; telabels - matrix tesize
;;;
;;;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(de build-database-dir (ddir)
  (let ((examples (read-dir ddir)))
    ;;make and return the example dsource
    (new dsource-idx3l-narrow
             (new dsource-mnist
                  (read-image-matrix examples ddir)
                  (read-label-matrix examples)
		  *net-size* *net-size* 0 0.01)
             (length examples) 0)
    )
)

;;Build friends test dataset
(de build-database-friends ()
    (let ((friends (read-dir *friends-dir*)))
     ;;now make the friends test dataset
	(setq frimages (read-image-matrix friends *friends-dir*))
	(setq frlabels (read-label-matrix friends))
	(setq friendsdb
        (new dsource-idx3l-narrow
             (new dsource-mnist
                  frimages
                  frlabels
		  *net-size* *net-size* 0 0.01)
             (length friends) 0))

))

;; Build training and test datasets
(de build-databases-jaffe (trsize tesize)
  ;; limited number of images, error if too big
  (if (> (+ trsize tesize) *train-size*) (error "not enough images")) 
  (let ((files (read-dir *jaffe-dir*))
	(tral 0) (tesl 0))
    ;; read in training set
;;    (printf "# of files =%d\n" (length files))
      (let ((trfiles (shuf (nfirst trsize files))))
	(setq trimages (read-image-matrix trfiles *jaffe-dir*))
	(setq trlabels (read-label-matrix trfiles))
	(setq trainingdb
        (new dsource-idx3l-narrow
             (new dsource-mnist
                  trimages
		  trlabels
		  *net-size* *net-size* 0 0.01)
             trsize 0)))
      (let ((tefiles (shuf (nfirst tesize (reverse files)))))
	(setq teimages (read-image-matrix tefiles *jaffe-dir*))
	(setq telabels (read-label-matrix tefiles))
	(setq testingdb
        (new dsource-idx3l-narrow
             (new dsource-mnist
                  teimages
                  telabels
		  *net-size* *net-size* 0 0.01)
             tesize 0)))

;;    (printf "\nidx-dim trainset 0 = %d\n" (idx-dim trainimgs 0))
;;    (printf "\nidx-dim trainset 0 = %d\n" (idx-dim testimgs 0))
    ))

;; read a list of files into a matrix lenx256x256
;; should return a list, a matrix of images, and matrix of labels
(de read-image-matrix (dfiles ddir)
  (let ((m (ubyte-matrix (length dfiles) *net-size* *net-size*)))
    (idx-bloop ((image m))
     (let ((f (image-read-ubim-fixed (concat-fname ddir (car dfiles)))))
       ;; mode 0, fit the image into width x height
       ;; high pass on 132x132 should return 128x128
       (idx-copy (ubim-high-pass (ubim-resize f (+ 4 *net-size*) (+ 4 *net-size*) 0)) image))
     (setq dfiles (cdr dfiles)))
   m))

;; read a list of files into a matrix of labels 0-5
(de read-label-matrix (dfiles)
  (let ((m (ubyte-matrix (length dfiles))))
    (idx-bloop ((dlabel m))
      (dlabel (read-label (car dfiles)))
      (setq dfiles (cdr dfiles)))
    (idx-bloop ((dlabel m))
      (printf "label - %d\n" (dlabel)))

  m))



;; parse a label from a filename string
;; encodings of expressions happen here.
;; kind of reliant on my poor regular expression skills.
;; fear is intentionally last because it made be
;; difficult to classify.
;;  ;;(regex-count "\\SU[0-9]" "N.YM.SU3.60.tiff.tif")
(de read-label (s)
  (let ((expressions 
	'("\\HA[0-9]"
	 "\\SU[0-9]"
	 "\\NE[0-9]"
	 "\\AN[0-9]"
	 "\\SA[0-9]"
	 "\\DI[0-9]"
	 "\\FE[0-9]")) (i 0) (label 7))
    ;; flip through regular expresions and find a label
    (each ((exp expressions))
      (if (= 1 (regex-count exp s)) 
	  ;; if not 7, then we already classified, problem!
	  (setq label i))
      ;; pop off the list
      (setq expressions (cdr expressions))
      (incr i))
    ;; error if no label found
    (if (= label 7) (error "bad label %s" s))
    ;;(printf "%d" label)
    label

    )
)


;; switch on a classification and return it's meaning
(de write-label (n)
  (let ((labels '("Happy" "Surprised" "Neutral" "Angry" "Sad" "Disgusted" "Afraid")))
    (nth n labels)
    )
)

;; function should read the image names into a list and shuffle them
(de read-dir (ddir)
  ;;(let ((x 0))
  ;;  (each ((x (ls "jaffedata")))
  ;;   (printf x)))
  ;; read the list of files and randomize it
  ;; pop off names from files and in
  (let ((dfiles (ls ddir)) (i 1))
      ;;print the names, for debugging purposes
      ;;(each ((x files)) (printf "%d - %l\n" i x) (incr i))
      ;;now randomize
      ;;(shuf dfiles)
      ;;(each ((x files)) (printf "%d - %l\n" i x) (incr i))
      dfiles
      )
)

;; shuffles a list, courtesy of Yann LeCun
(de shuf (l)
  (let ((n (length l))
	(p l))
    (while p
      (let ((tmp (car p))
	    (c2 (nthcdr (rand 0 (- n 1)) l)))
	(rplaca p (car c2))
	(rplaca c2 tmp)
	(setq p (cdr p))))
    l))

;; Fixed version to read TIFF images
(de image-read-ubim-fixed (s)
  ((-str-) s)
  (let* ((f (popen (sprintf "%s %s PPM:-" 
                            "/usr/bin/convert -compress lossless" s) "r")))
    (when (not f)
      (error "couldn't open convert pipe file") )
    (prog1
        (pnm-fread-ubim f)
      (pclose f) ) ) )



;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;
;;;
;;; The following is for construction of the network. This
;;; should define the # of classifications
;;;
;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;;

(de run-facex (trsize tesize niter)

;; number of classes
(setq nclasses 7)
;; the target values for mean-squared error training
;; are +target for the output corresponding to the class 
;; being shown and -target for the other outputs.
(setq target 1)

;; fill matrix with 1-of-n code
(setq labels (int-matrix nclasses))
(setq targets (float-matrix nclasses nclasses))
(idx-f2dotc targets 1.5 targets)
(targets ()() (- target))
(for (i 0 (- nclasses 1)) 
     (targets i i target)
     (labels i i))

;; create trainable parameter
(setq theparam (new idx1-ddparam 0 60000))

;; create the network
(setq thenet
      (new idx3-supervised-module
	   ;; height width conx2 sigx2 "" "" hidden #class param
	   (new-lenet5 *net-size* *net-size* 9 9 2 2 9 9 2 2 10 7 theparam)
	   (new edist-cost labels 1 1 targets)
	   (new max-classer labels)))

;; create the trainer
(setq thetrainer
      (new supervised-gradient thenet theparam))

;; a classifier-meter measures classification errors
(setq trainmeter (new classifier-meter))
(setq testmeter  (new classifier-meter))
(setq friendsmeter (new classifier-meter))

;; initialize the network weights
(==> :thenet:machine forget 1 2)

(when *theparam*
  (==> theparam load *theparam*))

;; build databases
;;(build-databases-jaffe 319 107)
(build-databases-jaffe trsize tesize)
(build-database-friends)


;; do training iterations 
(printf "training with %d training samples and %d test samples\n" 
	(==> trainingdb size)
	(==> testingdb size))

;; this goes at about 25 examples per second on a PIIIM 800MHz
(de doit (n)
  (repeat n
;; .0001
    ;; estimate second derivative on 100 iterations, using mu=0.02
    ;; and set individual espilons
    (printf "computing diagonal hessian and learning rates\n")
    (==> thetrainer compute-diaghessian trainingdb 50 0.02)
    (printf "training: ") (flush)
    (==> thetrainer train trainingdb trainmeter 0.0001 0)
    (==> thetrainer test trainingdb trainmeter)
    (==> trainmeter display)
    (printf " testing: ") (flush)
    (==> thetrainer test testingdb testmeter)
    (==> testmeter display)
 ;;run the test on my friends
   (printf " friends: ") (flush)
   (==> thetrainer test friendsdb friendsmeter)
   (==> friendsmeter display)

    ()))

(print (time (doit niter)))

)


;; Perform a high-pass filter on a UBIM greyscale image.
;; Returns a new matrix
(de ubim-high-pass (img)
  (let ((input (matrix (idx-dim img 0) (idx-dim img 1)))
	(dresult (matrix (- (idx-dim img 0) 4) (- (idx-dim img 1) 4)))
	(uresult (ubyte-matrix (- (idx-dim img 0) 4) (- (idx-dim img 1) 4)))
	(kernel [[-1 0 0 0 -1]
		 [0 -1 0 -1 0]
		 [0 0 8 0 0]
		 [0 -1 0 -1 0]
		 [-1 0 0 0 -1]])
	(low 0))

    ;; performs numerical conversion automatically
    (idx-copy img input)
    ;; perform high-pass
    (midx-m2convol input kernel dresult)
    ;; need to find the lowest negative value to re-center values to 0
    (idx-bloop ((row dresult))
      (idx-bloop ((value row))
	(when (< (value) 0)
	  (value 0))
	  ;;	  (value (+ 256 (value)))
	(when (> (value) 255)
	  (value 255))
      ))

    ;;    (pause)
    (idx-copy dresult uresult)
    ;; return result
    uresult))

;; show an example of the high-pass filter
(de run-example ()
  (let ((image1 (image-read-ubim-fixed "example/original.tif"))
	(image2 (image-read-ubim-fixed "example/original2.tif"))
	(exdir "./example/class/") ;; probably unnecessary
	)
    (show-all image1 
	      (ubim-high-pass image1) image2 (ubim-high-pass image2) 
	      "high pass example")
    (let ((exfiles (read-dir exdir)) ;; probably unnecessary
	  (exdata (build-database-dir exdir)))
      ;; here I need to fprop over the datasource
      ;;(==> exdata seek 1)
      ;;(setq img1
      (printf " testing: ") (flush)
      (==> thetrainer test exdata testmeter)
      (==> testmeter display)
      (pause)
      (setq offset 0)
      ;; piecemeal testing
      ;; 0 is -1/1 for correct or not, 1 is the output, 2 is the desired
      (for (i 0 7)
	   (let ((result (==> thetrainer test-sample exdata testmeter i)))
	     (if (> 0 (nth 0 result))
		 (printf "\nfailed ") (printf "\npassed "))
	     (printf "label: %d,%s" (nth 2 result) (write-label (nth 2 result)))
	     (printf "\tguess: %d,%s" (nth 1 result) (write-label (nth 1 result)))
	     )
	   (let ((input (new idx3-state 64 64 64)) (desired (int-matrix)))
	     (==> exdata seek i)
	     (==> exdata fprop input desired)
	     (gray-draw-matrix 0 offset (select :input:x 0 0) 0 2.55)
	     )
	   (setq offset (+ 64 offset))
	   )   

;;      (show-all image1 
;;	      (ubim-high-pass image1) image2 (ubim-high-pass image2) 
;;	      "high pass example")
)))