def __init__(self): # self.image_final = np.array(1, 224, 224, 3) self.bridge = CvBridge() self.count = 0 self.X = tf.placeholder(tf.float32, [1, 224, 224, 3], name="X") self.Y = tf.unstack(inference(self.X), axis=1) dg = {} lg = [ 'w1', 'b1', 'w2', 'b2', 'w3', 'b3', 'w4', 'b4', 'w5', 'b5', 'w_fc1', 'b_fc1', 'w_fc2', 'b_fc2', 'w_output', 'b_output' ] # lg = ['w1', 'b1', 'w2', 'b2', 'w3', 'b3', 'w4', 'b4', 'w5', 'b5', 'w_fc1', 'b_fc1', 'w_fc2', 'b_fc2', 'w_fc3', 'b_fc3', 'w_output', 'b_output'] for i in lg: dg[i] = [ v for v in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) if v.name == i + ':0' ][0] self.saver_g = tf.train.Saver(dg) self._session = tf.InteractiveSession() self.saver_g.restore( self._session, '/home/juna/catkin_ws/src/ros_grasp_detection/src/models/grasp/m4/m4.ckpt' ) # init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) # init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) # self._session.run(init_op) self.object_pub = rospy.Publisher("/objects", positionNpose, queue_size=1) self.image_sub = rospy.Subscriber("/croppedRoI", Image, self.image_callback, queue_size=1, buff_size=2**24)
def run_training(): print(FLAGS.train_or_validation) if FLAGS.train_or_validation == 'train': print('distorted_inputs') data_files_ = TRAIN_FILE images, bboxes = grasp_img_proc.distorted_inputs( [data_files_], FLAGS.num_epochs, batch_size=FLAGS.batch_size) else: print('inputs') data_files_ = VALIDATE_FILE images, bboxes = grasp_img_proc.inputs([data_files_]) x, y, rad = bboxes_to_grasps(bboxes) # images_np = np.array(images) degree = inference(images) # list # tangent of 85 degree is 11 cost = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits(logits=logits, labels=Y)) optimizer = tf.train.AdamOptimizer(learning_rate=learning_rate).minimize(cost) rad_hat_confined = tf.minimum(11., tf.maximum(-11., rad_hat)) rad_confined = tf.minimum(11., tf.maximum(-11., rad)) # Loss function gamma = tf.constant(10.) loss = tf.reduce_mean(tf.pow(x_hat -x, 2) +tf.pow(y_hat -y, 2) + gamma*tf.pow(rad_hat_confined - rad_confined, 2)) train_op = tf.train.AdamOptimizer(epsilon=0.1).minimize(loss) init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) sess = tf.Session() sess.run(init_op) coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess=sess, coord=coord) #save/restore model d={} l = ['w1', 'b1', 'w2', 'b2', 'w3', 'b3', 'w4', 'b4', 'w5', 'b5', 'w_fc1', 'b_fc1', 'w_fc2', 'b_fc2'] for i in l: d[i] = [v for v in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) if v.name == i+':0'][0] dg={} lg = ['w1', 'b1', 'w2', 'b2', 'w3', 'b3', 'w4', 'b4', 'w5', 'b5', 'w_fc1', 'b_fc1', 'w_fc2', 'b_fc2', 'w_output', 'b_output'] for i in lg: dg[i] = [v for v in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) if v.name == i+':0'][0] saver = tf.train.Saver(d) saver_g = tf.train.Saver(dg, max_to_keep=20) #saver.restore(sess, "/root/grasp/grasp-detection/models/imagenet/m2/m2.ckpt") saver_g.restore(sess, FLAGS.model_path) try: count = 0 step = 0 start_time = time.time() while not coord.should_stop(): start_batch = time.time() #train if FLAGS.train_or_validation == 'train': _, loss_value, images_np, x_value, x_model, y_value, y_model, rad_value, rad_model = \ sess.run([train_op, loss, images, x, x_hat, y, y_hat, rad, rad_hat]) duration = time.time() - start_batch if step % 100 == 0: print('Step %d | loss = %s\n | x = %s\n | x_hat = %s\n | y = %s\n | y_hat = %s\n | rad = %s\n | rad_hat = %s\n | (%.3f sec/batch\n') \ %(step, loss_value, x_value[:3], x_model[:3], y_value[:3], y_model[:3], rad_value[:3], rad_model[:3], duration) if step % 1000 == 0: filename = MODEL_SAVE_PATH + '/step_' + str(step) if not os.path.exists(filename): os.mkdir(filename) saver_g.save(sess, filename + '/m4.ckpt') # cv2.imshow('bbox', images_np[1]) ################################### # cv2.waitKey(0) # print(images_np[0]) if step == 100000: sess.close() else: bbox_hat = grasp_to_bbox(x_hat, y_hat, rad_hat) bbox_value, bbox_model, x_value, x_model, y_value, y_model, rad_value, rad_model = sess.run([bboxes, bbox_hat, x, x_hat, y, y_hat, rad, rad_hat]) bbox_value = np.reshape(bbox_value, -1) bbox_value = [(bbox_value[0],bbox_value[1]),(bbox_value[2],bbox_value[3]),(bbox_value[4],bbox_value[5]),(bbox_value[6],bbox_value[7])] # bbox_value = [(bbox_value[0]*0.35,bbox_value[1]*0.47),(bbox_value[2]*0.35,bbox_value[3]*0.47),(bbox_value[4]*0.35,bbox_value[5]*0.47),(bbox_value[6]*0.35,bbox_value[7]*0.47)] p1 = Polygon(bbox_value) p2 = Polygon(bbox_model) iou = p1.intersection(p2).area / (p1.area +p2.area -p1.intersection(p2).area) angle_diff = np.abs(np.arctan(rad_model)*180/np.pi -np.arctan(rad_value)*180/np.pi) duration = time.time() -start_batch # if angle_diff < 30. and iou >= 0.25: count+=1 print('image: %d | duration = %.2f | count = %d | iou = %.2f | angle_difference = %.2f' %(step, duration, count, iou, angle_diff)) print('x=',x_value,x_model,' y=',y_value, y_model,' rad=',rad_value, rad_model,') step +=1 except tf.errors.OutOfRangeError: print('Done training for %d epochs, %d steps, %.1f min.' % (FLAGS.num_epochs, step, (time.time()-start_time)/60)) finally: coord.request_stop() coord.join(threads) sess.close()
def run_training(): tf.reset_default_graph() print(FLAGS.train_or_validation) if FLAGS.train_or_validation == 'train': print('distorted_inputs') data_files_ = TRAIN_FILE #datafiles is the actual tfrdata images, bboxes = grasp_img_proc.distorted_inputs( [data_files_], FLAGS.num_epochs, batch_size=FLAGS.batch_size) else: #Validation dataset print('inputs') data_files_ = VALIDATE_FILE images, bboxes = grasp_img_proc.inputs([data_files_]) # These are the labels (uses the grasp_img_proc module with the disorted images), also everything is processed as bboxes as that the labels are in xy coords x, y, tan, h, w = bboxes_to_grasps(bboxes) # These are the outputs of the model - used for the error minimisation (uses the grasp_inf model) - outputs the grasp representation directly x_hat, y_hat, tan_hat, h_hat, w_hat = tf.unstack(inference(images), axis=1) # tangent of 85 degree is 11 tan_hat_confined = tf.minimum(11., tf.maximum(-11., tan_hat)) tan_confined = tf.minimum(11., tf.maximum(-11., tan)) # Loss function gamma = tf.constant(10.) # A custom cost function for the box regression - essentially just a custom MSE loss = tf.reduce_sum( tf.pow(x_hat - x, 2) + tf.pow(y_hat - y, 2) + gamma * tf.pow(tan_hat_confined - tan_confined, 2) + tf.pow(h_hat - h, 2) + tf.pow(w_hat - w, 2)) # Instead of stochastic gradient descent train_op = tf.train.AdamOptimizer(epsilon=0.1).minimize(loss) # Initiliases the variables to have the values they have been parametised by init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) #Create a session sess = tf.Session() sess.run(init_op) #Allows for thread qeueing coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess=sess, coord=coord) #Some stuff related to plotting #save/restore model #d are the weights that were used to pre-train on Imagenet! - hence why it doesnt contain w_output and b_output (those would have been classification outputs) d = {} l = [ 'w1', 'b1', 'w2', 'b2', 'w3', 'b3', 'w4', 'b4', 'w5', 'b5', 'w_fc1', 'b_fc1', 'w_fc2', 'b_fc2' ] # Iterates through the list l, if its in the GraphKeys, store it in the tuple d for i in l: d[i] = [ v for v in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) if v.name == i + ':0' ][0] dg = {} lg = [ 'w1', 'b1', 'w2', 'b2', 'w3', 'b3', 'w4', 'b4', 'w5', 'b5', 'w_fc1', 'b_fc1', 'w_fc2', 'b_fc2', 'w_output', 'b_output' ] for i in lg: dg[i] = [ v for v in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) if v.name == i + ':0' ][0] # This essentially just saves the current models #saver = tf.train.Saver(d) saver_g = tf.train.Saver(dg) #my_models/test_crash/test_crash_continuation if FLAGS.continue_from_trained_model == 'yes': #saver.restore(sess, "/root/grasp/grasp-detection/models/imagenet/m2/m2.ckpt") saver_g.restore(sess, "my_models/test_save/test_save") '''Restores a previously trained model''' #Just a quick code to determine the ecurrent epoch steps_per_epoch = int(708 / FLAGS.batch_size) model_no = 0 try: count = 0 step = 0 epoch = 0 start_time = time.time() while not coord.should_stop(): start_batch = time.time() if FLAGS.train_or_validation == 'train': #so everytime the loss, x, x_hat, tan etc are called, their operations as well as their operation flows are also called implicitly - graph flow _, loss_value, x_value, x_model, tan_value, tan_model, h_value, h_model, w_value, w_model = sess.run( [ train_op, loss, x, x_hat, tan, tan_hat, h, h_hat, w, w_hat ]) duration = time.time() - start_batch #if step % 100 == 0: #print("Step %d | loss = %s\n | x = %s\n | x_hat = %s\n | tan = %s\n | tan_hat = %s\n | h = %s\n | h_hat = %s\n | w = %s\n | w_hat = %s\n | (%.3f sec/batch\n"%(step, loss_value, x_value[:3], x_model[:3], tan_value[:3], tan_model[:3], h_value[:3], h_model[:3], w_value[:3], w_model[:3], duration)) #How come the y values are not included? - does that not matter because the x's are already being called? if step % 110 == 0: saver_g.save( sess, "my_models/test_save/test_save" + str(model_no)) model_no += 1 if step % steps_per_epoch == 0: print( "Step %d | loss = %s\n | x = %s\n | x_hat = %s\n | tan = %s\n | tan_hat = %s\n | h = %s\n | h_hat = %s\n | w = %s\n | w_hat = %s\n | (%.3f sec/batch\n)" % (step, loss_value, x_value[:3], x_model[:3], tan_value[:3], tan_model[:3], h_value[:3], h_model[:3], w_value[:3], w_model[:3], duration)) print("Current numbers of epoch: %d" % (epoch)) ani = FuncAnimation(fig, animate(epoch, loss_value), 1) epoch += 1 plt.tight_layout() plt.show() #MAKE A LIVE GRAPH MAKE IT EASIER TO SEE else: #VALIDATION '''wont work yet as I have not edited the grasp_img_proc file yet''' #Converts output of NN to four corner vertices bbox_hat = grasp_to_bbox(x_hat, y_hat, tan_hat, h_hat, w_hat) #Gets the value of the actual vertices (randomly), bbox from NN, the actual tan and the predicted tan bbox_value, bbox_model, tan_value, tan_model = sess.run( [bboxes, bbox_hat, tan, tan_hat]) #Turn the bbox value into a 1D array bbox_value = np.reshape(bbox_value, -1) #Rescale it to the size of the 224x224 output of the neural net bbox_value = [(bbox_value[0] * 0.35, bbox_value[1] * 0.47), (bbox_value[2] * 0.35, bbox_value[3] * 0.47), (bbox_value[4] * 0.35, bbox_value[5] * 0.47), (bbox_value[6] * 0.35, bbox_value[7] * 0.47)] #Takes in the x,y coordinates of the vertices, and creates rectangles from vertices p1 = Polygon(bbox_value) p2 = Polygon(bbox_model) #Jaccard Index/ if area is greater than 25% then it counds jaccard = p1.intersection(p2).area / (p1.area + p2.area - p1.intersection(p2).area) #Also if the angle is within 30 degrees of the randomly picked rectangle then: angle_diff = np.abs( np.arctan(tan_model) * 180 / np.pi - np.arctan(tan_value) * 180 / np.pi) duration = time.time() - start_batch if angle_diff < 30. and jaccard >= 0.25: #Add to the count of the 'correct' count += 1 print( 'image: %d | duration = %.2f | count = %d | jaccard = %.2f | angle_difference = %.2f' % (step, duration, count, jaccard, angle_diff)) step += 1 except tf.errors.OutOfRangeError: #some error saver_g.save(sess, "my_models/test_save/test_save" + str(model_no)) #Best to save it again at the end time if FLAGS.train_or_validation == 'train': print('Done training for %d epochs, %d steps, %.1f min.' % (FLAGS.num_epochs, step, (time.time() - start_time) / 60)) else: #print("Number of validation data: ", step) error = ((1 - (count / step)) * 100) print("\nError of %.2f%%" % (error)) finally: coord.request_stop() #stops the threading/queueing coord.join(threads) #rejoins the threads sess.close() return error
def predict_grasp(filepath, filepath_depth, tp): tf.reset_default_graph() init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) #just initialises variables like before sess = tf.Session() sess.run(init_op) img_raw_data = tf.gfile.FastGFile(filepath, 'rb').read() #reads the image img_raw_depth = tf.gfile.FastGFile(filepath_depth, 'rb').read() img_show = cv2.imread(filepath) #this is just to help display it then? #This bit decodes and augments the normal image img_data = tf.image.decode_png(img_raw_data) img_data = tf.image.convert_image_dtype(img_data, dtype=tf.float32) img_reshape = tf.image.resize_images(img_data, [224, 224]) img_reshape = tf.reshape(img_reshape, shape=[1, 224, 224, 3]) #This bit is for decoding the depth data img_depth = tf.image.decode_png(img_raw_depth) img_depth = tf.image.convert_image_dtype(img_depth, dtype=tf.float32) img_depth_reshape = tf.image.resize_images(img_depth, [224, 224]) img_depth_reshape = tf.reshape(img_depth_reshape, shape=[1, 224, 224, 3]) img_depth_reshape = tf.image.rgb_to_grayscale(img_depth_reshape) rgbd_image = tf.concat([img_reshape, img_depth_reshape], 3) #print(rgbd_image.get_shape()) #rgbd_image = tf.reshape(rgbd_image, shape=[1, 224, 224, 4]) rgbd_image = normalize_depth(rgbd_image) #should in theory be RGBD with pixel values between 0-1 # Convert from 0-1 to -1 to 1 rgbd_image = tf.subtract(rgbd_image, 0.5) rgbd_image = tf.multiply(rgbd_image, 2.0) #normalise the pixels of the images as well so that it's between -1 and 1? because surely if its been trained on that, then itll work better when its like that #Up to this point this essentially, just scales the images #Just operations x_hat, y_hat, tan_hat, h_hat, w_hat = tf.unstack(inference(rgbd_image), axis=1) #uses the inference function/module bbox_hat = grasp_to_bbox(x_hat, y_hat, tan_hat, h_hat, w_hat) #some function that converts to a bounding box dg={} lg = ['w1', 'b1', 'w2', 'b2', 'w3', 'b3', 'w4', 'b4', 'w5', 'b5', 'w_fc1', 'b_fc1', 'w_fc2', 'b_fc2', 'w_output', 'b_output'] for i in lg: dg[i] = [v for v in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) if v.name == i+':0'][0] saver_g = tf.train.Saver(dg) saver_g.restore(sess, '/home/narly/zed_ur_ws_yf/src/ur_rg2/ur_control/scripts/my_models/rgbdstep-6500') #uses this model to make the grasping predictions #test = sess.run(dg[1]) #print(test) bbox_model = sess.run(bbox_hat) draw_bbox(img_show, bbox_model) #so just overlays the bounding box over the original image (all in one frame) #print(sess.run(x_hat.eval())) #this should be the bit that is given to the robot inverse kine script. grasp_representation = transfer_to_robot_inv_kine(x_hat, y_hat, tan_hat, h_hat, w_hat, sess) save_im = '/home/narly/zed_ur_ws_yf/src/ur_rg2/ur_control/scripts/live_test_image/' + str(tp) +'imagegrasp.png' cv2.imwrite(save_im, img_show) #cv2.imshow('bbox', img_show) #cv2.waitKey(0) #cv2.destroyAllWindows() return grasp_representation
def run_training(): print(FLAGS.train_or_validation) if FLAGS.train_or_validation == 'train': print('distorted_inputs') data_files_ = TRAIN_FILE images, bboxes = grasp_img_proc.distorted_inputs( [data_files_], FLAGS.num_epochs, batch_size=FLAGS.batch_size) else: print('inputs') data_files_ = VALIDATE_FILE images, bboxes = grasp_img_proc.inputs([data_files_]) x, y, degree, h, w = bboxes_to_grasps( bboxes) # x, y, tan, h, w = bboxes_to_grasps(bboxes) # images_np = np.array(images) x_hat, y_hat, degree_hat, h_hat, w_hat = tf.unstack(inference(images), axis=1) # list # x_hat, y_hat, tan_hat, h_hat, w_hat = tf.unstack(inference(images), axis=1) # list # tangent of 85 degree is 11 # tan_hat_confined = tf.minimum(11., tf.maximum(-11., tan_hat)) # tan_confined = tf.minimum(11., tf.maximum(-11., tan)) # Loss function gamma = tf.constant(0.001) loss = tf.reduce_sum( tf.pow(x_hat - x, 2) + tf.pow(y_hat - y, 2) + gamma * tf.pow(degree_hat - degree, 2) + tf.pow(h_hat - h, 2) + tf.pow(w_hat - w, 2)) train_op = tf.train.AdamOptimizer(epsilon=0.1).minimize(loss) init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) sess = tf.Session() sess.run(init_op) coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess=sess, coord=coord) #save/restore model d = {} l = [ 'w1', 'b1', 'w2', 'b2', 'w3', 'b3', 'w4', 'b4', 'w5', 'b5', 'w_fc1', 'b_fc1', 'w_fc2', 'b_fc2', 'w_fc2', 'b_fc2', 'w_fc3', 'b_fc3', 'w_output', 'b_output' ] for i in l: d[i] = [ v for v in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) if v.name == i + ':0' ][0] dg = {} lg = [ 'w1', 'b1', 'w2', 'b2', 'w3', 'b3', 'w4', 'b4', 'w5', 'b5', 'w_fc1', 'b_fc1', 'w_fc2', 'b_fc2', 'w_output', 'b_output' ] for i in lg: dg[i] = [ v for v in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) if v.name == i + ':0' ][0] saver = tf.train.Saver(d) saver_g = tf.train.Saver(dg, max_to_keep=20) #saver.restore(sess, "/root/grasp/grasp-detection/models/imagenet/m2/m2.ckpt") saver_g.restore(sess, FLAGS.model_path) try: count = 0 step = 0 start_time = time.time() while not coord.should_stop(): start_batch = time.time() #train if FLAGS.train_or_validation == 'train': _, loss_value, images_np, x_value, x_model, y_value, y_model, degree_value, degree_model, h_value, h_model, w_value, w_model = sess.run( [ train_op, loss, images, x, x_hat, y, y_hat, degree, degree_hat, h, h_hat, w, w_hat ]) # _, loss_value, images_np, x_value, x_model, y_value, y_model, tan_value, tan_model, h_value, h_model, w_value, w_model = sess.run([train_op, loss, images, x, x_hat, y, y_hat, tan, tan_hat, h, h_hat, w, w_hat]) duration = time.time() - start_batch if step % 100 == 0: print( 'Step %d | loss = %s\n | x = %s\n | x_hat = %s\n | y = %s\n | y_hat = %s\n | degree = %s\n | degree_hat = %s\n | h = %s\n | h_hat = %s\n | w = %s\n | w_hat = %s\n | (%.3f sec/batch\n' ) % (step, loss_value, x_value[:3], x_model[:3], y_value[:3], y_model[:3], degree_value[:3], degree_model[:3], h_value[:3], h_model[:3], w_value[:3], w_model[:3], duration) # print('Step %d | loss = %s\n | x = %s\n | x_hat = %s\n | y = %s\n | y_hat = %s\n | tan = %s\n | tan_hat = %s\n | h = %s\n | h_hat = %s\n | w = %s\n | w_hat = %s\n | (%.3f sec/batch\n')%(step, loss_value, x_value[:3], x_model[:3], y_value[:3], y_model[:3], tan_value[:3], tan_model[:3], h_value[:3], h_model[:3], w_value[:3], w_model[:3], duration) # cv2.imshow('bbox', images_np[1]) ################################### # print(images_np[1]) # cv2.imwrite("./"+str(step)+".png", images_np[1]); cv2.waitKey(1) if step % 1000 == 0: filename = MODEL_SAVE_PATH + '/step_' + str(step) if not os.path.exists(filename): os.mkdir(filename) saver.save(sess, filename + '/m4.ckpt') # cv2.imshow('bbox', images_np[1]) ################################### # cv2.waitKey(0) # print(images_np[0]) # if step == 100000: # sess.close() # else: # bbox_hat = grasp_to_bbox(x_hat, y_hat, tan_hat, h_hat, w_hat) # bbox_value, bbox_model, x_value, x_model, y_value, y_model, tan_value, tan_model, h_value, h_model, w_value, w_model = sess.run([bboxes, bbox_hat, x, x_hat, y, y_hat, tan, tan_hat, h, h_hat, w, w_hat]) # bbox_value = np.reshape(bbox_value, -1) # bbox_value = [(bbox_value[0],bbox_value[1]),(bbox_value[2],bbox_value[3]),(bbox_value[4],bbox_value[5]),(bbox_value[6],bbox_value[7])] # bbox_value = [(bbox_value[0]*0.35,bbox_value[1]*0.47),(bbox_value[2]*0.35,bbox_value[3]*0.47),(bbox_value[4]*0.35,bbox_value[5]*0.47),(bbox_value[6]*0.35,bbox_value[7]*0.47)] # p1 = Polygon(bbox_value) # p2 = Polygon(bbox_model) # iou = p1.intersection(p2).area / (p1.area +p2.area -p1.intersection(p2).area) # angle_diff = np.abs(np.arctan(tan_model)*180/np.pi -np.arctan(tan_value)*180/np.pi) # duration = time.time() -start_batch # # if angle_diff < 30. and iou >= 0.25: # count+=1 # print('image: %d | duration = %.2f | count = %d | iou = %.2f | angle_difference = %.2f' %(step, duration, count, iou, angle_diff)) # print('x=',x_value,x_model,' y=',y_value, y_model,' tan=',tan_value, tan_model,' h=', h_value, h_model,' w=', w_value, w_model) step += 1 except tf.errors.OutOfRangeError: print('Done training for %d epochs, %d steps, %.1f min.' % (FLAGS.num_epochs, step, (time.time() - start_time) / 60)) finally: coord.request_stop() coord.join(threads) sess.close()
cv2.line(img, p2, p3, (0, 0, 255)) cv2.line(img, p3, p4, (0, 0, 255)) cv2.line(img, p4, p1, (0, 0, 255)) init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) sess = tf.Session() sess.run(init_op) dg = {} img_raw_data = tf.placeholder(tf.string) img_data = tf.image.decode_jpeg(img_raw_data) img_data = tf.image.convert_image_dtype(img_data, dtype=tf.float32) img_reshape = tf.image.resize_images(img_data, [224, 224]) img_reshape = tf.reshape(img_reshape, shape=[1, 224, 224, 3]) x_hat, y_hat, tan_hat, w_hat, h_hat = tf.unstack(inference(img_reshape), axis=1) bbox_hat = grasp_to_bbox(x_hat, y_hat, tan_hat, h_hat, w_hat) while True: with contextlib.closing( mmap.mmap(-1, 1024, tagname='grasp_det', access=mmap.ACCESS_READ)) as m: filename = m.read(1024).decode().replace('\x00', '') if filename != prev_filename: img_show = cv2.imread(filename) with open(filename, 'rb') as f: img = f.read() print(filename + " received") if dg == {}:
def run_training(): print(FLAGS.train_or_validation) if FLAGS.train_or_validation == 'train': print('distorted_inputs') data_files_ = TRAIN_FILE images, bboxes = grasp_img_proc.distorted_inputs( [data_files_], FLAGS.num_epochs, batch_size=FLAGS.batch_size) else: print('inputs') data_files_ = VALIDATE_FILE images, bboxes = grasp_img_proc.inputs([data_files_]) x, y, tan, h, w = bboxes_to_grasps(bboxes) x_hat, y_hat, tan_hat, h_hat, w_hat = tf.unstack(inference(images), axis=1) # list # tangent of 85 degree is 11 tan_hat_confined = tf.minimum(11., tf.maximum(-11., tan_hat)) tan_confined = tf.minimum(11., tf.maximum(-11., tan)) # Loss function gamma = tf.constant(10.) loss = tf.reduce_sum( tf.pow(x_hat - x, 2) + tf.pow(y_hat - y, 2) + gamma * tf.pow(tan_hat_confined - tan_confined, 2) + tf.pow(h_hat - h, 2) + tf.pow(w_hat - w, 2)) train_op = tf.train.AdamOptimizer(epsilon=0.1).minimize(loss) init_op = tf.group(tf.global_variables_initializer(), tf.local_variables_initializer()) sess = tf.Session() sess.run(init_op) coord = tf.train.Coordinator() threads = tf.train.start_queue_runners(sess=sess, coord=coord) #save/restore model d = {} l = [ 'w1', 'b1', 'w2', 'b2', 'w3', 'b3', 'w4', 'b4', 'w5', 'b5', 'w_fc1', 'b_fc1', 'w_fc2', 'b_fc2' ] for i in l: d[i] = [ v for v in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) if v.name == i + ':0' ][0] dg = {} lg = [ 'w1', 'b1', 'w2', 'b2', 'w3', 'b3', 'w4', 'b4', 'w5', 'b5', 'w_fc1', 'b_fc1', 'w_fc2', 'b_fc2', 'w_output', 'b_output' ] for i in lg: dg[i] = [ v for v in tf.get_collection(tf.GraphKeys.GLOBAL_VARIABLES) if v.name == i + ':0' ][0] saver = tf.train.Saver(d) saver_g = tf.train.Saver(dg) #saver.restore(sess, "/root/grasp/grasp-detection/models/imagenet/m2/m2.ckpt") saver_g.restore(sess, FLAGS.model_path) try: count = 0 step = 0 start_time = time.time() while not coord.should_stop(): start_batch = time.time() #train if FLAGS.train_or_validation == 'train': _, loss_value, x_value, x_model, tan_value, tan_model, h_value, h_model, w_value, w_model = sess.run( [ train_op, loss, x, x_hat, tan, tan_hat, h, h_hat, w, w_hat ]) duration = time.time() - start_batch if step % 100 == 0: print( 'Step %d | loss = %s\n | x = %s\n | x_hat = %s\n | tan = %s\n | tan_hat = %s\n | h = %s\n | h_hat = %s\n | w = %s\n | w_hat = %s\n | (%.3f sec/batch\n' ) % (step, loss_value, x_value[:3], x_model[:3], tan_value[:3], tan_model[:3], h_value[:3], h_model[:3], w_value[:3], w_model[:3], duration) if step % 1000 == 0: saver_g.save(sess, FLAGS.model_path) else: bbox_hat = grasp_to_bbox(x_hat, y_hat, tan_hat, h_hat, w_hat) bbox_value, bbox_model, tan_value, tan_model = sess.run( [bboxes, bbox_hat, tan, tan_hat]) bbox_value = np.reshape(bbox_value, -1) bbox_value = [(bbox_value[0] * 0.35, bbox_value[1] * 0.47), (bbox_value[2] * 0.35, bbox_value[3] * 0.47), (bbox_value[4] * 0.35, bbox_value[5] * 0.47), (bbox_value[6] * 0.35, bbox_value[7] * 0.47)] p1 = Polygon(bbox_value) p2 = Polygon(bbox_model) iou = p1.intersection(p2).area / (p1.area + p2.area - p1.intersection(p2).area) angle_diff = np.abs( np.arctan(tan_model) * 180 / np.pi - np.arctan(tan_value) * 180 / np.pi) duration = time.time() - start_batch if angle_diff < 30. and iou >= 0.25: count += 1 print( 'image: %d | duration = %.2f | count = %d | iou = %.2f | angle_difference = %.2f' % (step, duration, count, iou, angle_diff)) step += 1 except tf.errors.OutOfRangeError: print('Done training for %d epochs, %d steps, %.1f min.' % (FLAGS.num_epochs, step, (time.time() - start_time) / 60)) finally: coord.request_stop() coord.join(threads) sess.close()