def main(argv=None):
print ('Use seg', FLAGS.use_seg)
input_img = tf.placeholder(tf.float32, [1, FLAGS.height, FLAGS.width, 3])
# Create network.
net = DeepLabResNetModel({'data': input_img}, is_training=False, num_classes=NUM_CLASSES)
# Which variables to load.
restore_var = tf.global_variables()
# Predictions.
raw_output = net.layers['fc_out']
raw_output_up = tf.image.resize_bilinear(raw_output, tf.shape(input_img)[1:3,])
raw_output_up = tf.argmax(raw_output_up, dimension=3)
pred = tf.expand_dims(raw_output_up, dim=3)
# Set up TF session and initialize variables.
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
init = tf.global_variables_initializer()
sess.run(init)
# Load weights.
ckpt = tf.train.get_checkpoint_state(RESTORE_PATH)
if ckpt and ckpt.model_checkpoint_path:
loader = tf.train.Saver(var_list=restore_var)
load_step = int(os.path.basename(ckpt.model_checkpoint_path).split('-')[1])
load(loader, sess, ckpt.model_checkpoint_path)
else:
print('No checkpoint file found.')
load_step = 0
p = predictor.getPredictor(load_path=FLAGS.path, transform=False)
# create server
# TODO: Remove socket dependency
# ---begin---
server = serverSock(name=FLAGS.name)
server.create(port=FLAGS.port)
server.waitForClient()
# ---end---
img_counter = 0
while True:
print('wait for task')
# TODO: img from ros callback function
img = server.recv()
print('receive task')
img = np.fromstring(img, np.uint8)
img = cv2.imdecode(img, 1)
hi, wi, _ = img.shape
store_img = img
img_counter += 1
img = cv2.resize(img, (FLAGS.width, FLAGS.height)).astype(float)
raw_img_resize = img
img = img - IMG_MEAN
img = np.expand_dims(img, axis=0)
#### segmentation
preds = sess.run(pred, feed_dict={input_img: img})
####
if FLAGS.use_seg:
msk = decode_labels(preds, num_classes=NUM_CLASSES)
else:
msk = raw_img_resize
msk = cv2.resize(msk, (84, 84))
#z = np.zeros(shape=(84, 84, 4), dtype=np.float32)
#z[:, :, 0:3] = msk
if FLAGS.rtimg == True:
msk = cv2.resize(msk, (wi, hi))
_, msk = cv2.imencode('.png', msk)
msk = msk.tostring()
# TODO: No use
server.send(msk)
else:
msk = cv2.resize(msk, (84, 84))
act = p(msk)
cv2.imwrite('seg_image/raw_img_%05d_%d.png' % (img_counter, act), store_img)
print('predict: {}'.format(act))
# if img_counter==5:
# time.sleep(5)
act = pickle.dumps(act)
# TOOD: Send action by ROS
server.send(act)
print('task done')
# TODO: No use
server.close()
def main(argv=None):
# Input placeholder
input_img = tf.placeholder(tf.float32, [1, FLAGS.height, FLAGS.width, 3])
# Create network.
net = DeepLabResNetModel({'data': input_img},
is_training=False,
num_classes=NUM_CLASSES)
# Which variables to load.
restore_var = tf.global_variables()
# Predictions.
raw_output = net.layers['fc_out']
raw_output_up = tf.image.resize_bilinear(raw_output,
tf.shape(input_img)[1:3, ])
raw_output_up = tf.argmax(raw_output_up, dimension=3)
# Color transform
color_mat = label_colours[..., [2, 1, 0]]
color_mat = tf.constant(color_mat, dtype=tf.float32)
onehot_output = tf.one_hot(raw_output_up, depth=len(label_colours))
onehot_output = tf.reshape(onehot_output, (-1, len(label_colours)))
pred = tf.matmul(onehot_output, color_mat)
pred = tf.reshape(pred, (1, FLAGS.height, FLAGS.width, 3))
# Set up TF session and initialize variables.
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
init = tf.global_variables_initializer()
sess.run(init)
# Load weights.
ckpt = tf.train.get_checkpoint_state(RESTORE_PATH)
if ckpt and ckpt.model_checkpoint_path:
loader = tf.train.Saver(var_list=restore_var)
load_step = int(
os.path.basename(ckpt.model_checkpoint_path).split('-')[1])
load(loader, sess, ckpt.model_checkpoint_path)
else:
print('No checkpoint file found.')
load_step = 0
#p = predictor.getSingleFramePredictor(load_path=FLAGS.path, transform=False)
p = predictor.getPredictor(load_path=FLAGS.path, transform=False)
scale = 1.0
ros_rmp_motion_publisher = RosKobukiMotionPublisher(linear=FLAGS.lin,
angular=FLAGS.ang)
#ros_rmp_motion_publisher = RosRMPMotionPublisher(linear=FLAGS.lin, angular=FLAGS.ang)
def actor(stop_event):
global image_queue
while len(image_queue) == 0:
print('Wait for image queue filled ... (sleep 1)')
time.sleep(1)
pass
timestep = 0
while not stop_event.is_set():
start = time.time()
img = image_queue[0]
raw_img = img
#raw_img = cv2.resize(img, (FLAGS.width, FLAGS.height)).astype(float)
img = raw_img - IMG_MEAN
img = np.expand_dims(img, axis=0)
if not FLAGS.use_depth:
preds = sess.run(pred, feed_dict={input_img: img})
if FLAGS.use_seg:
s = preds[0].astype(np.uint8)
msk_img = s
s = cv2.resize(s, (84, 84))
else:
s = raw_img
msk_img = s
s = cv2.resize(s, (84, 84))
if FLAGS.use_depth:
s = np.expand_dims(s, axis=-1)
else:
s = cv2.cvtColor(s, cv2.COLOR_BGR2RGB)
act = p(s)
ros_rmp_motion_publisher.publish(act)
end = time.time()
print('Inference time = %f' % (end - start))
if timestep < 3000:
#raw_img_resize = cv2.resize(raw_img, (84, 84))
s = cv2.cvtColor(s, cv2.COLOR_RGB2BGR)
cv2.imwrite('imgs/raw_img_%05d.png' % (timestep), raw_img)
cv2.imwrite('imgs/msk_img_%05d.png' % (timestep), msk_img)
timestep += 1
print("STEP: {}".format(timestep))
def image_subscribe_callback(img):
global image_queue
image_queue.append(img)
if FLAGS.use_depth:
ros_camera_image_subscriber = RosDepthImageSubscriber(
user_callback=image_subscribe_callback)
else:
ros_camera_image_subscriber = RosImageSubscriber(
user_callback=image_subscribe_callback)
rospy.init_node('agent', anonymous=True)
actor_thread_stop_event = threading.Event()
try:
actor_thread = threading.Thread(target=actor,
args=(actor_thread_stop_event, ))
actor_thread.start()
rospy.spin()
except KeyboardInterrupt:
print("Shutting down")
finally:
actor_thread_stop_event.set()
actor_thread.join()
def main(argv=None):
# Input placeholder
input_img = tf.placeholder(tf.float32, [1, FLAGS.height, FLAGS.width, 3])
# Create network.
# net = DeepLabResNetModel({'data': input_img}, is_training=False, num_classes=NUM_CLASSES)
net = ICNet({'data': input_img}, num_classes=NUM_CLASSES)
# Which variables to load.
restore_var = tf.global_variables()
# Predictions.
# raw_output = net.layers['fc_out']
raw_output = net.layers['conv6_cls']
raw_output_up = tf.image.resize_bilinear(raw_output,
tf.shape(input_img)[1:3, ])
raw_output_up = tf.argmax(raw_output_up, dimension=3)
# Color transform
color_mat = label_colours[..., [2, 1, 0]]
color_mat = tf.constant(color_mat, dtype=tf.float32)
onehot_output = tf.one_hot(raw_output_up, depth=len(label_colours))
onehot_output = tf.reshape(onehot_output, (-1, len(label_colours)))
pred = tf.matmul(onehot_output, color_mat)
pred = tf.reshape(pred, (1, FLAGS.height, FLAGS.width, 3))
# Set up TF session and initialize variables.
config = tf.ConfigProto()
config.gpu_options.allow_growth = True
sess = tf.Session(config=config)
init = tf.global_variables_initializer()
sess.run(init)
# Load weights.
net.load(RESTORE_PATH, sess)
#p = predictor.getSingleFramePredictor(load_path=FLAGS.path, transform=False)
p = predictor.getPredictor(load_path=FLAGS.path, transform=False)
scale = 1.0
#ros_rmp_motion_publisher = RosKobukiMotionPublisher()
ros_rmp_motion_publisher = RosRMPMotionPublisher()
def actor(stop_event):
global image_queue
while len(image_queue) == 0:
print('Wait for image queue filled ... (sleep 1)')
time.sleep(1)
pass
timestep = 0
while not stop_event.is_set():
start = time.time()
img = image_queue[0]
raw_img = img
#raw_img = cv2.resize(img, (FLAGS.width, FLAGS.height)).astype(float)
img = raw_img - IMG_MEAN
img = np.expand_dims(img, axis=0)
preds = sess.run(pred, feed_dict={input_img: img})
if FLAGS.use_seg:
s = preds[0].astype(np.uint8)
else:
s = raw_img
msk_bgr = cv2.resize(s, (84, 84))
# Change BGR -> RGB
msk = cv2.cvtColor(msk_bgr, cv2.COLOR_BGR2RGB)
act = p(msk)
ros_rmp_motion_publisher.publish(act)
end = time.time()
print('Timestep = %d: Inference time = %f' %
(timestep, end - start))
if timestep < 50:
cv2.imwrite('raw_img_%05d_%d.png' % (timestep, act), raw_img)
cv2.imwrite('msk_img_%05d_%d.png' % (timestep, act), msk_bgr)
timestep += 1
def image_subscribe_callback(img):
global image_queue
image_queue.append(img)
ros_camera_image_subscriber = RosImageSubscriber(
user_callback=image_subscribe_callback)
rospy.init_node('agent', anonymous=True)
actor_thread_stop_event = threading.Event()
try:
actor_thread = threading.Thread(target=actor,
args=(actor_thread_stop_event, ))
actor_thread.start()
rospy.spin()
except KeyboardInterrupt:
print("Shutting down")
finally:
actor_thread_stop_event.set()
actor_thread.join()