Remote reading system by M5Stack and K210

#coding:utf-8
import tensorflow as tf
from tensorflow.examples.tutorials.mnist import input_data
import time

"""
  weight init
0  init to a small number close to zero
"""
def weight_variable(shape):
    initial = tf.truncated_normal(shape, stddev = 0.1)
    return tf.Variable(initial)

def bias_variable(shape):
    initial = tf.constant(0.1, shape = shape)
    return tf.Variable(initial)

"""
1stride size,0padding size
2x2max pooling
conv and pool
"""
def conv2d(x, W):
    return tf.nn.conv2d(x, W, strides=[1, 1, 1, 1], padding = 'SAME')
    # tf.nn.conv2d(input, filter, strides, padding, use_cudnn_on_gpu=None, data_format=None, name=None)
    # x(input)  : [batch, in_height, in_width, in_channels]
    # W(filter) : [filter_height, filter_width, in_channels, out_channels]
    # strides   : The stride of the sliding window for each dimension of input.
    #             For the most common case of the same horizontal and vertices strides, strides = [1, stride, stride, 1]

def max_pool_2x2(x):
    return tf.nn.max_pool(x, ksize = [1, 2, 2, 1],
                          strides = [1, 2, 2, 1], padding = 'SAME')
    # tf.nn.max_pool(value, ksize, strides, padding, data_format='NHWC', name=None)
    # x(value)              : [batch, height, width, channels]
    # ksize(pool)        : A list of ints that has length >= 4. The size of the window for each dimension of the input tensor.
    # strides(pool)   : A list of ints that has length >= 4. The stride of the sliding window for each dimension of the input tensor.

start = time.clock() #
mnist = input_data.read_data_sets("MNIST_data/", one_hot=True) #MNIST

"""
    1st layer  conv layer
x_image(batch, 28, 28, 1) -> h_pool1(batch, 14, 14, 16)
"""
x = tf.placeholder(tf.float32,[None, 784])
x_image = tf.reshape(x, [-1, 28, 28, 1])  
W_conv1 = weight_variable([3, 3, 1, 16])
b_conv1 = bias_variable([16])

h_conv1 = tf.nn.relu(conv2d(x_image, W_conv1) + b_conv1)
# x_image -> [batch, in_height, in_width, in_channels]
#            [batch, 28, 28, 1]
# W_conv1 -> [filter_height, filter_width, in_channels, out_channels]
#            [3, 3, 1, 16]
# output  -> [batch, out_height, out_width, out_channels]
#            [batch, 28, 28, 16]
h_pool1 = max_pool_2x2(h_conv1)
# h_conv1 -> [batch, in_height, in_weight, in_channels]
#            [batch, 28, 28, 16]
# output  -> [batch, out_height, out_weight, out_channels]
#            [batch, 14, 14, 16]

"""
   2nd layer  conv layer
h_pool1(batch, 14, 14, 16) -> h_pool2(batch, 7, 7, 32)
"""
W_conv2 = weight_variable([3, 3, 16, 32])
b_conv2 = bias_variable([32])

h_conv2 = tf.nn.relu(conv2d(h_pool1, W_conv2) + b_conv2)
# h_pool1 -> [batch, 14, 14, 16]
# W_conv2 -> [5, 5, 16, 32]
# output  -> [batch, 14, 14, 32]
h_pool2 = max_pool_2x2(h_conv2)
# h_conv2 -> [batch, 14, 14, 32]
# output  -> [batch, 7, 7, 32]

"""
    3rd layer  full connect layer 
h_pool2(batch, 7, 7, 32) -> h_fc1(1, 1)
"""
W_fc1 = weight_variable([7 * 7 * 32, 32])
b_fc1 = bias_variable([32])

h_pool2_flat = tf.layers.flatten(h_pool2)
h_fc1 = tf.nn.relu(tf.matmul(h_pool2_flat, W_fc1) + b_fc1)

"""
Dropout
h_fc1 -> h_fc1_drop
"""
keep_prob = tf.placeholder("float")
h_fc1_drop = tf.nn.dropout(h_fc1, keep_prob)

"""
 Softmax   4th layer  Softmax output
"""
W_fc2 = weight_variable([32, 10])
b_fc2 = bias_variable([10])

y_conv = tf.nn.softmax(tf.matmul(h_fc1_drop, W_fc2) + b_fc2)

"""
		train and eval the module
ADAM,feed_dictkeep_probdropout
"""
y_ = tf.placeholder("float", [None, 10])
cross_entropy = -tf.reduce_sum(y_ * tf.log(y_conv)) #  
train_step = tf.train.AdamOptimizer(1e-4).minimize(cross_entropy) #adam0.0001
correct_prediction = tf.equal(tf.argmax(y_conv,1), tf.argmax(y_,1)) #
accuracy = tf.reduce_mean(tf.cast(correct_prediction,"float"))
saver = tf.train.Saver()

sess = tf.Session() #
#sess.run(tf.initialize_all_variables()) #
sess.run(tf.global_variables_initializer()) #
merged = tf.summary.merge_all() 
writer = tf.summary.FileWriter('logs',sess.graph)

for i in range(20000): #5000
    batch = mnist.train.next_batch(50) #batch50
    if i % 100 == 0:
        train_accuracy = accuracy.eval(session = sess,
                                       feed_dict = {x:batch[0], y_:batch[1], keep_prob:1.0})
        print("step %d, train_accuracy %g" %(i, train_accuracy))
        saver.save(sess, 'model/mnist.ckpt')
    train_step.run(session = sess, feed_dict = {x:batch[0], y_:batch[1],
                   keep_prob:0.5}) # keep_prob 0.5

sess = tf.InteractiveSession()
sess.run(tf.global_variables_initializer())
saver = tf.train.Saver(tf.global_variables())
saver.restore(sess, 'model/mnist.ckpt')
print( "test accuracy %g" % accuracy.eval(feed_dict={
    x:mnist.test.images, y_:mnist.test.labels, keep_prob:1.0}))


#get parms
WC1 = W_conv1.eval()
BC1 = b_conv1.eval()
WC2 = W_conv2.eval()
BC2 = b_conv2.eval()
WF1 = W_fc1.eval()
BF1 = b_fc1.eval()
WF2 = W_fc2.eval()
BF2 = b_fc2.eval()
#restruct new graph to save
g = tf.Graph()
with g.as_default():
    x_image=tf.placeholder("float", shape=[None,28,28,1], name="inputs")

    WC1 = tf.constant(WC1, name="WC1")
    BC1 = tf.constant(BC1, name="BC1")
    CONV1 = tf.nn.relu(conv2d(x_image, WC1) + BC1)
    MAXPOOL1 = max_pool_2x2(CONV1)

    WC2 = tf.constant(WC2, name="WC2")
    BC2 = tf.constant(BC2, name="BC2")
    CONV2 = tf.nn.relu(conv2d(MAXPOOL1, WC2) + BC2)
    MAXPOOL2 = max_pool_2x2(CONV2)

    WF1 = tf.constant(WF1, name="WF1")
    BF1 = tf.constant(BF1, name="BF1")
    FC1 = tf.layers.flatten(MAXPOOL2)   
    FC1 = tf.nn.relu(tf.matmul(FC1, WF1) + BF1)

    WF2 = tf.constant(WF2, name="WF2")
    BF2 = tf.constant(BF2, name="BF2")
    OUTPUT = tf.nn.softmax(tf.matmul(FC1, WF2) + BF2,name="output")
    
    sess = tf.Session()
    sess.run(tf.global_variables_initializer())

    graph_def = g.as_graph_def()
    tf.train.write_graph(graph_def, "./", "mnist.pb", as_text=False)

"""
k210
mnist.kmodelsdsdsd
TODO
1
2
"""
from fpioa_manager import fm, board_info
from machine import UART
import utime
fm.register(board_info.PIN9,fm.fpioa.UART2_TX)
fm.register(board_info.PIN10,fm.fpioa.UART2_RX)
uart_B = UART(UART.UART2, 115200, 8, None, 1, timeout=10)
import sensor, image, time, lcd, math
import KPU as kpu
#task = kpu.load("/sd/paste_mnist.kmodel")
task = kpu.load("/sd/mnist.kmodel")
info=kpu.netinfo(task)

lcd.init(freq=15000000)
sensor.reset()                      # Reset and initialize the sensor. It will
                                    # run automatically, call sensor.run(0) to stop
sensor.set_pixformat(sensor.RGB565) # Set pixel format to RGB565 (or GRAYSCALE)
sensor.set_framesize(sensor.QVGA)   # Set frame size to QVGA (320x240)
sensor.set_vflip(True)
sensor.set_auto_gain(True)
sensor.set_auto_whitebal(True)
sensor.set_gainceiling(8)
sensor.skip_frames(time = 2000)     # Wait for settings take effect.
clock = time.clock()                # Create a clock object to track the FPS.

def mnist_run(img, dx, dy, dis, x00 =0, y00 = 80, nnn = 2):
    if nnn == 4:
        x00 = x00
        dy = dy
    img0 = img.copy((x00+dis*nnn,y00+nnn*0, dx, dy))
    #img0.mean(2, threshold=True, offset=1, invert=True)  #A
    img0.median(2, percentile=0.3, threshold=True, offset=-3, invert=True)
    #img0.midpoint(2, bias=0.3, threshold=True, offset=0, invert=True)
    #img0.mode(2, threshold=True, offset=0, invert=True)  #B

    #img0.binary([(110,255)], invert = True)
    for dx0 in range(dx):
        for dy0 in range(dy):
            a0 = img0.get_pixel(dx0,dy0)
            img.set_pixel(x00+dis*nnn+dx0,y00+nnn*0+dy0,a0)
    #img1 = img0.copy((1,1, dx-1, dy-1))
    img1 = img0
    img1 = img1.resize(28,28)
    img1 = img1.to_grayscale(1)
    img1.pix_to_ai()
    fmap=kpu.forward(task,img1)
    plist=fmap[:]
    pmax=max(plist)
    max_index=plist.index(pmax)
    kpu.fmap_free(fmap)
    return max_index, pmax

def search_col(x_input, y_input, img, width = 320, height = 240):
    x_l = []
    y_l = []
    for x in range(x_input - 32,x_input + 32):
        for y in range(y_input - 32,y_input + 32):
            if math.sqrt((x-x_input)*(x-x_input) + (y-y_input)*(y-y_input))<32 and math.sqrt((x-x_input)*(x-x_input) + (y-y_input)*(y-y_input))>14:
                col = img.get_pixel(x,y)
                if col[0]>120 and col[1]<100 and col[2]<100:
                    x_l.append(x-x_input)
                    y_l.append(-y+y_input)
                    #img.set_pixel(x,y,(255,255,255))
                #else:
                    #img.set_pixel(x,y,(0,0,0))
    angle_count = 0
    le = 0
    x_c = 0
    y_c = 0
    for i in range(len(x_l)):
        leng = math.sqrt(x_l[i]**2 + y_l[i]**2)
        le = le + leng
        angle_count = angle_count + math.acos(y_l[i]/leng)*leng
        x_c = x_c + x_l[i]
        y_c = y_c + y_l[i]
    if le == 0:
        angle = 0
    else:
        angle = angle_count/le
    dx = 0
    dy = 0
    dx = int(30 * math.sin(angle))
    dy = int(30 * math.cos(angle))
    if x_c < 0:
        angle = -angle + 2*math.pi
        dx = -dx
    img.draw_line((x_input, y_input,x_input+dx, y_input-dy), thickness = 2, color=(0,0,255))
    return angle/math.pi*180

num_list = [0, 0, 0, 0, 5]
p_list = [0,0,0,0,0]
angle_list = [0,0,0,0]
while(True):
    count_0 = 0
    count_4 = 0
    clock.tick()                    # Update the FPS clock.
    img = sensor.snapshot()         # Take a picture and return the image.
    #img.mean(1, threshold=True, offset=5, invert=True)
    #img.binary([(100,255)], invert = True)
    #img.erode(1)

    x00 = 91
    y00 = 4
    dx = 20
    dy = 20
    dis = 25
    p_thre = 0.95
    for i in range(0,5):
        class_num, pmax = mnist_run(img, dx, dy, dis,\
            x00 =x00, y00 = y00,\
            nnn=i)
        if pmax > p_thre:
            num_list[i] = class_num
            p_list[i] = pmax

    for i in range(0,5):
        if i == 4:
            x00 = x00
            dy = dy
        img.draw_rectangle((x00+dis*i,y00+i*0, dx, dy), color=255)
    R_list = []
    c_color = []
    x_list = [101+3, 175+2, 241, 263]
    y_list = [176-6, 193-6, 156-6, 84-6]

    angle_list[0] = search_col(x_list[0], y_list[0], img, width = 320, height = 240)
    angle_list[1] = search_col(x_list[1], y_list[1], img, width = 320, height = 240)
    angle_list[2] = search_col(x_list[2], y_list[2], img, width = 320, height = 240)
    angle_list[3] = search_col(x_list[3], y_list[3], img, width = 320, height = 240)
    print(num_list)
    print(p_list)
    #print(angle_list)
    R = 32
    img.draw_circle(x_list[0], y_list[0], R, color = (255, 0, 0), thickness = 2, fill = False)
    img.draw_circle(x_list[1], y_list[1], R, color = (255, 0, 0), thickness = 2, fill = False)
    img.draw_circle(x_list[2], y_list[2], R, color = (255, 0, 0), thickness = 2, fill = False)
    img.draw_circle(x_list[3], y_list[3], R, color = (255, 0, 0), thickness = 2, fill = False)

    # R-G-B 180-60-60
    r = 3
    img.draw_circle(x_list[0], y_list[0], r, color = (255, 255, 0), thickness = 1, fill = False)
    img.draw_circle(x_list[1], y_list[1], r, color = (255, 255, 0), thickness = 1, fill = False)
    img.draw_circle(x_list[2], y_list[2], r, color = (255, 255, 0), thickness = 1, fill = False)
    img.draw_circle(x_list[3], y_list[3], r, color = (255, 255, 0), thickness = 1, fill = False)
    utime.sleep_ms(250)
    #str(num_list[0])
    uart_B.write(str(0)+ str(0)+ str(0)+ str(0)+ str(6)+\
        '.' + str(int(angle_list[3]/36)) + str(int(angle_list[2]/36)) + str(int(angle_list[1]/36)) + str(int(angle_list[0]/36)))

    lcd.display(img)                # Display on LCD

    #lcd.draw_string(20,20,"%d: %.3f"%(max_index,pmax),lcd.WHITE,lcd.BLACK)