wpvnteam/vietnamese-handwriting
0
1var log = console.log;2var ctx = null;3var canvas = null;4var RNN_SIZE = 400;5var VOCAB_SIZE = 165;6var NUM_ATT_HEADS=10;7var NUM_GMM_HEADS=20;8var cur_run = 0;9var scale_factor = 1.;10 11var randn = function() {12 // Standard Normal random variable using Box-Muller transform.13 var u = Math.random() * 0.999 + 1e-5;14 var v = Math.random() * 0.999 + 1e-5;15 return Math.sqrt(-2.0 * Math.log(u)) * Math.cos(2.0 * Math.PI * v);16}17 18var rand_truncated_normal = function(low, high) {19 while (true) {20 r = randn();21 if (r >= low && r <= high)22 break;23 // rejection sampling.24 }25 return r;26}27 28var softplus = function(x) {29 const m = tf.maximum(x, 0.0);30 return tf.add(m, tf.log(tf.add(tf.exp(tf.neg(m)), tf.exp(tf.sub(x, m)))));31}32 33var char2idx = {'\x00': 0, ' ': 1, '!': 2, '"': 3, '#': 4, '%': 5, '&': 6, "'": 7, '(': 8, ')': 9, '*': 10, ',': 11, '-': 12, '.': 13, '/': 14, '0': 15, '1': 16, '2': 17, '3': 18, '4': 19, '5': 20, '6': 21, '7': 22, '8': 23, '9': 24, ':': 25, ';': 26, '?': 27, 'A': 28, 'B': 29, 'C': 30, 'D': 31, 'E': 32, 'F': 33, 'G': 34, 'H': 35, 'I': 36, 'J': 37, 'K': 38, 'L': 39, 'M': 40, 'N': 41, 'O': 42, 'P': 43, 'Q': 44, 'R': 45, 'S': 46, 'T': 47, 'U': 48, 'V': 49, 'W': 50, 'X': 51, 'Y': 52, 'a': 53, 'b': 54, 'c': 55, 'd': 56, 'e': 57, 'f': 58, 'g': 59, 'h': 60, 'i': 61, 'j': 62, 'k': 63, 'l': 64, 'm': 65, 'n': 66, 'o': 67, 'p': 68, 'q': 69, 'r': 70, 's': 71, 't': 72, 'u': 73, 'v': 74, 'w': 75, 'x': 76, 'y': 77, 'z': 78, 'À': 79, 'Á': 80, 'Â': 81, 'Ô': 82, 'Ú': 83, 'Ý': 84, 'à': 85, 'á': 86, 'â': 87, 'ã': 88, 'è': 89, 'é': 90, 'ê': 91, 'ì': 92, 'í': 93, 'ò': 94, 'ó': 95, 'ô': 96, 'õ': 97, 'ù': 98, 'ú': 99, 'ý': 100, 'Ă': 101, 'ă': 102, 'Đ': 103, 'đ': 104, 'ĩ': 105, 'ũ': 106, 'Ơ': 107, 'ơ': 108, 'Ư': 109, 'ư': 110, 'ạ': 111, 'Ả': 112, 'ả': 113, 'Ấ': 114, 'ấ': 115, 'Ầ': 116, 'ầ': 117, 'ẩ': 118, 'ẫ': 119, 'ậ': 120, 'ắ': 121, 'ằ': 122, 'ẳ': 123, 'ẵ': 124, 'ặ': 125, 'ẹ': 126, 'ẻ': 127, 'ẽ': 128, 'ế': 129, 'Ề': 130, 'ề': 131, 'Ể': 132, 'ể': 133, 'ễ': 134, 'Ệ': 135, 'ệ': 136, 'ỉ': 137, 'ị': 138, 'ọ': 139, 'ỏ': 140, 'Ố': 141, 'ố': 142, 'Ồ': 143, 'ồ': 144, 'ổ': 145, 'ỗ': 146, 'ộ': 147, 'ớ': 148, 'ờ': 149, 'Ở': 150, 'ở': 151, 'ỡ': 152, 'ợ': 153, 'ụ': 154, 'Ủ': 155, 'ủ': 156, 'ứ': 157, 'ừ': 158, 'ử': 159, 'ữ': 160, 'ự': 161, 'ỳ': 162, 'ỷ': 163, 'ỹ': 164};34 35var gru_core = function(input, weights, state, hidden_size) {36 var [w_h,w_i,b] = weights;37 var [w_h_z,w_h_a] = tf.split(w_h, [2 * hidden_size, hidden_size], 1);38 var [b_z,b_a] = tf.split(b, [2 * hidden_size, hidden_size], 0);39 gates_x = tf.matMul(input, w_i);40 [zr_x,a_x] = tf.split(gates_x, [2 * hidden_size, hidden_size], 1);41 zr_h = tf.matMul(state, w_h_z);42 zr = tf.add(tf.add(zr_x, zr_h), b_z);43 // fix this44 [z,r] = tf.split(tf.sigmoid(zr), 2, 1);45 a_h = tf.matMul(tf.mul(r, state), w_h_a);46 a = tf.tanh(tf.add(tf.add(a_x, a_h), b_a));47 next_state = tf.add(tf.mul(tf.sub(1., z), state), tf.mul(z, a));48 return [next_state, next_state];49};50 51 52var generate = function() {53 cur_run = cur_run + 1;54 setTimeout(function() { 55 var counter = 2000;56 tf.disposeVariables();57 58 tf.engine().startScope();59 ctx.clearRect(0, 0, canvas.width, canvas.height);60 ctx.beginPath();61 dojob(cur_run);62 }, 200);63 64 return false;65}66 67var dojob = function(run_id) {68 var text = document.getElementById("user-input").value;69 if (text.length == 0) {70 text = "Tất cả mọi người đều sinh ra có quyền bình đẳng.";71 }72 73 74 log(text);75 original_text = text;76 text = '' + text + ' ';77 78 text = Array.from(text).map(function(e) {79 return char2idx[e]80 })81 var text_embed = WEIGHTS['rnn/~/embed_1__embeddings'];82 indices = tf.tensor1d(text, 'int32');83 text = text_embed.gather(indices);84 85 var embed = text;86 87 var writer_embed = WEIGHTS['rnn/~/embed__embeddings'];88 var e = document.getElementById("writers");89 var wid = parseInt(e.value);90 log(wid);91 92 wid = tf.tensor1d([wid], 'int32');93 wid = writer_embed.gather(wid);94 embed = tf.add(wid, embed);95 96 97 98 filter = WEIGHTS['rnn/~/conv1_d__w'];99 embed = tf.conv1d(embed, filter, 1, 'same');100 bias = tf.expandDims(WEIGHTS['rnn/~/conv1_d__b'], 0);101 embed = tf.add(embed, bias);102 103 104 // initial state105 var gru0_hx = tf.zeros([1, RNN_SIZE]);106 var gru1_hx = tf.zeros([1, RNN_SIZE]);107 var gru2_hx = tf.zeros([1, RNN_SIZE]);108 109 var att_location = tf.zeros([1, NUM_ATT_HEADS]);110 var att_context = tf.zeros([1, VOCAB_SIZE]);111 112 var input = tf.tensor([[0., 0., 1.]]);113 114 gru0_w_h = WEIGHTS['rnn/~/attention_core/~/gru__w_h'];115 gru0_w_i = WEIGHTS['rnn/~/attention_core/~/gru__w_i'];116 gru0_bias = WEIGHTS['rnn/~/attention_core/~/gru__b'];117 118 gru1_w_h = WEIGHTS['rnn/~/attention_core/~/gru_1__w_h'];119 gru1_w_i = WEIGHTS['rnn/~/attention_core/~/gru_1__w_i'];120 gru1_bias = WEIGHTS['rnn/~/attention_core/~/gru_1__b'];121 122 gru2_w_h = WEIGHTS['rnn/~/attention_core/~/gru_2__w_h'];123 gru2_w_i = WEIGHTS['rnn/~/attention_core/~/gru_2__w_i'];124 gru2_bias = WEIGHTS['rnn/~/attention_core/~/gru_2__b'];125 126 att_w = WEIGHTS['rnn/~/attention_core/~/linear__w'];127 att_b = WEIGHTS['rnn/~/attention_core/~/linear__b'];128 gmm_w = WEIGHTS['rnn/~/linear__w'];129 gmm_b = WEIGHTS['rnn/~/linear__b'];130 131 var ruler = tf.tensor([...Array(text.shape[0]).keys()]);132 ruler = tf.expandDims(ruler, 1);133 var bias = parseInt(document.getElementById("bias").value) / 100 * 3;134 135 136 var cur_x = 20;137 var cur_y = innerHeight / 2 + 30;138 var path = [];139 var dx = 0.;140 var dy = 0;141 var eos = 1.;142 var counter = 0;143 144 145 function loop(my_run_id) {146 if (my_run_id < cur_run) {147 tf.disposeVariables();148 tf.engine().endScope();149 return;150 }151 152 counter++;153 if (counter < 2000) {154 [att_location,att_context,gru0_hx,gru1_hx, gru2_hx, input] = tf.tidy(function() {155 // Attention156 const inp_0 = tf.concat([att_context, input], 1);157 gru0_hx_ = gru0_hx;158 [out_0,gru0_hx] = gru_core(inp_0, [gru0_w_h, gru0_w_i, gru0_bias], gru0_hx, RNN_SIZE);159 tf.dispose(gru0_hx_);160 const att_inp = tf.concat([att_context, input, out_0], 1);161 const att_params = tf.add(tf.matMul(att_inp, att_w), att_b);162 [alpha,beta,kappa] = tf.split(softplus(att_params), 3, 1);163 att_location_ = att_location;164 att_location = tf.add(att_location, tf.div(kappa, 25.));165 tf.dispose(att_location_)166 167 var phi = tf.sum(tf.mul(alpha, tf.exp(tf.div(tf.neg(tf.square(tf.sub(att_location, ruler))), beta))), 1);168 phi = tf.expandDims(phi, 0);169 170 att_context_ = att_context;171 att_context = tf.sum(tf.mul(tf.expandDims(phi, 2), tf.expandDims(embed, 0)), 1)172 tf.dispose(att_context_);173 174 const inp_1 = tf.concat([input, out_0, att_context], 1);175 // tf.dispose(input);176 gru1_hx_ = gru1_hx;177 [out_1,gru1_hx] = gru_core(inp_1, [gru1_w_h, gru1_w_i, gru1_bias], gru1_hx, RNN_SIZE);178 tf.dispose(gru1_hx_);179 180 const inp_2 = tf.concat([input, out_1, att_context], 1);181 tf.dispose(input);182 gru2_hx_ = gru2_hx;183 [out_2, gru2_hx] = gru_core(inp_2, [gru2_w_h, gru2_w_i, gru2_bias], gru2_hx, RNN_SIZE);184 tf.dispose(gru2_hx_);185 186 // debugger;187 188 // GMM189 const gmm_params = tf.add(tf.matMul(out_2, gmm_w), gmm_b);190 [x,y,logstdx,logstdy,angle,log_weight,eos_logit] = tf.split(gmm_params, [NUM_GMM_HEADS, NUM_GMM_HEADS, NUM_GMM_HEADS, NUM_GMM_HEADS, NUM_GMM_HEADS, NUM_GMM_HEADS, 1], 1);191 // log_weight = tf.softmax(log_weight, 1);192 // log_weight = tf.log(log_weight);193 // log_weight = tf.mul(log_weight, 1. + bias);194 const idx = tf.argMax(log_weight, 1).dataSync()[0];195 // const idx = tf.multinomial(log_weight, 1).dataSync()[0];196 x = x.dataSync()[idx];197 y = y.dataSync()[idx];198 const stdx = tf.exp(tf.sub(logstdx, bias)).dataSync()[idx];199 const stdy = tf.exp(tf.sub(logstdy, bias)).dataSync()[idx];200 angle = angle.dataSync()[idx];201 e = tf.sigmoid(tf.mul(eos_logit, (1. + bias/5))).dataSync()[0];202 const rx = rand_truncated_normal(-5, 5) * stdx;203 const ry = rand_truncated_normal(-5, 5) * stdy;204 x = x + Math.cos(-angle) * rx - Math.sin(-angle) * ry;205 y = y + Math.sin(-angle) * rx + Math.cos(-angle) * ry;206 if (Math.random() < e) {207 e = 1.;208 } else {209 e = 0.;210 }211 input = tf.tensor([[x, y, e]]);212 return [att_location, att_context, gru0_hx, gru1_hx, gru2_hx, input];213 });214 215 [dx,dy,eos_] = input.dataSync();216 dy = -dy * 3. * scale_factor;217 dx = dx * 3. * scale_factor;218 if (eos == 0.) {219 ctx.beginPath();220 ctx.moveTo(cur_x, cur_y, 0, 0);221 ctx.lineTo(cur_x + dx, cur_y + dy);222 ctx.stroke();223 }224 eos = eos_;225 cur_x = cur_x + dx;226 cur_y = cur_y + dy;227 228 if (att_location.dataSync()[0] < original_text.length + 1.5) {229 setTimeout(function() {loop(my_run_id);}, 0);230 }231 }232 }233 234 loop(run_id);235}236 237 238window.onload = function(e) {239 //Setting up canvas240 canvas = document.getElementById("hw-canvas");241 ctx = canvas.getContext("2d");242 scale_factor = window.innerWidth / 1600;243 ctx.canvas.width = window.innerWidth - 50;244 ctx.canvas.height = window.innerHeight - 50;245 246}247