show_color_image.ml

open Owl_neural
open Owl_neural_graph
open Algodiff.S
open Owl_neural_neuron
open Plplot


let x, y = Dataset.load_cifar_train_data 1
let m = Dense.Matrix.S.row_num x
let num_test = 9
let a, b = Owl_neural_optimise.Utils.draw_samples (Mat x) (Mat y) num_test
let imgs = Algodiff.S.unpack_mat a

let img = Dense.Matrix.S.slice [[0]] imgs
let img = Dense.Matrix.S.to_ndarray img
let img = Dense.Ndarray.S.reshape img [|32;32;3|] 

let red = Dense.Ndarray.S.slice [[]; []; [0]] img 
let red = Dense.Ndarray.S.reshape red [|32;32|]

let green = Dense.Ndarray.S.slice [[]; []; [1]] img 
let green = Dense.Ndarray.S.reshape red [|32;32|]

let blue = Dense.Ndarray.S.slice [[]; []; [2]] img 
let blue = Dense.Ndarray.S.reshape red [|32;32|]

(* Apparently adding is not the right way to mix 3 colors*)
let index = Dense.Ndarray.S.(add red (add green blue))
let index = Dense.Ndarray.S.div_scalar index (255. *. 3.)

let red = Dense.Ndarray.S.div_scalar red 255.
let green = Dense.Ndarray.S.div_scalar red 255.
let blue = Dense.Ndarray.S.div_scalar red 255.

let arraytoplot x = 
    let x = x |> Dense.Ndarray.Generic.cast_s2d
              |> Dense.Matrix.D.of_ndarray
    in 
    Owl_dense_matrix.D.rotate x 90 |> Owl_dense_matrix.D.to_arrays 

let x = arraytoplot red (*index/green/blue*)

let _ = 
plparseopts Sys.argv [PL_PARSE_FULL];
plinit ();

plcol0 2; (* draw a yellow plot box, useful for diagnostics *)
plenv 0.0 35.0 0.0 35.0 1 (-1);


plspal0 "cmap0_black_on_white.pal";
plspal1 "cmap1_blue_red.pal" true;
(* plspal1 "cmap1_blue_yellow.pal" true; *)


plimage x 1.0 32.0 1.0 32.0 0.0 1.0 1.0 32.0 1.0 32.0;

plend (); ()