def train(studies_dir, list_of_dirs=False):
tf.reset_default_graph()
tracker = Tracker().build()
studies_handler = StudiesHandler(studies_dir, list_of_dirs=list_of_dirs)
create_dir(checkpoint_dir)
saver = tf.train.Saver()
init = tf.global_variables_initializer()
keep_prob = 0.8
# update_prob_rate = epoch_num*0.9 / 6
with tf.Session() as sess:
sess.run(init)
if restore_model:
print('Loading Model...')
saver.restore(sess, checkpoint_file)
for i in range(epoch_num):
# if i == update_prob_rate:
# update_prob_rate *= 2
# keep_prob += 0.05
study = studies_handler.get_study()
print("\n#####################################################")
print("{}) study: {}".format(i, study.study_dir))
print("#####################################################")
while not study.is_end_of_study():
minibatch_size = 0
images = []
points = []
while not study.is_end_of_study() and minibatch_size < batch_size:
image1, image2, point = study.next_couple()
if image2 is None:
break
minibatch_size += 1
images.append(image1.reshape((image_height, image_width, 3)))
images.append(image2.reshape((image_height, image_width, 3)))
points.append([point.x, point.y])
if len(images) == 0:
continue
_, loss, p = sess.run([tracker.optimizer, tracker.loss, tracker.predict],
feed_dict={tracker.minibatch_size: minibatch_size,
tracker.input_frames: images,
tracker.target: points,
tracker.keep_prob: keep_prob})
p = np.array(p)
print("------------------------------------------------------------------")
print("loss = {}".format(loss))
for i in range(len(points)):
print("predict: {}, relative: {}".format(p[i], points[i]))
saver.save(sess, checkpoint_file)
def train(studies_dir, list_of_dirs=False):
tf.reset_default_graph()
tracker = Tracker().build()
studies_handler = StudiesHandler(studies_dir, list_of_dirs=list_of_dirs)
create_dir(checkpoint_dir)
saver = tf.train.Saver()
init = tf.global_variables_initializer()
keep_prob = 1
# update_prob_rate = epoch_num*0.9 / 6
with tf.Session() as sess:
sess.run(init)
if restore_model:
print('Loading Model...')
saver.restore(sess, checkpoint_file)
for i in range(epoch_num):
# if i == update_prob_rate:
# update_prob_rate *= 2
# keep_prob += 0.05
study = studies_handler.get_study()
print("\n#####################################################")
print("{}) study: {}".format(i, study.study_dir))
print("#####################################################")
current_cell_state = np.zeros((batch_size, rnn_state_size))
current_hidden_state = np.zeros((batch_size, rnn_state_size))
warm_frames = study.get_warm_frames()
for frame in warm_frames:
next_state = sess.run(tracker.out_state, feed_dict={tracker.input_frames: [frame],
tracker.keep_prob: keep_prob,
tracker.cell_state: current_cell_state,
tracker.hidden_state: current_hidden_state})
current_cell_state, current_hidden_state = next_state
# old_state = current_cell_state
while not study.is_end_of_study():
image, point = study.next()
point = [point.x, point.y]
# _, loss, next_state, p, check = sess.run([tracker.optimizer, tracker.loss, tracker.out_state, tracker.predict, tracker.check],
check = sess.run(tracker.check,
feed_dict={tracker.input_frames: image,
tracker.target: [point],
tracker.keep_prob: keep_prob,
tracker.cell_state: current_cell_state,
tracker.hidden_state: current_hidden_state})
print(check)
# old_state = current_cell_state
# print("predict: {}, relative: {}, loss = {}".format(np.array(p), point, loss))
# print("check: {}".format(check))
saver.save(sess, checkpoint_file)
def __init__(self, work_dir, fingers_functions=False):
pygame.init()
self.screen = pygame.display.set_mode(
(0, 0), pygame.NOFRAME | pygame.FULLSCREEN)
self.screen.fill(Board.WHITE)
self.functionGen = FunctionGenerator(fingers_functions)
create_dir(work_dir)
self.work_folder = work_dir
self.cap = cv2.VideoCapture(0)
self.frameWriter = None
def extract_frames_from_video(video_path,
amount_of_frames=None,
start_index_name=0,
warm_frames=None,
fps=None):
out_dir = os.path.abspath(
os.path.join(os.path.join(video_path, os.pardir), "frames"))
create_dir(out_dir)
cap = cv2.VideoCapture(video_path)
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
if amount_of_frames is not None:
taken_frames = total_frames // amount_of_frames
if taken_frames == 0:
taken_frames = 1
remainder = total_frames / amount_of_frames - taken_frames
elif fps is not None:
amount = (total_frames * fps) / cap.get(cv2.CAP_PROP_FPS)
taken_frames = total_frames // amount
if taken_frames == 0:
taken_frames = 1
remainder = total_frames / amount - taken_frames
else:
taken_frames = 1
remainder = 0
if warm_frames is not None:
warm_frames -= 1
success, image = cap.read()
index = start_index_name
count = 0
remainder_count = 0
success = True
while success:
if count == 0:
remainder_count += remainder
if remainder_count >= 1:
count = taken_frames + 1
remainder_count -= 1
else:
count = taken_frames
cv2.imwrite(os.path.join(out_dir, "{}.jpg".format(index)), image)
count -= 1
success, image = cap.read()
index += 1
if warm_frames is not None and warm_frames > 0:
count = 0
remainder_count = 0
warm_frames -= 1
def test(studies_dir, list_of_dirs=False):
tf.reset_default_graph()
tracker = Tracker().build()
studies_handler = StudiesHandler(studies_dir, list_of_dirs=list_of_dirs)
create_dir(checkpoint_dir)
saver = tf.train.Saver()
init = tf.global_variables_initializer()
keep_prob = 1
with tf.Session() as sess:
sess.run(init)
print('Loading Model...')
saver.restore(sess, checkpoint_file)
while not studies_handler.epoch_done:
study = studies_handler.get_study()
print("\n#####################################################")
print("study: {}".format(study.study_dir))
print("#####################################################")
current_cell_state = np.zeros((batch_size, rnn_state_size))
current_hidden_state = np.zeros((batch_size, rnn_state_size))
warm_frames = study.get_warm_frames()
for frame in warm_frames:
next_state = sess.run(tracker.out_state, feed_dict={tracker.input_frames: [frame],
tracker.keep_prob: keep_prob,
tracker.cell_state: current_cell_state,
tracker.hidden_state: current_hidden_state})
current_cell_state, current_hidden_state = next_state
old_state = current_cell_state
while not study.is_end_of_study():
image, point = study.next()
point = [point.x, point.y]
loss, next_state, p, check = sess.run([tracker.loss, tracker.out_state, tracker.predict, tracker.check],
feed_dict={tracker.input_frames: image,
tracker.target: [point],
tracker.keep_prob: keep_prob,
tracker.cell_state: current_cell_state,
tracker.hidden_state: current_hidden_state})
current_cell_state, current_hidden_state = next_state
def __init__(self, studies_dir, stats_dir, list_of_dirs=False):
self.WRITE_STATS_STEPS = 100
if list_of_dirs:
studies = []
for dir in studies_dir:
studies += dir_to_subdir_list(dir)
self.studies = studies
else:
self.studies = dir_to_subdir_list(studies_dir)
self.current_study_index = 0
self.current_study = None
self.rewards_counter = 0
self.completed_paths_counter = 0
delete_dir(stats_dir)
create_dir(stats_dir)
self.steps_counter = 0
self.rewards = []
self.path_len = []
self.path_completed = []
self.rewards_file = os.path.join(stats_dir, "rewards.txt")
self.path_len_file = os.path.join(stats_dir, "path_len.txt")
self.path_completed_file = os.path.join(stats_dir, "path_completed.txt")
def record_path(self, work_folder, path_index):
done = False
curr_t = 0.0
start = False
path = []
while not done:
if start:
if time.time() * 1000 - curr_t >= 1000 / Board.FILM_FPS:
curr_t = time.time() * 1000
pos = pygame.mouse.get_pos()
pygame.draw.circle(self.screen, Board.RED, pos, 3)
path.append(pos)
for event in pygame.event.get():
if event.type == pygame.MOUSEBUTTONUP:
start = True
if event.type == pygame.KEYUP:
if event.key == pygame.K_LCTRL:
start = False
base_dir = os.path.join(work_folder, str(path_index))
create_dir(base_dir)
Board.write_path_to_file(
os.path.join(base_dir, "Path.txt"), path)
pygame.image.save(
self.screen,
os.path.join(base_dir, "screenshot.jpg"))
self.screen.fill(Board.BACKGROUND)
path_index += 1
path = []
if event.key == pygame.K_ESCAPE:
done = True
pygame.display.flip()
def train(studies_dir, list_of_dirs=False):
tf.reset_default_graph()
# We define the cells for the main and target q-networks
cell = tf.contrib.rnn.BasicLSTMCell(num_units=units_size, state_is_tuple=True)
cell_t = tf.contrib.rnn.BasicLSTMCell(num_units=units_size, state_is_tuple=True)
main_net, main_net_vars = DRQN(cell, 'main_DRQN', reuse_encoder_variables=False).build_graph()
target_net, target_net_vars = DRQN(cell_t, 'target_DRQN', reuse_encoder_variables=True).build_graph()
init = tf.global_variables_initializer()
encoder_loader = DRQN.load_encoder_weights()
saver = tf.train.Saver(main_net_vars + target_net_vars)
create_dir(checkpoint_dir)
target_ops = DRQN.update_target_graph(tf.trainable_variables(), update_target_ratio)
experienced_studies_buffer = ExperiencedStudiesBuffer(500)
environment = Environment(studies_dir, stats_dir, list_of_dirs=list_of_dirs)
# Set the rate of random action decrease.
e = startE
step_drop = (startE - endE) / annealing_steps
total_steps = 0
with tf.Session() as sess:
sess.run(init)
encoder_loader.restore(sess, encoder_checkpoint_file)
if restore_model:
print('Loading Model...')
saver.restore(sess, checkpoint_file)
DRQN.update_target(target_ops, sess) # Set the target network to be equal to the main network.
for i in range(epoch_num):
print("study number: {}".format(i))
# Reset environment and get warm frames
warm_frames = environment.reset()
# Warm up the recurrent layer's hidden state
lstm_state = (np.zeros([1, units_size]), np.zeros([1, units_size]))
lstm_state = DRQN.warm_lstm_state(sess, main_net, lstm_state, warm_frames)
# Get first state
s = environment.start()
j = 0
# The Q-Network
while j < max_study_len:
j += 1
# Choose an action randomly from the Q-network
if np.random.rand(1) < e or total_steps < pre_train_steps:
lstm_state1 = sess.run(main_net.rnn_state, feed_dict={main_net.input_frames: s,
main_net.trace_len: 1,
main_net.state_in: lstm_state,
main_net.batch_size: 1})
a = np.random.randint(0, actions_num)
else:
a, lstm_state1 = sess.run([main_net.predict, main_net.rnn_state],
feed_dict={main_net.input_frames: s,
main_net.trace_len: 1,
main_net.state_in: lstm_state,
main_net.batch_size: 1})
a = a[0]
s1, r, d = environment.step(a)
total_steps += 1
if total_steps > pre_train_steps:
if e > endE:
e -= step_drop
if total_steps % update_freq == 0:
DRQN.update_target(target_ops, sess)
experienced_studies_buffer.select_study()
warm_frames = experienced_studies_buffer.get_warm_frames()
done_trace = False
# Reset the recurrent layer's hidden state
state_train = (np.zeros([mb_size, units_size]), np.zeros([mb_size, units_size]))
main_state_train = DRQN.warm_lstm_state(sess, main_net, state_train, warm_frames)
target_state_train = DRQN.warm_lstm_state(sess, target_net, state_train, warm_frames)
while not done_trace:
# train_batch contains list of [s, a, r, s1, d]
train_batch, trace_length, done_trace = experienced_studies_buffer.get_trace()
if train_batch is None:
break
# update the Double-DQN to the target Q-values
main_next_state, Q1 = sess.run([main_net.rnn_state, main_net.predict],
feed_dict={main_net.input_frames: np.vstack(train_batch[:, 3]),
main_net.trace_len: trace_length,
main_net.state_in: main_state_train,
main_net.batch_size: mb_size})
target_next_state, Q2 = sess.run([target_net.rnn_state, target_net.Qout],
feed_dict={target_net.input_frames: np.vstack(train_batch[:, 3]),
target_net.trace_len: trace_length,
target_net.state_in: target_state_train,
target_net.batch_size: mb_size})
# end_multiplier = -(train_batch[:, 4] - 1)
# doubleQ = Q2[range(mb_size * trace_length), Q1]
targetQ = train_batch[:, 2] # + (gamma * doubleQ * end_multiplier)
# Update the network with our target values.
_, loss = sess.run([main_net.optimizer, main_net.loss], feed_dict={main_net.input_frames: np.vstack(train_batch[:, 0]),
main_net.target_Q: targetQ,
main_net.actions: train_batch[:, 1],
main_net.trace_len: trace_length,
main_net.state_in: state_train,
main_net.batch_size: mb_size})
print("len: {},\tloss: {}".format(trace_length, loss))
main_state_train = main_next_state
target_state_train = target_next_state
s = s1
lstm_state = lstm_state1
if d:
break
experienced_studies_buffer.append(environment.current_study)
if i % 100 == 0 and i != 0 and i > pre_train_steps:
saver.save(sess, checkpoint_file)
saver.save(sess, checkpoint_file)
path_len = environment.current_study.image_index - warm_frames_num
print("{}:\t len {},\t done: {}".format(study_name, path_len, environment.current_study.finished_successfully))
environment.write_stats()
if __name__ == '__main__':
# if len(sys.argv) == 1:
# print("Error: missing run mode type (hand/fingers)")
# sys.exit(1)
#
# run_mode_type = sys.argv[1]
# set_run_config(run_mode_type)
sys.stdout = StdoutLog(log_dir, sys.stdout, print_to_log, print_to_stdout)
create_dir(checkpoint_dir)
if mode == "train":
if platform.system() == 'Linux':
if mode_type == "hand":
print("use hands dataset..")
train(["/home/ami/handsTrack/studies/full_hands/17.12.17",
"/home/ami/handsTrack/studies/full_hands/18.12.17",
"/home/ami/handsTrack/studies/full_hands/19.12.17"],
list_of_dirs=True)
else:
print("use fingers dataset..")
train("/home/ami/handsTrack/studies/fingers/20.12.17")
else:
if mode_type == "hand":
def play(self, index):
i = index
done = False
while not done:
ret, frame = self.cap.read()
if self.frameWriter is not None and ret:
frame = cv2.flip(frame, 1)
self.frameWriter.write(frame)
for event in pygame.event.get():
if event.type == pygame.KEYUP:
if event.key == pygame.K_LCTRL:
self.screen.fill(Board.WHITE)
x_points, y_points = self.functionGen.generate_function(
)
i += 1
video_index = 1
base_dir = os.path.join(self.work_folder, str(i))
create_dir(base_dir)
create_dir(os.path.join(base_dir, "frames"))
self._draw_function(
x_points, y_points,
os.path.join(base_dir, Board.PATH_NAME))
pygame.image.save(
self.screen,
os.path.join(base_dir, "screenshot.jpg"))
if self.frameWriter is not None:
self.frameWriter.release()
self.frameWriter = None
if self.frameWriter is None:
pygame.draw.rect(self.screen, Board.RED,
(0, 0, 20, 20))
else:
pygame.draw.rect(self.screen, Board.GREEN,
(0, 0, 20, 20))
if event.key == pygame.K_LSHIFT:
self.frameWriter = cv2.VideoWriter(
os.path.join(self.work_folder, str(i),
"video_{}.avi".format(video_index)),
cv2.VideoWriter_fourcc(*"XVID"), Board.FILM_FPS,
(640, 480))
pygame.draw.rect(self.screen, Board.GREEN,
(0, 0, 20, 20))
if event.key == pygame.K_z:
if self.frameWriter is not None:
self.frameWriter.release()
video_index += 1
self.frameWriter = cv2.VideoWriter(
os.path.join(self.work_folder, str(i),
"video_{}.avi".format(video_index)),
cv2.VideoWriter_fourcc(*"XVID"), Board.FILM_FPS,
(640, 480))
pygame.draw.rect(self.screen, Board.GREEN,
(0, 0, 20, 20))
if event.type == pygame.MOUSEBUTTONUP or event.type == pygame.K_ESCAPE:
done = True
self.cap.release()
cv2.destroyAllWindows()
if self.frameWriter is not None:
self.frameWriter.release()
self.frameWriter = None
pygame.draw.rect(self.screen, Board.RED,
(0, 0, 20, 20))
pygame.display.flip()
def train(studies_dir, list_of_dirs=False):
tf.reset_default_graph()
tracker = Tracker().build()
studies_handler = StudiesHandler(studies_dir, list_of_dirs=list_of_dirs)
create_dir(checkpoint_dir)
saver = tf.train.Saver()
init = tf.global_variables_initializer()
keep_prob = 1
with tf.Session() as sess:
sess.run(init)
if restore_model:
print('Loading Model...')
saver.restore(sess, checkpoint_file)
for i in range(epoch_num):
study = studies_handler.get_study()
print("{}) study: {}".format(i, study.study_dir))
current_cell_state = np.zeros((batch_size, rnn_state_size))
current_hidden_state = np.zeros((batch_size, rnn_state_size))
warm_frames = study.get_warm_frames()
for frame in warm_frames:
next_state = sess.run(tracker.out_state,
feed_dict={
tracker.input_frames: [frame],
tracker.keep_prob:
keep_prob,
tracker.cell_state:
current_cell_state,
tracker.hidden_state:
current_hidden_state
})
current_cell_state, current_hidden_state = next_state
while not study.is_end_of_study():
len = 0
images = []
points = []
while not study.is_end_of_study() and len < seq_len:
len += 1
image, point = study.next()
images.append(image.reshape(
(image_height, image_width, 3)))
points.append([point.x, point.y])
# _, loss, next_state, p = sess.run([tracker.optimizer, tracker.loss, tracker.out_state, tracker.predict],
check = sess.run(tracker.check,
feed_dict={
tracker.input_frames: images,
tracker.target: points,
tracker.keep_prob: keep_prob,
tracker.cell_state: current_cell_state,
tracker.hidden_state: current_hidden_state
})
# current_cell_state, current_hidden_state = next_state
# print("predict: {}, loss = {}".format(np.array(p), loss))
print(np.array(check).shape)
saver.save(sess, checkpoint_file)
def play(self, path_folder, work_folder):
done = start = False
t = 0.0
path_files = dir_to_subdir_list(path_folder)
path = []
path_index = 0
study_index = 132
# <-
repeat_index = Board.REPEAT
while not done:
if start:
if time.time() * 1000 - t >= 1000 / Board.FILM_FPS:
t = time.time() * 1000
_, frame = self.cap.read()
self.frameWriter.write(frame)
self.screen.fill(Board.BACKGROUND)
pygame.draw.circle(self.screen, Board.RED,
[*path[path_index]], 8)
path_index = path_index + 1
if path_index == len(path):
pygame.draw.rect(self.screen, Board.RED,
(800, 0, 20, 20))
start = False
for event in pygame.event.get():
if event.type == pygame.KEYUP:
if event.key == pygame.K_LSHIFT:
study_index += 1
base_dir = os.path.join(work_folder, str(study_index))
create_dir(base_dir)
Board.write_path_to_file(
os.path.join(base_dir, "Path.txt"), path)
if self.frameWriter is not None:
self.frameWriter.release()
self.frameWriter = cv2.VideoWriter(
os.path.join(work_folder, str(study_index),
"video.avi"),
cv2.VideoWriter_fourcc(*"XVID"), Board.FILM_FPS,
(640, 480))
self.cap.read() # clear old frame
for i in range(Board.warm_frames_amount):
_, frame = self.cap.read()
self.frameWriter.write(frame)
start = True
if event.key == pygame.K_LCTRL:
self.screen.fill(Board.BACKGROUND)
path_index = 0
start = False
if repeat_index == Board.REPEAT:
repeat_index = 0
current_file = np.random.choice(path_files)
repeat_index += 1
path = Board.read_path_from_file(
os.path.join(current_file, "Path.txt"))
pygame.draw.circle(self.screen, Board.RED, [*path[0]],
8)
for i in range(len(path) - 1):
pygame.draw.line(self.screen, Board.BLUE,
[*path[i]], [*path[i + 1]], 3)
if event.key == pygame.K_d:
self.screen.fill(Board.BACKGROUND)
path_index = 0
start = False
create_dir(
os.path.join(work_folder, str(study_index),
"DELETE"))
path = Board.read_path_from_file(
os.path.join(current_file, "Path.txt"))
pygame.draw.circle(self.screen, Board.RED, [*path[0]],
8)
for i in range(len(path) - 1):
pygame.draw.line(self.screen, Board.BLUE,
[*path[i]], [*path[i + 1]], 3)
if event.key == pygame.K_ESCAPE:
done = True
self.cap.release()
cv2.destroyAllWindows()
pygame.display.flip()
def train(studies_dir, list_of_dirs=False):
tf.reset_default_graph()
tracker = Tracker().build()
studies_handler = StudiesHandler(studies_dir, list_of_dirs=list_of_dirs)
create_dir(checkpoint_dir)
saver = tf.train.Saver()
init = tf.global_variables_initializer()
keep_prob = 0.8
with tf.Session() as sess:
sess.run(init)
if restore_model:
print('Loading Model...')
saver.restore(sess, checkpoint_file)
for i in range(epoch_num):
study = studies_handler.get_study()
print("\n#####################################################")
print("{}) study: {}".format(i, study.study_dir))
print("#####################################################")
while not study.is_end_of_study():
minibatch_size = 0
images = []
points = []
while not study.is_end_of_study(
) and minibatch_size < batch_size:
image1, image2, point = study.next_couple()
if image2 is None:
break
minibatch_size += 1
images.append(
image1.reshape((image_height, image_width, 3)))
images.append(
image2.reshape((image_height, image_width, 3)))
if point.x == 0 and point.y == 0:
points.append([number_of_directions])
continue
point_np = np.array([point.x, point.y])
point_np = point_np / np.sqrt(point.x**2 + point.y**2)
direction_index = get_direction_index(directions, point_np)
points.append([direction_index])
if len(images) == 0:
continue
_, loss = sess.run(
[tracker.optimizer, tracker.loss],
feed_dict={
tracker.minibatch_size: minibatch_size,
tracker.input_frames: images,
tracker.target: points,
tracker.keep_prob: keep_prob
})
print("loss = {}".format(loss))
saver.save(sess, checkpoint_file)
def train(studies_dir, list_of_dirs=False):
tf.reset_default_graph()
tracker = Tracker(is_training=True).build()
studies_handler = StudiesHandler(studies_dir, list_of_dirs=list_of_dirs)
create_dir(checkpoint_dir)
saver = tf.train.Saver()
init = tf.global_variables_initializer()
keep_prob = 1
with tf.Session() as sess:
sess.run(init)
if restore_model:
print('Loading Model...')
saver.restore(sess, checkpoint_file)
for i in range(epoch_num):
loss_array = []
study = studies_handler.get_study()
print("\n#####################################################")
print("{}) study: {}".format(i, study.study_dir))
print("#####################################################")
while not study.is_end_of_study():
minibatch_size = 0
images = []
targets = []
while not study.is_end_of_study(
) and minibatch_size < batch_size:
image1, image2, point = study.next_couple()
if image2 is None:
break
minibatch_size += 1
image1 = image1.reshape(
(image_height, image_width, input_channels))
image2 = image2.reshape(
(image_height, image_width, input_channels))
images.append(np.concatenate((image1, image2), axis=2))
targets.append([point.x, point.y])
if len(images) == 0:
continue
_, loss, predict = sess.run(
[tracker.optimizer, tracker.loss, tracker.predict],
feed_dict={
tracker.minibatch_size: minibatch_size,
tracker.input_frames: images,
tracker.target: targets,
tracker.keep_prob: keep_prob
})
print(
"------------------------------------------------------------------"
)
print("loss = {}".format(loss))
loss_array.append(loss)
for j in range(len(targets)):
print("predict: {}, relative: {}".format(
predict[j], targets[j]))
write_list_to_file(loss_array, loss_file)
if i % 100 == 0:
saver.save(sess, checkpoint_file)