def image_demo():
"""Detect image."""
assert FLAGS.demo_net == 'squeezeDet' or FLAGS.demo_net == 'squeezeDet+', \
'Selected nueral net architecture not supported: {}'.format(FLAGS.demo_net)
with tf.Graph().as_default():
# Load model
if FLAGS.demo_net == 'squeezeDet':
mc = kitti_squeezeDet_config()
mc.BATCH_SIZE = 1
# model parameters will be restored from checkpoint
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDet(mc, FLAGS.gpu)
elif FLAGS.demo_net == 'squeezeDet+':
mc = kitti_squeezeDetPlus_config()
mc.BATCH_SIZE = 1
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDetPlus(mc, FLAGS.gpu)
saver = tf.train.Saver(model.model_params)
with tf.Session(config=tf.ConfigProto(
allow_soft_placement=True)) as sess:
saver.restore(sess, FLAGS.checkpoint)
for f in glob.iglob(FLAGS.input_path):
im = cv2.imread(f)
im = im.astype(np.float32, copy=False)
im = cv2.resize(im, (mc.IMAGE_WIDTH, mc.IMAGE_HEIGHT))
input_image = im - mc.BGR_MEANS
# Detect
det_boxes, det_probs, det_class = sess.run(
[model.det_boxes, model.det_probs, model.det_class],
feed_dict={model.image_input: [input_image]})
# Filter
final_boxes, final_probs, final_class = model.filter_prediction(
det_boxes[0], det_probs[0], det_class[0])
keep_idx = [idx for idx in range(len(final_probs)) \
if final_probs[idx] > mc.PLOT_PROB_THRESH]
final_boxes = [final_boxes[idx] for idx in keep_idx]
final_probs = [final_probs[idx] for idx in keep_idx]
final_class = [final_class[idx] for idx in keep_idx]
# Draw boxes
_draw_box(
im, final_boxes,
[mc.CLASS_NAMES[idx]+': (%.2f)'% prob \
for idx, prob in zip(final_class, final_probs)],
cdict=mc.cls2clr
)
file_name = os.path.split(f)[1]
out_file_name = os.path.join(FLAGS.out_dir, 'out_' + file_name)
cv2.imwrite(out_file_name, im)
print(
'Image detection output saved to {}'.format(out_file_name))
def image_demo():
"""Detect image."""
with tf.Graph().as_default():
# Load model
mc = kitti_squeezeDet_config()
mc.BATCH_SIZE = 1
# model parameters will be restored from checkpoint
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDet(mc, FLAGS.gpu)
saver = tf.train.Saver(model.model_params)
with tf.Session(config=tf.ConfigProto(
allow_soft_placement=True)) as sess:
saver.restore(sess, FLAGS.checkpoint)
for f in glob.iglob(FLAGS.input_path):
im = cv2.imread(f)
im = im.astype(np.float32, copy=False)
im = cv2.resize(im, (mc.IMAGE_WIDTH, mc.IMAGE_HEIGHT))
input_image = im - mc.BGR_MEANS
# Detect
det_boxes, det_probs, det_class = sess.run(
[model.det_boxes, model.det_probs, model.det_class],
feed_dict={model.image_input: [input_image]})
# Filter
final_boxes, final_probs, final_class = model.filter_prediction(
det_boxes[0], det_probs[0], det_class[0])
keep_idx = [idx for idx in range(len(final_probs)) \
if final_probs[idx] > mc.PLOT_PROB_THRESH]
final_boxes = [final_boxes[idx] for idx in keep_idx]
final_probs = [final_probs[idx] for idx in keep_idx]
final_class = [final_class[idx] for idx in keep_idx]
# TODO(bichen): move this color dict to configuration file
cls2clr = {
'car': (255, 191, 0),
'cyclist': (0, 191, 255),
'pedestrian': (255, 0, 191)
}
# Draw boxes
_draw_box(
im, final_boxes,
[mc.CLASS_NAMES[idx]+': (%.2f)'% prob \
for idx, prob in zip(final_class, final_probs)],
cdict=cls2clr,
)
file_name = os.path.split(f)[1]
out_file_name = os.path.join(FLAGS.out_dir, 'out_' + file_name)
cv2.imwrite(out_file_name, im)
print(
'Image detection output saved to {}'.format(out_file_name))
def image_demo():
assert FLAGS.demo_net == 'squeezeDet' or FLAGS.demo_net == 'squeezeDet+',\
'Selected nueral net architecture not supported: {}'.format(FLAGS.demo_net)
with tf.Graph().as_default():
#Load model
if FLAGS.demo_net == 'squeezeDet':
m = kitti_squeezeDet_config()
m.BATCH_SIZE = 1
m.LOAD_PRETRAINED_MODEL = False
model = SqueezeDet(m, FLAGS.gpu)
elif FLAGS.demo_net == 'squeezeDet+':
m = kitti_squeezeDetPlus_config()
m.BATCH_SIZE = 1
m.LOAD_PRETRAINED_MODEL = False
saver = tf.train.Saver(model.model_params)
with tf.Session(config=tf.ConfigProto(
allow_soft_placement=True)) as sess:
saver.restore(sess, FLAGS.checkpoint)
for f in glob.iglob(FLAGS.input_path):
print f
im = cv2.imread(f)
im = im.astype(np.float32, copy=False)
im = cv2.resize(im, (m.IMAGE_WIDTH, m.IMAGE_HEIGHT))
input_image = im - m.BGR_MEANS
time_start = time.time()
det_boxes, det_probs, det_class = sess.run(
[model.det_boxes, model.det_probs, model.det_class],
feed_dict={model.image_input: [input_image]})
print 'time run {} s'.format(time.time() - time_start)
print 'boxes shape', det_boxes.shape, 'probs shape', det_probs.shape, 'class shape', det_class.shape
final_boxes, final_probs, final_class = model.filter_prediction(
det_boxes[0], det_probs[0], det_class[0])
keep_idx = [idx for idx in range(len(final_probs)) \
if final_probs[idx] > m.PLOT_PROB_THRESH]
final_boxes = [final_boxes[idx] for idx in keep_idx]
final_probs = [final_probs[idx] for idx in keep_idx]
final_class = [final_class[idx] for indx in keep_idx]
print final_probs
print final_class
cls2clr = {
'car': (255, 191, 0),
'cyclist': (0, 191, 255),
'pedestrain': (255, 0, 191)
}
_draw_box(im,final_boxes,
[m.CLASS_NAMES[idx]+': (% 2f)'% prob \
for idx, prob in zip(final_class,final_probs)],
cdict = cls2clr,
)
file_name = os.path.split(f)[1]
out_file_name = os.path.join(FLAGS.out_dir, 'out_' + file_name)
cv2.imwrite(out_file_name, im)
def classify(im_path, conf, prob_thresh):
(sess, mc, model) = conf
im = cv2.imread(im_path)
im = im.astype(np.float32, copy=False)
im = cv2.resize(im, (mc.IMAGE_WIDTH, mc.IMAGE_HEIGHT))
input_image = im - mc.BGR_MEANS
# Detect
det_boxes, det_probs, det_class = sess.run(
[model.det_boxes, model.det_probs, model.det_class],
feed_dict={model.image_input: [input_image]})
# Filter
final_boxes, final_probs, final_class = model.filter_prediction(
det_boxes[0], det_probs[0], det_class[0])
#keep_idx = [idx for idx in range(len(final_probs)) \
# if final_probs[idx] > mc.PLOT_PROB_THRESH]
keep_idx = [idx for idx in range(len(final_probs)) \
if final_probs[idx] > prob_thresh]
final_boxes = [final_boxes[idx] for idx in keep_idx]
final_probs = [final_probs[idx] for idx in keep_idx]
final_class = [final_class[idx] for idx in keep_idx]
# Extract labels + confidence values
res = []
for label, confidence, box in zip(final_class, final_probs, final_boxes):
res.append((label, confidence, box))
cls2clr = {
'car': (255, 191, 0),
'cyclist': (0, 191, 255),
'pedestrian': (255, 0, 191)
}
# Draw boxes
_draw_box(
im, final_boxes,
[mc.CLASS_NAMES[idx]+': (%.2f)'% prob \
for idx, prob in zip(final_class, final_probs)],
cdict=cls2clr,
)
# out_file_name = os.path.join('./', 'out.png')
# cv2.imwrite(out_file_name, im)
# print ('Image detection output saved to {}'.format(out_file_name))
return res
def write_output(index, im, det_boxes, det_probs, det_class):
# Filter
final_boxes, final_probs, final_class = MODEL.filter_prediction(
det_boxes, det_probs, det_class)
keep_idx = [idx for idx in range(len(final_probs)) \
if final_probs[idx] > MODEL_CONFIG.PLOT_PROB_THRESH]
final_boxes = [final_boxes[idx] for idx in keep_idx]
final_probs = [final_probs[idx] for idx in keep_idx]
final_class = [final_class[idx] for idx in keep_idx]
# Draw boxes
_draw_box(
im, final_boxes,
[MODEL_CONFIG.CLASS_NAMES[idx]+': (%.2f)'% prob \
for idx, prob in zip(final_class, final_probs)],
cdict=COLORS
)
cv2.imwrite('output_%d.png' % index, im)
def make_detection_video(self, bag_file, demo_net):
cls2clr = {
'car': (255, 191, 0),
'cyclist': (0, 191, 255),
'pedestrian':(255, 0, 191)
}
video_maker = video.VideoMaker(FLAGS.out_dir)
generator = generate_detections(bag_file, demo_net = demo_net, skip_null = True)
for im, boxes, probs, classes, lidar in generator:
train._draw_box(
im, boxes,
[self.mc.CLASS_NAMES[idx]+': (%.2f)'% prob \
for idx, prob in zip(classes, probs)],
cdict=cls2clr,
)
video_maker.add_image(im)
video_filename = get_filename(bag_file) + '.mp4'
video_maker.make_video(video_filename)
def video_demo():
"""Detect videos."""
cap = cv2.VideoCapture(FLAGS.input_path)
fps = cap.get(cv2.CAP_PROP_FPS)
width = cap.get(cv2.CAP_PROP_FRAME_WIDTH) # float
height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT) # float
out_name = FLAGS.input_path.split('/')[-1:][0]
out_name = out_name.split('.')[0]
#out_name = os.path.join(FLAGS.out_dir, 'det_'+out_name+'.avi')
out_name = os.path.join(FLAGS.out_dir, 'det_' + out_name + '.x264')
#print(out_name)
#fourcc = cv2.VideoWriter_fourcc(*'XVID')
fourcc = cv2.VideoWriter_fourcc(*'X264')
video = cv2.VideoWriter(out_name, fourcc, fps, (int(width), int(height)),
True)
#cap = skvideo.io.VideoCapture(FLAGS.input_path)
# Define the codec and create VideoWriter object
# fourcc = cv2.cv.CV_FOURCC(*'XVID')
# fourcc = cv2.cv.CV_FOURCC(*'MJPG')
# in_file_name = os.path.split(FLAGS.input_path)[1]
# out_file_name = os.path.join(FLAGS.out_dir, 'out_'+in_file_name)
# out = cv2.VideoWriter(out_file_name, fourcc, 30.0, (375,1242), True)
# out = VideoWriter(out_file_name, frameSize=(1242, 375))
# out.open()
assert FLAGS.demo_net == 'squeezeDet' or FLAGS.demo_net == 'tinyYolo' or FLAGS.demo_net == 'tinyYolo', \
'Selected nueral net architecture not supported: {}'.format(FLAGS.demo_net)
with tf.Graph().as_default():
# Load model
if FLAGS.demo_net == 'squeezeDet':
mc = kitti_squeezeDet_config()
mc.BATCH_SIZE = 1
mc.IS_TRAINING = False
# model parameters will be restored from checkpoint
mc.LOAD_PRETRAINED_MODEL = False
#model = SqueezeDet(mc, FLAGS.gpu)
model = SqueezeDet(mc, gpu_id=1)
elif FLAGS.demo_net == 'squeezeDet+':
mc = kitti_squeezeDetPlus_config()
mc.BATCH_SIZE = 1
mc.LOAD_PRETRAINED_MODEL = False
mc.IS_TRAINING = False
model = SqueezeDetPlus(mc, FLAGS.gpu)
elif FLAGS.demo_net == 'tinyYolo':
mc = kitti_tinyYolo_config()
mc.BATCH_SIZE = 1
mc.LOAD_PRETRAINED_MODEL = False
mc.IS_TRAINING = False
model = TinyYolov3(mc, FLAGS.gpu)
saver = tf.train.Saver(model.model_params)
with tf.Session(config=tf.ConfigProto(
allow_soft_placement=True)) as sess:
saver.restore(sess, FLAGS.checkpoint)
# ==============================================================================
# Store Graph
# ==============================================================================
tf.train.write_graph(sess.graph_def,
"/tmp/tensorflow",
"test.pb",
as_text=False)
tf.train.write_graph(sess.graph_def,
"/tmp/tensorflow",
"test.pbtx",
as_text=True)
# ==============================================================================
print("Graph store is done")
times = {}
count = 0
det_last = [0, 0, 0]
top_last = [0, 0, 0]
while cap.isOpened():
t_start = time.time()
count += 1
# Load images from video and crop
ret, frame = cap.read()
print(ret)
if ret:
frame = frame[:, :, ::-1] # <--- convert to BGR
orig_h, orig_w, _ = [float(v) for v in frame.shape]
#print(orig_h, orig_w)
y_start = int(orig_h / 2 - mc.IMAGE_HEIGHT * 2 / 2)
x_start = int(orig_w - mc.IMAGE_WIDTH * 2)
im = frame[y_start:y_start + mc.IMAGE_HEIGHT * 2,
x_start:x_start + mc.IMAGE_WIDTH * 2]
#im = im.astype(np.float32)
#im = frame.astype(np.float32)
#im = im[y_start:y_start+mc.IMAGE_HEIGHT, x_start:x_start+mc.IMAGE_WIDTH]
im = cv2.resize(im, (mc.IMAGE_WIDTH, mc.IMAGE_HEIGHT))
im_input = im.astype(
np.float32
) - mc.BGR_MEANS # <---------------------------------------------------------------------!!!!!!
#im = cv2.resize(im, (mc.IMAGE_WIDTH*1, mc.IMAGE_HEIGHT*1))
#im_gray = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY) # gray color instead of RGB
#im_input = im_gray - np.array([[[128]]])
#im_input = im_gray
#im_input = im_input.reshape((mc.IMAGE_HEIGHT, mc.IMAGE_WIDTH, 1))
else:
print('Done')
break
t_reshape = time.time()
times['reshape'] = t_reshape - t_start
# Detect
det_boxes, det_probs, det_class = sess.run(
[model.det_boxes, model.det_probs, model.det_class],
feed_dict={model.image_input: [im_input]})
t_detect = time.time()
times['detect'] = t_detect - t_reshape
# Extract class only - mine :)
top_idx = det_probs[0].argsort()[:-2:-1] # top probability only
#print("top_idx=", top_idx)
top_prob = det_probs[0][top_idx]
#print("top_prob=", top_prob)
top_class = det_class[0][top_idx]
#print('top class=', top_class)
if (top_prob > mc.PLOT_PROB_THRESH):
new_top_last = [top_last[1], top_last[2], 1]
else:
new_top_last = [top_last[1], top_last[2], 0]
# End of mine
# Filter
final_boxes, final_probs, final_class = model.filter_prediction(
det_boxes[0], det_probs[0], det_class[0])
keep_idx = [idx for idx in range(len(final_probs)) \
if final_probs[idx] > mc.PLOT_PROB_THRESH]
frame = frame[:, :, ::-1]
#im_show_i = im[:,:,::-1] # convert back to RGB
#im_show = im_show_i.astype(np.uint8).copy() # to solve known bug of cv2
im_show = frame[y_start:y_start + mc.IMAGE_HEIGHT * 2,
x_start:x_start + mc.IMAGE_WIDTH * 2]
if (len(keep_idx) != 0):
final_boxes = [final_boxes[idx] for idx in keep_idx]
final_probs = [final_probs[idx] for idx in keep_idx]
final_class = [final_class[idx] for idx in keep_idx]
t_filter = time.time()
times['filter'] = t_filter - t_detect
# Draw boxes
# TODO(bichen): move this color dict to configuration file
cls2clr = {
'car': (255, 191, 0),
'cyclist': (0, 191, 255),
'pedestrian': (255, 0, 191)
}
print(final_boxes)
if (sum(det_last) != 0): # filter
#if(True): # filter
_draw_box(
im_show, final_boxes,
[mc.CLASS_NAMES[idx]+': (%.2f)'% prob \
for idx, prob in zip(final_class, final_probs)], scale=2
)
#t_draw = time.time()
#times['draw']= t_draw - t_filter
#im_show_exp = cv2.resize(im_show, (mc.IMAGE_WIDTH*2, mc.IMAGE_HEIGHT*2))
im_show_exp = im_show
frame[y_start:y_start + mc.IMAGE_HEIGHT * 2,
x_start:x_start + mc.IMAGE_WIDTH * 2] = im_show_exp
cv2.rectangle(
frame, (x_start, y_start),
(x_start + mc.IMAGE_WIDTH * 2, y_start + mc.IMAGE_HEIGHT * 2),
(255, 0, 255), 4)
if (top_prob > mc.PLOT_PROB_THRESH and sum(top_last) != 0):
font = cv2.FONT_HERSHEY_SIMPLEX
print('top_class=', top_class[0])
label = mc.CLASS_NAMES[top_class[0]] #+': (%.2f)'% top_prob[0]
label = label[-2:]
cv2.putText(frame, label, (x_start, y_start), font, 1.5,
(0, 255, 0), 2)
#cv2.imshow('video', im_show) # <--- RGB input
cv2.imshow('video', frame) # <--- RGB input
video.write(frame)
if (len(keep_idx) != 0 and sum(det_last) != 0):
#if(len(keep_idx) !=0):
for x in range(10): # slow down in demo video
video.write(frame)
#cv2.imwrite(out_im_name, im_show) # <--- BGR input
#times['total']= time.time() - t_start
# time_str = ''
# for t in times:
# time_str += '{} time: {:.4f} '.format(t[0], t[1])
# time_str += '\n'
#time_str = 'Total time: {:.4f}, detection time: {:.4f}, filter time: '\
# '{:.4f}'. \
# format(times['total'], times['detect'], times['filter'])
#print (time_str)
#if((len(keep_idx) != 0 and det_last != 0) or True):
# cv2.waitKey()
if cv2.waitKey(5) & 0xFF == ord('q'):
break
new_det_last = [det_last[1], det_last[2], len(keep_idx)]
det_last = new_det_last
top_last = new_top_last
#print(det_last)
# Release everything if job is finished
cap.release()
video.release()
# out.release()
cv2.destroyAllWindows()
def image_demo():
"""Detect image."""
assert FLAGS.demo_net == 'squeezeDet' or FLAGS.demo_net == 'tinyYolo', \
'Selected nueral net architecture not supported: {}'.format(FLAGS.demo_net)
with tf.Graph().as_default():
# Load model
if FLAGS.demo_net == 'squeezeDet':
mc = kitti_squeezeDet_config()
mc.BATCH_SIZE = 1
# model parameters will be restored from checkpoint
mc.LOAD_PRETRAINED_MODEL = False
mc.IS_TRAINING = False
#model = SqueezeDet(mc, FLAGS.gpu)
model = SqueezeDet(mc, gpu_id=0)
elif FLAGS.demo_net == 'tinyYolo':
mc = kitti_tinyYolo_config()
mc.BATCH_SIZE = 1
# model parameters will be restored from checkpoint
mc.LOAD_PRETRAINED_MODEL = False
mc.IS_TRAINING = False
#model = SqueezeDet(mc, FLAGS.gpu)
model = TinyYolov3(mc, gpu_id=0)
saver = tf.train.Saver(model.model_params)
with tf.Session(config=tf.ConfigProto(
allow_soft_placement=True)) as sess:
saver.restore(sess, FLAGS.checkpoint)
if True:
# ==============================================================================
# Store Graph
# ==============================================================================
tf.train.write_graph(sess.graph_def,
"./logs/tensorflow",
"test.pb",
as_text=False)
tf.train.write_graph(sess.graph_def,
"./logs/tensorflow",
"test.pbtxt",
as_text=True)
# ==============================================================================
print('graph was written in pbtxt udner logs/tensorflow')
print(FLAGS.input_path)
for f in glob.iglob(FLAGS.input_path):
print('file name:' + f)
im = cv2.imread(f) # <---------------------------- BGR format
im = im.astype(np.float32, copy=False)
#im = cv2.resize(im, (mc.IMAGE_WIDTH, mc.IMAGE_HEIGHT), interpolation=cv2.INTER_AREA)
orig_h, orig_w, _ = [float(v) for v in im.shape]
org_im = im.copy()
im -= mc.BGR_MEANS # <---------------------------------------------------------------------!!!!!!
im /= 128.0
np.set_printoptions(threshold=np.inf) #'nan')
# Detect
det_boxes, det_probs, det_class, conv12 = sess.run(
[
model.det_boxes, model.det_probs, model.det_class,
model.preds
],
feed_dict={model.image_input: [im]})
conv12 = np.reshape(conv12, (1, 42, 14, 14))
#print('shape of conv12={}'.format(conv12.shape))
#print('conv12={}'.format(conv12))
#print('det_boxes={}'.format(det_boxes))
#print('det_probs={}'.format(det_probs))
# Filter
final_boxes, final_probs, final_class = model.filter_prediction(
det_boxes[0], det_probs[0], det_class[0])
keep_idx = [idx for idx in range(len(final_probs)) \
if final_probs[idx] > mc.PLOT_PROB_THRESH]
final_boxes = [final_boxes[idx] for idx in keep_idx]
final_probs = [final_probs[idx] for idx in keep_idx]
final_class = [final_class[idx] for idx in keep_idx]
#print('keep_idx={}'.format(keep_idx))
#print('final_boxes={}'.format(final_boxes))
#print('final_probs={}'.format(final_probs))
# TODO(bichen): move this color dict to configuration file
'''
cls2clr = {
'car': (255, 191, 0),
'cyclist': (0, 191, 255),
'pedestrian':(255, 0, 191)
}
'''
# Draw boxes
#print('# of final boxes=', len(keep_idx),final_boxes)
_draw_box(#im_gray, final_boxes,
org_im, final_boxes,
[mc.CLASS_NAMES[idx]+': (%.2f)'% prob \
for idx, prob in zip(final_class, final_probs)] #,
#cdict=cls2clr,
)
file_name = os.path.split(f)[1]
out_file_name = os.path.join(FLAGS.out_dir, 'out_' + file_name)
cv2.imwrite(out_file_name, org_im) # <----- BGR format
#cv2.imwrite(out_file_name, im_gray) # <----- BGR format
print(
'Image detection output saved to {}'.format(out_file_name))
def video_demo():
"""Detect videos."""
cap = cv2.VideoCapture(FLAGS.input_path)
# Define the codec and create VideoWriter object
# fourcc = cv2.cv.CV_FOURCC(*'XVID')
# fourcc = cv2.cv.CV_FOURCC(*'MJPG')
# in_file_name = os.path.split(FLAGS.input_path)[1]
# out_file_name = os.path.join(FLAGS.out_dir, 'out_'+in_file_name)
# out = cv2.VideoWriter(out_file_name, fourcc, 30.0, (375,1242), True)
# out = VideoWriter(out_file_name, frameSize=(1242, 375))
# out.open()
assert FLAGS.demo_net == 'squeezeDet' or FLAGS.demo_net == 'squeezeDet+', \
'Selected nueral net architecture not supported: {}'.format(FLAGS.demo_net)
with tf.Graph().as_default():
# Load model
if FLAGS.demo_net == 'squeezeDet':
mc = kitti_squeezeDet_config()
mc.BATCH_SIZE = 1
# model parameters will be restored from checkpoint
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDet(mc, FLAGS.gpu)
elif FLAGS.demo_net == 'squeezeDet+':
mc = kitti_squeezeDetPlus_config()
mc.BATCH_SIZE = 1
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDetPlus(mc, FLAGS.gpu)
saver = tf.train.Saver(model.model_params)
with tf.Session(config=tf.ConfigProto(
allow_soft_placement=True)) as sess:
saver.restore(sess, FLAGS.checkpoint)
times = {}
count = 0
while cap.isOpened():
t_start = time.time()
count += 1
out_im_name = os.path.join(FLAGS.out_dir,
str(count).zfill(6) + '.jpg')
# Load images from video and crop
ret, frame = cap.read()
if ret == True:
# crop frames
frame = frame[500:-205, 239:-439, :]
im_input = frame.astype(np.float32) - mc.BGR_MEANS
else:
break
t_reshape = time.time()
times['reshape'] = t_reshape - t_start
# Detect
det_boxes, det_probs, det_class = sess.run(
[model.det_boxes, model.det_probs, model.det_class],
feed_dict={model.image_input: [im_input]})
t_detect = time.time()
times['detect'] = t_detect - t_reshape
# Filter
final_boxes, final_probs, final_class = model.filter_prediction(
det_boxes[0], det_probs[0], det_class[0])
keep_idx = [idx for idx in range(len(final_probs)) \
if final_probs[idx] > mc.PLOT_PROB_THRESH]
final_boxes = [final_boxes[idx] for idx in keep_idx]
final_probs = [final_probs[idx] for idx in keep_idx]
final_class = [final_class[idx] for idx in keep_idx]
t_filter = time.time()
times['filter'] = t_filter - t_detect
# Draw boxes
# TODO(bichen): move this color dict to configuration file
cls2clr = {
'car': (255, 191, 0),
'cyclist': (0, 191, 255),
'pedestrian': (255, 0, 191)
}
_draw_box(
frame, final_boxes,
[mc.CLASS_NAMES[idx]+': (%.2f)'% prob \
for idx, prob in zip(final_class, final_probs)],
cdict=cls2clr
)
t_draw = time.time()
times['draw'] = t_draw - t_filter
cv2.imwrite(out_im_name, frame)
# out.write(frame)
times['total'] = time.time() - t_start
# time_str = ''
# for t in times:
# time_str += '{} time: {:.4f} '.format(t[0], t[1])
# time_str += '\n'
time_str = 'Total time: {:.4f}, detection time: {:.4f}, filter time: '\
'{:.4f}'. \
format(times['total'], times['detect'], times['filter'])
print(time_str)
if cv2.waitKey(1) & 0xFF == ord('q'):
break
# Release everything if job is finished
cap.release()
# out.release()
cv2.destroyAllWindows()
def cam_demo():
"""Detect image."""
assert FLAGS.demo_net == 'squeezeDet' or FLAGS.demo_net == 'squeezeDet+', \
'Selected nueral net architecture not supported: {}'.format(FLAGS.demo_net)
with tf.Graph().as_default():
# Load model
if FLAGS.demo_net == 'squeezeDet':
mc = kitti_squeezeDet_config()
mc.BATCH_SIZE = 1
# model parameters will be restored from checkpoint
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDet(mc, FLAGS.gpu)
elif FLAGS.demo_net == 'squeezeDet+':
mc = kitti_squeezeDetPlus_config()
mc.BATCH_SIZE = 1
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDetPlus(mc, FLAGS.gpu)
saver = tf.train.Saver(model.model_params)
with tf.Session(config=tf.ConfigProto(
allow_soft_placement=True)) as sess:
saver.restore(sess, FLAGS.checkpoint)
idx_frame = 1
#for f in sorted(glob.iglob(FLAGS.input_path)):
cam = cv2.VideoCapture(
'../dataset/test1.mp4'
) #'/home/siiva/RUDY/SIIVA/GoalCam/annotations/mi5_toysbball/vid5/original_data/vid5.mp4') #'/home/siiva/RUDY/SIIVA/GoalCam/annotations/toys2/original_data/test1.MOV')
print(cam.read())
while True:
ret_val, img1 = cam.read()
#print(cam.read())
#print(ret_val)
img1 = rotate_bound(img1, 0)
#img1 = cv2.imread(f)
im = img1.astype(np.float32, copy=True)
im = cv2.resize(im, (mc.IMAGE_WIDTH, mc.IMAGE_HEIGHT))
img = cv2.resize(img1, (mc.IMAGE_WIDTH, mc.IMAGE_HEIGHT))
input_image = im - mc.BGR_MEANS
# Detect
det_boxes, det_probs, det_class = sess.run(
[model.det_boxes, model.det_probs, model.det_class],
feed_dict={model.image_input: [input_image]})
# Filter
final_boxes, final_probs, final_class = model.filter_prediction(
det_boxes[0], det_probs[0], det_class[0])
keep_idx = [idx for idx in range(len(final_probs)) \
if final_probs[idx] >mc.PLOT_PROB_THRESH] #mc.PLOT_PROB_THRESH
final_boxes = [final_boxes[idx] for idx in keep_idx]
final_probs = [final_probs[idx] for idx in keep_idx]
final_class = [final_class[idx] for idx in keep_idx]
# TODO(bichen): move this color dict to configuration file
cls2clr = {
'01ball': (255, 191, 0),
'02basket': (0, 191, 255),
'03person': (191, 0, 255),
#'pedestrian':(255, 0, 191)
}
# Draw boxes
event = _draw_box(
img, final_boxes,
[mc.CLASS_NAMES[idx] + ': (%.2f)' % prob \
for idx, prob in zip(final_class, final_probs)],
cdict=cls2clr,
)
# file_name = os.path.split(f)[1]
# out_file_name = os.path.join(FLAGS.out_dir, 'out_'+file_name)
# cv2.imwrite(out_file_name, img)
# print ('Image detection output saved to {}'.format(out_file_name))
# if event:
# cv2.imwrite('/home/siiva/RUDY/SIIVA/GoalCam/annotations/NBA_CBA_CUBA/190918/001/original_data/out/%06d.jpg' % idx_frame, img)
# else:
# cv2.imwrite('/media/rudy/C0E28D29E28D252E/goalCAM/CODE/out/%06d.jpg' % idx_frame, img)
# cv2.imwrite(
# '/home/siiva/RUDY/SIIVA/GoalCam/annotations/2018-09-02/frames_VID_20180622_153554/%06d.jpg' % idx_frame, img1)
idx_frame += 1
cv2.imshow('my webcam', img)
if cv2.waitKey(1) == 27:
break # esc to quit
cv2.destroyAllWindows()
def image_demo():
"""Detect image."""
assert FLAGS.demo_net == 'squeezeDet' or FLAGS.demo_net == 'squeezeDet+', \
'Selected nueral net architecture not supported: {}'.format(FLAGS.demo_net)
with tf.Graph().as_default():
# Load model
if FLAGS.demo_net == 'squeezeDet':
mc = kitti_squeezeDet_config()
mc.BATCH_SIZE = 1
# model parameters will be restored from checkpoint
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDet(mc, FLAGS.gpu)
elif FLAGS.demo_net == 'squeezeDet+':
mc = kitti_squeezeDetPlus_config()
mc.BATCH_SIZE = 1
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDetPlus(mc, FLAGS.gpu)
saver = tf.train.Saver(model.model_params)
with tf.Session(config=tf.ConfigProto(
allow_soft_placement=True)) as sess:
saver = tf.train.import_meta_graph(
'./logs/squeezedet/model.ckpt-802500.meta')
#saver.restore(sess, FLAGS.checkpoint)
saver.restore(sess,
tf.train.latest_checkpoint('./logs/squeezedet/'))
for f in glob.iglob(FLAGS.input_path):
im = cv2.imread(f)
im = im.astype(np.float32, copy=False)
im = cv2.resize(im, (mc.IMAGE_WIDTH, mc.IMAGE_HEIGHT))
input_image = im - mc.BGR_MEANS
# Detect
det_boxes, det_probs, det_class = sess.run(
[model.det_boxes, model.det_probs, model.det_class],
feed_dict={model.image_input: [input_image]})
print('class name:', det_class[0])
print('probablility:', det_probs[0])
# Filter
final_boxes, final_probs, final_class = model.filter_prediction(
det_boxes[0], det_probs[0], det_class[0])
keep_idx = [idx for idx in range(len(final_probs)) \
if final_probs[idx] > mc.PLOT_PROB_THRESH]
final_boxes = [final_boxes[idx] for idx in keep_idx]
final_probs = [final_probs[idx] for idx in keep_idx]
final_class = [final_class[idx] for idx in keep_idx]
# TODO(bichen): move this color dict to configuration file
cls2clr = {
'boat1': (255, 191, 0),
'car6': (0, 191, 255),
'person5': (255, 0, 191),
'group7': (254, 194, 5),
'car2': (0, 190, 250),
'group2': (253, 0, 195),
'boat4': (252, 192, 0),
'car6': (0, 191, 255),
'riding10': (251, 0, 193),
'boat5': (252, 194, 0),
'car23': (0, 198, 255),
'group3': (250, 0, 198),
'boat7': (22, 192, 153),
}
# Draw boxes
_draw_box(
im, final_boxes,
[mc.CLASS_NAMES[idx]+': (%.2f)'% prob \
for idx, prob in zip(final_class, final_probs)],
cdict=cls2clr,
)
file_name = os.path.split(f)[1]
out_file_name = os.path.join(FLAGS.out_dir, 'out_' + file_name)
cv2.imwrite(out_file_name, im)
#cv2.imshow('image',im)
print(
'Image detection output saved to {}'.format(out_file_name))
def video_demo_inference_graph(mask_parameterization_now,
log_anchors_now,
encoding_type_now,
checkpoint_path,
softnms=False):
assert FLAGS.demo_net == 'squeezeDet', 'Selected neural net architecture not supported: {}'.format(
FLAGS.demo_net)
if FLAGS.dataset_inf == 'CITYSCAPE':
CLASS_NAMES = tuple(
sorted(('person', 'rider', 'car', 'truck', 'bus', 'motorcycle',
'bicycle')))
else:
CLASS_NAMES = tuple(sorted(('car', 'pedestrian', 'cyclist')))
# Class specific color definitions
_cdict = {}
_cdict['person'] = (0, 204, 102)
_cdict['pedestrian'] = (0, 204, 102)
_cdict['rider'] = (102, 0, 204)
_cdict['car'] = (0, 000, 204)
_cdict['truck'] = (153, 153, 0)
_cdict['bus'] = (0, 153, 153)
_cdict['motorcycle'] = (204, 0, 102)
_cdict['bicycle'] = (255, 128, 0)
_cdict['cyclist'] = (255, 128, 0)
detection_graph = tf.Graph()
if not '.png' in FLAGS.input_path:
cap = cv2.VideoCapture(FLAGS.input_path)
if (cap.isOpened() == False):
print("Error opening video stream or file")
with detection_graph.as_default():
od_graph_def = tf.GraphDef()
with tf.gfile.GFile(checkpoint_path, 'rb') as fid:
serialized_graph = fid.read()
od_graph_def.ParseFromString(serialized_graph)
tf.import_graph_def(od_graph_def, name='')
with tf.Session() as sess:
# Get handles to input and output tensors
ops = tf.get_default_graph().get_operations()
all_tensor_names = {
output.name
for op in ops for output in op.outputs
}
tensor_dict = {}
for key in ['conv12/bias_add']:
tensor_name = key + ':0'
if tensor_name in all_tensor_names:
tensor_dict[key] = tf.get_default_graph(
).get_tensor_by_name(tensor_name)
image_tensor = tf.get_default_graph().get_tensor_by_name(
'image_input_1:0')
if '.png' in FLAGS.input_path:
print("Processing image sequences!")
for image_path in glob.iglob(FLAGS.input_path):
start = time.time()
read_img = cv2.imread(image_path)
frame_read_time = time.time() - start
image_np_orig = copy.deepcopy(read_img)
# Run inference
final_boxes, final_probs, final_class, final_probs, image_np_orig = process_frame(read_img, mask_parameterization_now, \
log_anchors_now, encoding_type_now, checkpoint_path, sess, tensor_dict, image_tensor, \
softnms)
end = time.time()
fps = 1 / (end - start)
fps_text = "FPS counter: {:0.2f} fps/Per frame read time: {:0.2f} ms".format(
round(fps, 2), round(frame_read_time * 1000, 2))
_draw_box(
image_np_orig, final_boxes,
[CLASS_NAMES[idx]+': (%.2f)'% prob \
for idx, prob in zip(final_class, final_probs)],
(0, 0, 255), draw_masks=(mask_parameterization_now == 8), fill=False, cdict=_cdict, fps_text=fps_text)
cv2.imshow('Frame', image_np_orig)
if FLAGS.write_to_disk:
out_file_name = os.path.join(
FLAGS.out_dir,
image_path.split('\\')[-1][:-4] + ".png")
cv2.imwrite(out_file_name, image_np_orig)
if cv2.waitKey(25) & 0xFF == ord('q'):
break
else:
frame_cnt = 0
print("Processing video!")
while (cap.isOpened()):
start = time.time()
ret, read_img = cap.read()
frame_read_time = time.time() - start
if ret == True:
# Run inference
final_boxes, final_probs, final_class, final_probs, image_np_orig = process_frame(read_img, mask_parameterization_now, \
log_anchors_now, encoding_type_now, checkpoint_path, sess, tensor_dict, image_tensor, \
softnms)
end = time.time()
fps = 1 / (end - start)
fps_text = "FPS counter: {:0.2f} fps/Per frame read time: {:0.2f} ms".format(
round(fps, 2), round(frame_read_time * 1000, 2))
_draw_box(
image_np_orig, final_boxes,
[CLASS_NAMES[idx]+': (%.2f)'% prob \
for idx, prob in zip(final_class, final_probs)],
(0, 0, 255), draw_masks=(mask_parameterization_now == 8), fill=False, cdict=_cdict, fps_text=fps_text)
cv2.imshow('Frame', image_np_orig)
frame_cnt += 1
if FLAGS.write_to_disk:
out_file_name = os.path.join(
FLAGS.out_dir,
FLAGS.input_path.split('\\')[-1][:-4] + "_" +
str(frame_cnt) + ".png")
cv2.imwrite(out_file_name, image_np_orig)
if cv2.waitKey(25) & 0xFF == ord(
'q') or frame_cnt == 1200:
break
else:
break
if not '.png' in FLAGS.input_path:
cap.release()
cv2.destroyAllWindows()
def TsDet_callback(rgb,pointcloud):
global count, sess, model, mc
print ('I here rgb and pointcloud !',count)
count = count + 1
bridge = CvBridge()
try:
cv_image_rgb = bridge.imgmsg_to_cv2(rgb, desired_encoding="bgr8")
#cv_image_depth = bridge.imgmsg_to_cv2(depth, desired_encoding="16UC1")
except CvBridgeError as e:
print(e)
#position_pub.publish(rgb)
#visualize
#cv2.imshow("Image window", cv_image_depth)
#cv2.waitKey(3)
times = {}
t_start = time.time()
im = cv_image_rgb
im = im.astype(np.float32, copy=False)
im = cv2.resize(im, (mc.IMAGE_WIDTH, mc.IMAGE_HEIGHT))
input_image = im - mc.BGR_MEANS
t_reshape = time.time()
times['reshape']= t_reshape - t_start
# Detect
det_boxes, det_probs, det_class = sess.run(
[model.det_boxes, model.det_probs, model.det_class],
feed_dict={model.image_input:[input_image]})
t_detect = time.time()
times['detect']= t_detect - t_reshape
# Filter
final_boxes, final_probs, final_class = model.filter_prediction(
det_boxes[0], det_probs[0], det_class[0])
# keep_idx = [idx for idx in range(len(final_probs)) \
# if final_probs[idx] > mc.PLOT_PROB_THRESH]
# final_boxes = [final_boxes[idx] for idx in keep_idx]
# final_probs = [final_probs[idx] for idx in keep_idx]
# final_class = [final_class[idx] for idx in keep_idx]
keep_idx = np.squeeze(np.argwhere(np.array(final_probs) > mc.PLOT_PROB_THRESH))
final_boxes = np.array(final_boxes)[keep_idx, :]
final_probs = np.array(final_probs)[keep_idx]
final_class = np.array(final_class)[keep_idx]
#print(final_boxes, final_probs, final_class)
#print("keep", len(keep_idx))
#keep_idx_ = keep_idx.tolist()
#print(keep_idx)
if keep_idx.shape == ():
final_boxes = final_boxes[np.newaxis, :]
final_probs = np.array([final_probs])
final_class = np.array([final_class])
#print(final_boxes, final_probs, final_class)
avgX = 0
avgY = 0
avgZ = 0
robo_position = []
print(final_class)
if len(final_boxes)>0 and (np.any(final_class == 0) or np.any(final_class == 2)) and np.any(final_class) == 1:
# pointcloud
#print('find enemy')
area = [bbox[2] * bbox[3] for bbox in final_boxes]
max_area_idx = np.argmax(area)
#print(final_boxes, max_area_idx)
robo_bbox = final_boxes[max_area_idx, :]
# find index of robo (which index != 1)
robo_idx = np.where(final_class != 1)
print('robo_idx:',robo_idx)
cx = robo_bbox[0]
cy = robo_bbox[1]
pointcloud_w = 5
pointcloud_h = 5
if pointcloud_w > robo_bbox[2]:
pointcloud_w = robo_bbox[2]
if pointcloud_h > robo_bbox[3]:
pointcloud_h = robo_bbox[3]
x_ = np.arange(int(cx - pointcloud_w/2), int(cx + pointcloud_w/2), 1)
y_ = np.arange(int(cy - pointcloud_h/2), int(cy + pointcloud_h/2), 1)
roi = [[x, y] for x in x_ for y in y_]
points = list(pc2.read_points(pointcloud, skip_nans=False, field_names = ("x", "y", "z"), uvs=roi))
# points_old = [[p[0],p[1],p[2]] for p in pc2.read_points(pointcloud, skip_nans=False, field_names=('x', 'y', 'z'))]
# points_old = np.array(points_old)
# points_old = np.resize(points_old, (240, 424 , 3))
#print(points)
robo_pointcloud = np.array(points)
avgX = 0
avgY = 0
avgZ = 0
idx = 0
# robo_pointcloud = [if for pointcloud_ in robo_pointcloud]
positionX = robo_pointcloud[:,0].reshape(-1,25).squeeze()
positionY = robo_pointcloud[:,1].reshape(-1,25).squeeze()
positionZ = robo_pointcloud[:,2].reshape(-1,25).squeeze()
#print(positionX)
positionX = positionX[np.logical_not(np.isnan(positionX))]
positionY = positionY[np.logical_not(np.isnan(positionY))]
positionZ = positionZ[np.logical_not(np.isnan(positionZ))]
#print(positionX, positionY, positionZ)
avgX = np.mean(positionX)
avgY = np.mean(positionY)
avgZ = np.mean(positionZ)
robo_position.append([avgX, avgY, avgZ])
print('position:', avgX, avgY, avgZ)
br = tf2_ros.TransformBroadcaster()
t = TransformStamped()
t.header.stamp = rospy.Time.now()
t.header.frame_id = 'base_link'
t.child_frame_id = 'enemy_0'
t.transform.translation.x = avgZ
t.transform.translation.y = -avgX
t.transform.translation.z = 0
t.transform.rotation.x = 0
t.transform.rotation.y = 0
t.transform.rotation.z = 0
t.transform.rotation.w = 1
br.sendTransform(t)
# enemy_position = Odometry()
# enemy_position.pose.pose.position.x = avgZ
# enemy_position.pose.pose.position.y = -avgX
# enemy_position.pose.pose.position.z = avgY
# pub.publish(enemy_position)
robo_position = np.array(robo_position)
enemy_position = ObjectList()
enemy_position.header.stamp = rospy.Time.now()
enemy_position.header.frame_id = 'enemy'
enemy_position.num = robo_position.shape[0]
if robo_position.shape[0] == 1:
robo_position = np.array(robo_position)
red_idx = 0
blue_idx = 0
for object_idx in range(robo_position.shape[0]):
enemy = Object()
enemy.pose.position.x = robo_position[object_idx, 2]
enemy.pose.position.y = -robo_position[object_idx, 0]
enemy.pose.position.z = robo_position[object_idx, 1]
enemy_position.object.append(enemy)
#sending tf information
t.header.stamp = rospy.Time.now()
t.header.frame_id = 'base_link'
print('np.array(robo_idx)', np.array(robo_idx))
if mc.CLASS_NAMES[final_class[np.array(robo_idx)[0][0]]] == 'red':
t.child_frame_id = mc.CLASS_NAMES[np.array(robo_idx)[0][0]] + str(red_idx)
enemy.team.data = mc.CLASS_NAMES[np.array(robo_idx)[0][0]] + str(red_idx)
red_idx = red_idx + 1
if mc.CLASS_NAMES[final_class[np.array(robo_idx)[0][0]]] == 'blue':
t.child_frame_id = mc.CLASS_NAMES[np.array(robo_idx)[0][0]] + str(blue_idx)
enemy.team.data = mc.CLASS_NAMES[np.array(robo_idx)[0][0]] + str(blue_idx)
blue_idx = blue_idx + 1
t.transform.translation.x = robo_position[object_idx, 2]
t.transform.translation.y = -robo_position[object_idx, 0]
t.transform.translation.z = robo_position[object_idx, 1]
t.transform.rotation.x = 0
t.transform.rotation.y = 0
t.transform.rotation.z = 0
t.transform.rotation.w = 1
br.sendTransform(t)
pub.publish(enemy_position)
mc.DRAW_BOX = True
else:
print('No enemy!!!')
mc.DRAW_BOX = True
t_filter = time.time()
times['filter']= t_filter - t_detect
times['total']= time.time() - t_start
time_str = 'Total time: {:.4f}, detection time: {:.4f}, filter time: '\
'{:.4f}'. \
format(times['total'], times['detect'], times['filter'])
print (time_str)
# # TODO(bichen): move this color dict to configuration file
cls2clr = {
'red': (0, 0, 255),
'wheel': (0, 255, 0),
'blue':(255, 0, 0)
}
if mc.DRAW_Video:
video = cv2.VideoWriter('/home/ubuntu/catkin_ws/src/robo_perception/scripts/visual/demo.avi',
cv2.VideoWriter_fourcc(*"MJPG"),
30,
(424, 240))
im = _draw_box(
im,
final_boxes,
[mc.CLASS_NAMES[idx]+': (%.2f)'% prob \
for idx, prob in zip(final_class, final_probs)],
cdict=cls2clr,
)
video.write(im)
# for idx, prob in zip(final_class, final_probs):
# print(mc.CLASS_NAMES[idx])
# print(idx)
# Draw boxes
if mc.DRAW_BOX:
#if mc.DRAW_BOX and len(final_boxes):
im = _draw_box(
im,
final_boxes,
[mc.CLASS_NAMES[idx]+':%.2f'% prob for idx, prob in zip(final_class, final_probs)],
cdict=cls2clr,
)
#
# cv2.imshow('demo', im)
# file_name = os.path.split(f)[1]
position_str = 'x='+str(round(avgX,3)) + ' y='+ str(round(avgY,3)) + ' z=' + str(round(avgZ,3))
font = cv2.FONT_HERSHEY_SIMPLEX
cv2.putText(im, position_str, (10, 20), font, 0.7, (0,255,255), 2)
image_count = count
image_name = image_count % 1000
file_name = str(image_name) + '.jpg'
out_file_name = os.path.join('/home/ubuntu/catkin_ws/src/robo_perception/scripts/visual', 'out_'+file_name)
im = im.astype('uint8')
cv2.imshow('demo', im)
cv2.waitKey(3)
cv2.imwrite(out_file_name, im)
print(im.dtype, im[0,0,0])
Im_to_ros = Image()
try:
Im_to_ros = bridge.cv2_to_imgmsg(im, "bgr8")
except CvBridgeError as e:
print (e)
Im_to_ros.header.stamp = rospy.Time.now()
Im_to_ros.header.frame_id = 'camera_link'
pub_dr.publish(Im_to_ros)
print ('Image detection output saved to {}'.format(out_file_name))
def video_demo():
"""Detect videos."""
cap = cv2.VideoCapture(
r'/home/pi/Downloads/road.mp4') #r'/home/pi/Downloads/road.mp4')
fps = int(cap.get(cv2.CAP_PROP_FPS))
# Define the codec and create VideoWriter object
# fourcc = cv2.cv.CV_FOURCC(*'XVID')
# fourcc = cv2.cv.CV_FOURCC(*'MJPG')
# in_file_name = os.path.split(FLAGS.input_path)[1]
# out_file_name = os.path.join(FLAGS.out_dir, 'out_'+in_file_name)
# out = cv2.VideoWriter(out_file_name, fourcc, 30.0, (375,1242), True)
# out = VideoWriter(out_file_name, frameSize=(1242, 375))
# out.open()
assert FLAGS.demo_net == 'squeezeDet' or FLAGS.demo_net == 'squeezeDet+', \
'Selected nueral net architecture not supported: {}'.format(FLAGS.demo_net)
with tf.Graph().as_default():
# Load model
if FLAGS.demo_net == 'squeezeDet':
mc = kitti_squeezeDet_config()
mc.BATCH_SIZE = 1
# model parameters will be restored from checkpoint
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDet(mc, FLAGS.gpu)
elif FLAGS.demo_net == 'squeezeDet+':
mc = kitti_squeezeDetPlus_config()
mc.BATCH_SIZE = 1
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDetPlus(mc, FLAGS.gpu)
saver = tf.train.Saver(model.model_params)
with tf.Session(config=tf.ConfigProto(
allow_soft_placement=True)) as sess:
saver.restore(sess, FLAGS.checkpoint)
times = {}
count = 0
while cap.isOpened():
t_start = time.time()
count += 1
out_im_name = os.path.join(FLAGS.out_dir,
str(count).zfill(6) + '.jpg')
# Load images from video and crop
ret, frame = cap.read()
if ret == True:
# crop frames
#frame = frame[500:-205, 239:-439, :]
frame = cv2.resize(frame, (1248, 384))
im_input = frame.astype(np.float32) - mc.BGR_MEANS
else:
break
t_reshape = time.time()
times['reshape'] = t_reshape - t_start
start = time.time()
# Detect
det_boxes, det_probs, det_class = sess.run(
[model.det_boxes, model.det_probs, model.det_class],
feed_dict={model.image_input: [im_input]})
t_detect = time.time()
times['detect'] = t_detect - t_reshape
# Filter
final_boxes, final_probs, final_class = model.filter_prediction(
det_boxes[0], det_probs[0], det_class[0])
keep_idx = [idx for idx in range(len(final_probs)) \
if final_probs[idx] > mc.PLOT_PROB_THRESH]
final_boxes = [final_boxes[idx] for idx in keep_idx]
final_probs = [final_probs[idx] for idx in keep_idx]
final_class = [final_class[idx] for idx in keep_idx]
end = time.time()
t_filter = time.time()
times['filter'] = t_filter - t_detect
# Draw boxes
# TODO(bichen): move this color dict to configuration file
cls2clr = {
'car': (255, 191, 0),
'cyclist': (0, 191, 255),
'pedestrian': (255, 0, 191)
}
_draw_box(
frame, final_boxes,
[mc.CLASS_NAMES[idx]+': (%.2f)'% prob \
for idx, prob in zip(final_class, final_probs)],
cdict=cls2clr
)
for i, b in enumerate(det_boxes[0]):
if det_class[0][i] == 0 or det_class[0][
i] == 1 or det_class[0][i] == 2:
if det_probs[0][i] >= 0.5:
mid_x = (det_boxes[0][i][1] +
det_boxes[0][i][3]) / 2
mid_y = (det_boxes[0][i][0] +
det_boxes[0][i][2]) / 2
apx_distance = round(
(1 -
(det_boxes[0][i][3] - det_boxes[0][i][1])), 1)
cv2.putText(frame, '{}'.format(apx_distance),
(int(mid_x * 800), int(mid_y * 450)),
cv2.FONT_HERSHEY_SIMPLEX, 0.7,
(255, 255, 255), 2)
if apx_distance <= 200 and apx_distance > 180:
#if mid_x > 130 and mid_x < 140:
pwm_motor.ChangeDutyCycle(25)
print("object distance:" + str(apx_distance) +
"motor is at 25% Duty Cycle")
time.sleep(0.5)
if apx_distance <= 180 and apx_distance > 160:
#if mid_x > 140 and mid_x < 150:
pwm_motor.ChangeDutyCycle(50)
print("object distance:" + str(apx_distance) +
"motor is at 50% Duty Cycle")
time.sleep(0.5)
if apx_distance <= 160 and apx_distance > 140:
#if mid_x > 176 and mid_x < 184:
pwm_motor.ChangeDutyCycle(75)
print("object distance:" + str(apx_distance) +
"motor is at 75% Duty Cycle")
time.sleep(0.5)
if apx_distance <= 140 and apx_distance >= 120:
#if mid_x > 168 and mid_x < 176:
pwm_motor.ChangeDutyCycle(100)
print("object distance:" + str(apx_distance) +
"motor is at 100% Duty Cycle")
cv2.putText(frame, 'WARNING!!!', (50, 50),
cv2.FONT_HERSHEY_SIMPLEX, 1.0,
(0, 0, 255), 3)
total_time_ms = (end - start) * 1000.
#print ('processing time:%.3fms'%(total_time_ms))
#frame /= 255.
cv2.putText(frame, 'fps:%.2f' % (1000 / total_time_ms),
(5, 20), cv2.FONT_HERSHEY_PLAIN, 1.3, (0, 255, 0),
2)
t_draw = time.time()
times['draw'] = t_draw - t_filter
#cv2.imwrite(out_im_name, frame)
# out.write(frame)
cv2.imshow('Video Detection', frame)
#times['total']= time.time() - t_start
# time_str = ''
# for t in times:
# time_str += '{} time: {:.4f} '.format(t[0], t[1])
# time_str += '\n'
'''time_str = 'Total time: {:.4f}, detection time: {:.4f}, filter time: '\
'{:.4f}'. \
format(times['total'], times['detect'], times['filter'])'''
if cv2.waitKey(1) & 0xFF == ord('q'):
break
# Release everything if job is finished
cap.release()
# out.release()
cv2.destroyAllWindows()
def image_demo():
"""Detect image."""
assert FLAGS.demo_net == 'squeezeDet' or FLAGS.demo_net == 'squeezeDet+', \
'Selected nueral net architecture not supported: {}'.format(FLAGS.demo_net)
with tf.Graph().as_default():
# Load model
if FLAGS.demo_net == 'squeezeDet':
mc = kitti_squeezeDet_config()
mc.BATCH_SIZE = 1
# model parameters will be restored from checkpoint
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDet(mc, FLAGS.gpu)
elif FLAGS.demo_net == 'squeezeDet+':
mc = kitti_squeezeDetPlus_config()
mc.BATCH_SIZE = 1
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDetPlus(mc, FLAGS.gpu)
saver = tf.train.Saver(model.model_params)
with tf.Session(config=tf.ConfigProto(
allow_soft_placement=True)) as sess:
saver.restore(sess, FLAGS.checkpoint)
cnt = 0
FPS = 0
FPS10 = 0
for f in glob.iglob(FLAGS.input_path):
# im = cv2.imread(f)
# im = im.astype(np.float32, copy=False)
# im = cv2.resize(im, (mc.IMAGE_WIDTH, mc.IMAGE_HEIGHT))
im = Image.open(f)
im = im.resize((mc.IMAGE_WIDTH, mc.IMAGE_HEIGHT),
Image.ANTIALIAS)
input_image = im - mc.BGR_MEANS
start = time.time()
# Detect
det_boxes, det_probs, det_class = sess.run(
[model.det_boxes, model.det_probs, model.det_class],
feed_dict={model.image_input: [input_image]})
cnt = cnt + 1
end = time.time()
FPS = (1 / (end - start))
FPS10 = FPS10 + FPS
# print ("FPS: " + str(FPS))
if cnt % 10 == 0:
print("FPS(mean), detection: " + str(FPS10 / 10))
FPS10 = 0
# Filter
final_boxes, final_probs, final_class = model.filter_prediction(
det_boxes[0], det_probs[0], det_class[0])
keep_idx = [idx for idx in range(len(final_probs)) \
if final_probs[idx] > mc.PLOT_PROB_THRESH]
final_boxes = [final_boxes[idx] for idx in keep_idx]
final_probs = [final_probs[idx] for idx in keep_idx]
final_class = [final_class[idx] for idx in keep_idx]
# TODO(bichen): move this color dict to configuration file
cls2clr = {
'car': (255, 191, 0),
'cyclist': (0, 191, 255),
'pedestrian': (255, 0, 191)
}
# Draw boxes
_draw_box(
im, final_boxes,
[mc.CLASS_NAMES[idx]+': (%.2f)'% prob \
for idx, prob in zip(final_class, final_probs)],
cdict=cls2clr,
)
file_name = os.path.split(f)[1]
out_file_name = os.path.join(FLAGS.out_dir, 'out_' + file_name)
# cv2.imwrite(out_file_name, im)
im.save(out_file_name)
print(
'Image detection output saved to {}'.format(out_file_name))
def callback(self, data):
self.cnt = self.cnt + 1
start_2 = time.time()
start_3 = time.time()
# load : 2*0.00136515457666 sec
try:
input_raw_image = self.bridge.imgmsg_to_cv2(data, "bgr8")
except CvBridgeError as e:
print(e)
(rows, cols, channels) = input_raw_image.shape
# if cols > 60 and rows > 60:
# cv2.circle(cv_image, (50, 50), 10, 255)
# demo.py
# im = cv_image
# print(input_raw_image.shape)
input_image = input_raw_image[int((rows - mc.IMAGE_HEIGHT) / 2) -
200:-200 +
int((rows - mc.IMAGE_HEIGHT) / 2) +
mc.IMAGE_HEIGHT,
int((cols - mc.IMAGE_WIDTH) /
2):int((cols - mc.IMAGE_WIDTH) / 2) +
mc.IMAGE_WIDTH, :]
# print(int((rows-mc.IMAGE_HEIGHT)/2)-200,int((cols-mc.IMAGE_WIDTH)/2))
end_3 = time.time()
sec_3 = (end_3 - start_3)
# print("f:crop image")
# print(sec_3)
self.sec_crop = self.sec_crop + sec_3
# print ("FPS: " + str(FPS))
if self.cnt % 10 == 0:
print("sec(crop image): " + str(self.sec_crop / 10))
self.sec_crop = 0
# print(type(im))
# print(type(self.BGR_MEANS))
# 2*9.16483911381e-04 sec
# im = im.astype(np.float32, copy=False)
# print (im.dtype)
# print(self.BGR_MEANS.dtype)
# print("-------")
# end_3 = time.time()
# sec_3 = (end_3 - start_3)
# print("g:im.astype")
# print(sec_3)
# start_3 = time.time()
# im = cv2.resize(im, (mc.IMAGE_WIDTH, mc.IMAGE_HEIGHT))
# print(mc.BGR_MEANS.dtype)
# print(mc.BGR_MEANS)
# input_image = im - mc.BGR_MEANS
# input_image = input_image - self.BGR_MEANS
# input_image = im
# print(im.shape)
# print(self.BGR_MEANS.shape)
# print (input_image.dtype)
# print(mc.BGR_MEANS.dtype)
# end_3 = time.time()
# sec_3 = (end_3 - start_3)
# print("e:im - mc.BGR_MEANS")
# print(sec_3)
# start_3 = time.time()
start = time.time()
# Detect
# det_boxes, det_probs, det_class = self._sess.run(
# [model.det_boxes, model.det_probs, model.det_class],
# feed_dict={model.image_input: [input_image]})
# print(model.det_boxes)
# print(model.image_input)
# print(model.ph_image_input)
# test = tf.get_default_graph().get_tensor_by_name("image_input:0")
# print(test) # Tensor("example:0", shape=(2, 2), dtype=float32)
det_boxes, det_probs, det_class = self._sess.run(
[model.det_boxes, model.det_probs, model.det_class],
feed_dict={model.ph_image_input: [input_image]})
# print(model.ph_raw_image_input)
# det_boxes, det_probs, det_class = self._sess.run(
# [model.det_boxes, model.det_probs, model.det_class],
# feed_dict={model.ph_ori_image_input: [input_raw_image]})
end = time.time()
self.FPS = (1 / (end - start))
# print("Detect:'")
# print((end - start))
# print(self.FPS)
self.FPS10 = self.FPS10 + self.FPS
# print ("FPS: " + str(FPS))
if self.cnt % 10 == 0:
print("FPS(detection): " + str(self.FPS10 / 10))
print("sec(detection): " + str(1 / self.FPS10 * 10))
self.FPS10 = 0
# Filter
# 0.208640098572 sec for 512 TOP_N_DETECTION
# 0.0111668109894 sec for 64 TOP_N_DETECTION
# 0.005 for 32
start_3 = time.time()
# # filter
final_boxes, final_probs, final_class = model.filter_prediction(
det_boxes[0], det_probs[0], det_class[0])
end_3 = time.time()
sec_3 = (end_3 - start_3)
# print("3:Filter")
# print(sec_3)
self.sec_nms = self.sec_nms + sec_3
# print ("FPS: " + str(FPS))
if self.cnt % 10 == 0:
print("sec(filter): " + str(self.sec_nms / 10))
self.sec_nms = 0
# 4.50611114502e-05 sec
"""Set plot threshold"""
start_3 = time.time()
keep_idx = [idx for idx in range(len(final_probs)) \
if final_probs[idx] > mc.PLOT_PROB_THRESH]
final_boxes = [final_boxes[idx] for idx in keep_idx]
final_probs = [final_probs[idx] for idx in keep_idx]
final_class = [final_class[idx] for idx in keep_idx]
end_3 = time.time()
sec_3 = (end_3 - start_3)
# print("final_probs[idx] > mc.PLOT_PROB_THRESH")
# print(sec_3)
start_3 = time.time()
# TODO(bichen): move this color dict to configuration file
cls2clr = {
'car': (255, 191, 0),
'cyclist': (0, 191, 255),
'pedestrian': (255, 0, 191)
}
# Draw boxes
# 0.00111293792725 sec
_draw_box(
input_image, final_boxes,
[mc.CLASS_NAMES[idx] + ': (%.2f)' % prob \
for idx, prob in zip(final_class, final_probs)],
cdict=cls2clr,
)
end_3 = time.time()
sec_3 = (end_3 - start_3)
# print("c:_draw_box")
# print(sec_3)
start_3 = time.time()
# file_name = os.path.split(f)[1]
file_name = "test.png"
out_file_name = os.path.join(FLAGS.out_dir, 'out_ros_' + file_name)
cv2.imwrite(out_file_name, input_image)
# print('Image detection output saved to {}'.format(out_file_name))
# show image
# cv2.imshow("Image window", im.astype(np.uint8, copy=False))
# cv2.waitKey(3)
# 2*5.80072402954e-04 sec
try:
# self.image_pub.publish(self.bridge.cv2_to_imgmsg(im, "16SC3"))
pass
except CvBridgeError as e:
print(e)
self.cnt_2 = self.cnt_2 + 1
end_2 = time.time()
self.FPS_2 = (1 / (end_2 - start_2))
self.sec_2 = self.sec_2 + (end_2 - start_2)
self.FPS10_2 = self.FPS10_2 + self.FPS_2
# print("callback")
# print(self.FPS_2)
# print("callback")
if self.cnt_2 % 10 == 0:
print("-----FPS(callback): " + str(self.FPS10_2 / 10))
print("-----Sec(callback): " + str(self.sec_2 / 10))
self.FPS10_2 = 0
self.sec_2 = 0
end_3 = time.time()
sec_3 = (end_3 - start_3)
# print("b:save and publish")
# print(sec_3)
print('---end-------------------------')
def webcam():
"""
Detect in a webcam
:return:
"""
with tf.Graph().as_default():
# Load model
mc = kitti_squeezeDet_config()
mc.BATCH_SIZE = 1
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDet(mc, FLAGS.gpu)
saver = tf.train.Saver(model.model_params)
with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess:
saver.restore(sess, FLAGS.checkpoint)
times = {}
count = 0
cap = cv2.VideoCapture(1)
while (True):
t_start = time.time()
count += 1
ret, frame = cap.read()
if ret==True:
# crop frames
#frame = frame[500:-205, 239:-439, :]
# Resize image to 375x1242
#dim = (1242, 375)
#frame = cv2.resize(frame, dim)
im_input = frame.astype(np.float32) - mc.BGR_MEANS
else:
break
t_reshape = time.time()
times['reshape'] = t_reshape - t_start
det_boxes, det_probs, det_class = sess.run(
[model.det_boxes, model.det_probs, model.det_class],
feed_dict={model.image_input: [im_input], model.keep_prob: 1.0})
t_detect = time.time()
times['detect'] = t_detect - t_reshape
# Filter
final_boxes, final_probs, final_class = model.filter_prediction(
det_boxes[0], det_probs[0], det_class[0])
keep_idx = [idx for idx in range(len(final_probs)) \
if final_probs[idx] > mc.PLOT_PROB_THRESH]
final_boxes = [final_boxes[idx] for idx in keep_idx]
final_probs = [final_probs[idx] for idx in keep_idx]
final_class = [final_class[idx] for idx in keep_idx]
t_filter = time.time()
times['filter'] = t_filter - t_detect
# TODO(bichen): move this color dict to configuration file
cls2clr = {
'car': (255, 191, 0),
'cyclist': (0, 191, 255),
'pedestrian': (255, 0, 191)
}
_draw_box(
frame, final_boxes,
[mc.CLASS_NAMES[idx] + ': (%.2f)' % prob \
for idx, prob in zip(final_class, final_probs)],
cdict=cls2clr
)
t_draw = time.time()
times['draw'] = t_draw - t_filter
cv2.imshow('Frame', frame)
times['total'] = time.time() - t_start
# time_str = ''
# for t in times:
# time_str += '{} time: {:.4f} '.format(t[0], t[1])
# time_str += '\n'
time_str = 'Total time: {:.4f}, detection time: {:.4f}, filter time: ' \
'{:.4f}'. \
format(times['total'], times['detect'], times['filter'])
print(time_str)
if cv2.waitKey(1) == 27:
break
cv2.destroyAllWindows()
def cam_demo(outDir):
"""Detect image."""
assert FLAGS.demo_net == 'squeezeDet' or FLAGS.demo_net == 'squeezeDet+', \
'Selected nueral net architecture not supported: {}'.format(FLAGS.demo_net)
with tf.Graph().as_default():
# Load model
if FLAGS.demo_net == 'squeezeDet':
mc = kitti_squeezeDet_config()
mc.BATCH_SIZE = 1
# model parameters will be restored from checkpoint
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDet(mc, FLAGS.gpu)
elif FLAGS.demo_net == 'squeezeDet+':
mc = kitti_squeezeDetPlus_config()
mc.BATCH_SIZE = 1
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDetPlus(mc, FLAGS.gpu)
saver = tf.train.Saver(model.model_params)
with tf.Session(config=tf.ConfigProto(
allow_soft_placement=True)) as sess:
saver.restore(sess, FLAGS.checkpoint)
idx_frame = 1
# for f in sorted(glob.iglob(FLAGS.input_path)):
cam = cv2.VideoCapture(0)
while True:
ret_val, img1 = cam.read()
img1 = rotate_bound(img1, 90)
# img1 = cv2.imread(f)
im = img1.astype(np.float32, copy=True)
im = cv2.resize(im, (mc.IMAGE_WIDTH, mc.IMAGE_HEIGHT))
img = cv2.resize(img1, (mc.IMAGE_WIDTH, mc.IMAGE_HEIGHT))
input_image = im - mc.BGR_MEANS
# Detect
det_boxes, det_probs, det_class = sess.run(
[model.det_boxes, model.det_probs, model.det_class],
feed_dict={model.image_input: [input_image]})
# Filter
final_boxes, final_probs, final_class = model.filter_prediction(
det_boxes[0], det_probs[0], det_class[0])
keep_idx = [idx for idx in range(len(final_probs)) \
if final_probs[idx] > 0.75] # mc.PLOT_PROB_THRESH
final_boxes = [final_boxes[idx] for idx in keep_idx]
final_probs = [final_probs[idx] for idx in keep_idx]
final_class = [final_class[idx] for idx in keep_idx]
# TODO(bichen): move this color dict to configuration file
cls2clr = {
'01ball': (255, 191, 0),
'02basket': (0, 191, 255),
# 'pedestrian':(255, 0, 191)
}
# Draw boxes
_draw_box(
img, final_boxes,
[mc.CLASS_NAMES[idx] + ': (%.2f)' % prob \
for idx, prob in zip(final_class, final_probs)],
cdict=cls2clr,
)
if not os.path.isdir(os.path.join(outDir, 'original_frames')):
os.makedirs(os.path.join(outDir, 'original_frames'))
if not os.path.isdir(os.path.join(outDir, 'result_frames')):
os.makedirs(os.path.join(outDir, 'result_frames'))
cv2.imwrite(
os.path.join(outDir, 'original_frames', '%06d.jpg') %
idx_frame, img1)
cv2.imwrite(
os.path.join(outDir, 'result_frames', '%06d.jpg') %
idx_frame, img)
idx_frame += 1
cv2.imshow('my webcam', img)
if cv2.waitKey(1) == 27:
break # esc to quit
cv2.destroyAllWindows()
def lj_caltech_test_demo():
"""Detect image."""
with tf.Graph().as_default():
# Load model
mc = caltech_vgg16_config()
mc.BATCH_SIZE = 1
mc.PLOT_PROB_THRESH = 0.01 ######################################## added by LJ
mc.NMS_THRESH = 0.4 ######################################## added by LJ
mc.PROB_THRESH = 0.01 ######################################## added by LJ
# model parameters will be restored from checkpoint
mc.LOAD_PRETRAINED_MODEL = False
model = VGG16ConvDet(mc, FLAGS.gpu)
saver = tf.train.Saver(model.model_params)
caltech_imdb = caltech(image_set='test', mc=mc)
with tf.Session(config=tf.ConfigProto(
allow_soft_placement=True)) as sess:
saver.restore(sess, FLAGS.checkpoint)
times = {}
for image_path in caltech_imdb.image_idx:
t_start = time.time()
im = cv2.imread(image_path)
scale_x = im.shape[0] / mc.IMAGE_HEIGHT
scale_y = im.shape[1] / mc.IMAGE_WIDTH
im = im.astype(np.float32, copy=False)
im = cv2.resize(im, (mc.IMAGE_WIDTH, mc.IMAGE_HEIGHT))
input_image = im - mc.BGR_MEANS
t_reshape = time.time()
times['reshape'] = t_reshape - t_start
# Detect
det_boxes, det_probs, det_class = sess.run(
[model.det_boxes, model.det_probs, model.det_class],
feed_dict={
model.image_input: [input_image],
model.keep_prob: 1.0
})
t_detect = time.time()
times['detect'] = t_detect - t_reshape
# Filter
final_boxes, final_probs, final_class = model.filter_prediction(
det_boxes[0], det_probs[0], det_class[0])
keep_idx = [idx for idx in range(len(final_probs)) \
if final_probs[idx] > mc.PLOT_PROB_THRESH]
final_boxes = [final_boxes[idx] for idx in keep_idx]
final_probs = [final_probs[idx] for idx in keep_idx]
final_class = [final_class[idx] for idx in keep_idx]
abs_path = '/home/bryant/MATLAB-tools/Matlab evaluation_labeling code3.2.1/data-USA/res/VGG+conv/'
txt_path = image_path[-21:-11] + '.txt'
txt_path = txt_path.replace('v', 'V')
txt_path = abs_path + txt_path
fram_id = int(image_path[-9:-4])
with open(txt_path, 'a') as f:
for bbox, label in zip(final_boxes, final_probs):
bbox = bbox_transform(bbox)
xmin, ymin, xmax, ymax = [b for b in bbox]
box_res = '{:d},{:.4f},{:.4f},{:.4f},{:.4f},{:.4f}\n'.format(
fram_id, xmin * scale_x, ymin * scale_y,
(xmax - xmin + 1) * scale_x,
(ymax - ymin + 1) * scale_y, label)
f.write(box_res)
f.close()
t_filter = time.time()
times['nms'] = t_filter - t_detect
# TODO(bichen): move this color dict to configuration file
cls2clr = {'pedestrian': (255, 0, 191)}
# Draw boxes
_draw_box(
im, final_boxes,
['%.2f'% prob \
for idx, prob in zip(final_class, final_probs)],
cdict=cls2clr,
)
t_draw = time.time()
times['draw'] = t_draw - t_filter
im = im.astype(np.uint8, copy=False)
cv2.imshow('img',
im) ############################### added by LJ
k = cv2.waitKey(1)
if k == 27:
pass #cv2.destroyWindow('img')
times['total'] = time.time() - t_start
# time_str = ''
# for t in times:
# time_str += '{} time: {:.4f} '.format(t[0], t[1])
# time_str += '\n'
time_str = 'Total time: {:.4f}, detection time: {:.4f}, filter time: '\
'{:.4f}'. \
format(times['total'], times['detect'], times['nms'])
print(time_str)
print('over')
def image_demo():
"""Detect image."""
assert FLAGS.demo_net == 'squeezeDet' or FLAGS.demo_net == 'squeezeDet+', \
'Selected nueral net architecture not supported: {}'.format(FLAGS.demo_net)
if flag_CSV:
csv = open("squeeze_th0.1_N512_v2.csv", "w") # squeeze, checkpoint999, random, plot > 0.005
columnTitleRow = "xmin,ymin,xmax,ymax,Frame,Label,Preview URL,confidence,random,y_loc,win_sizeX,win_sizeY\n"
csv.write(columnTitleRow)
with tf.Graph().as_default():
# Load model
if FLAGS.demo_net == 'squeezeDet':
mc = kitti_squeezeDet_config()
mc.BATCH_SIZE = 1
# model parameters will be restored from checkpoint
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDet(mc, FLAGS.gpu)
elif FLAGS.demo_net == 'squeezeDet+':
mc = kitti_squeezeDetPlus_config()
mc.BATCH_SIZE = 1
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDetPlus(mc, FLAGS.gpu)
saver = tf.train.Saver(model.model_params)
with tf.Session(config=tf.ConfigProto(allow_soft_placement=True)) as sess:
saver.restore(sess, FLAGS.checkpoint)
cnt = 0
FPS = 0
FPS10 = 0
for f in glob.iglob(FLAGS.input_path):
im = cv2.imread(f)
if flag_random:
"""random"""
randomX = random.randint(0, im.shape[1]-mc.IMAGE_WIDTH)
randomY = 400
im = im[randomY:randomY + mc.IMAGE_HEIGHT, randomX:randomX + mc.IMAGE_WIDTH, :]
else:
"""center"""
im = im[400:400 + mc.IMAGE_HEIGHT, 300:300 + mc.IMAGE_WIDTH]
im = im.astype(np.float32, copy=False)
# im = cv2.resize(im, (mc.IMAGE_WIDTH, mc.IMAGE_HEIGHT))
"""crop image"""
input_image = im - mc.BGR_MEANS
start = time.time()
# Detect
det_boxes, det_probs, det_class = sess.run(
[model.det_boxes, model.det_probs, model.det_class],
feed_dict={model.image_input:[input_image]})
print (mc.PLOT_PROB_THRESH)
cnt = cnt + 1
end = time.time()
FPS = (1 / (end - start))
FPS10 = FPS10 + FPS
# print ("FPS: " + str(FPS))
if cnt % 10 == 0:
print ("FPS(mean), detection: " + str(FPS10/10))
FPS10 = 0
# Filter
final_boxes, final_probs, final_class = model.filter_prediction(
det_boxes[0], det_probs[0], det_class[0])
keep_idx = [idx for idx in range(len(final_probs)) \
if final_probs[idx] > mc.PLOT_PROB_THRESH]
final_boxes = [final_boxes[idx] for idx in keep_idx]
final_probs = [final_probs[idx] for idx in keep_idx]
final_class = [final_class[idx] for idx in keep_idx]
# TODO(bichen): move this color dict to configuration file
cls2clr = {
'car': (255, 191, 0),
'cyclist': (0, 191, 255),
'pedestrian':(255, 0, 191)
}
# Draw boxes
_draw_box(
im, final_boxes,
[mc.CLASS_NAMES[idx]+': (%.2f)'% prob \
for idx, prob in zip(final_class, final_probs)],
cdict=cls2clr,
)
# print (final_boxes)
print (final_probs)
# print (final_class)
# print (mc.CLASS_NAMES)
# im2 = cv2.imread(f)
photo_name = f.split('/')[-1]
if flag_CSV:
for bbox_idx in range(len(final_boxes)):
bbox = bbox_transform(final_boxes[bbox_idx])
xmin, ymin, xmax, ymax = [int(b) for b in bbox]
csv.write(str(xmin+randomX))
csv.write(",")
csv.write(str(ymin+randomY))
csv.write(",")
csv.write(str(xmax+randomX))
csv.write(",")
csv.write(str(ymax+randomY))
csv.write(",")
"file name"
csv.write(photo_name)
csv.write(",")
"label"
csv.write(mc.CLASS_NAMES[final_class[bbox_idx]])
csv.write(",")
csv.write(",")
"confidence"
csv.write(str(final_probs[bbox_idx]))
csv.write(",")
"random selected window, x:"
csv.write(str(randomX))
csv.write(",")
"random selected window, Y:"
csv.write(str(randomY))
csv.write(",")
"random selected window, sizeX, X:"
csv.write(str(mc.IMAGE_WIDTH))
csv.write(",")
"random selected window, sizeY, Y:"
csv.write(str(mc.IMAGE_HEIGHT))
csv.write(",")
csv.write("\n")
# debug: offset random window size
# cv2.rectangle(im2, (xmin + randomX, ymin + randomY), (xmax + randomX, ymax + randomY), (0, 255, 0), 1)
if len(final_boxes) == 0:
print ("No detection: " + photo_name)
csv.write(",")
csv.write(",")
csv.write(",")
csv.write(",")
"file name"
csv.write(photo_name)
csv.write(",")
"label"
csv.write(",")
csv.write(",")
"confidence"
csv.write(",")
"random selected window, x:"
csv.write(str(randomX))
csv.write(",")
"random selected window, Y:"
csv.write(str(randomY))
csv.write(",")
"random selected window, sizeX, X:"
csv.write(str(mc.IMAGE_WIDTH))
csv.write(",")
"random selected window, sizeY, Y:"
csv.write(str(mc.IMAGE_HEIGHT))
csv.write(",")
csv.write("\n")
# for bbox, label in zip(final_boxes, label_list):
#
#
# xmin, ymin, xmax, ymax = [int(b) for b in bbox]
#
# l = label.split(':')[0] # text before "CLASS: (PROB)"
# if cdict and l in cdict:
# c = cdict[l]
# else:
# c = color
#
# # draw box
# cv2.rectangle(im, (xmin, ymin), (xmax, ymax), c, 1)
file_name = os.path.split(f)[1]
out_file_name = os.path.join(FLAGS.out_dir, 'out_'+file_name)
if cnt < 20:
cv2.imwrite(out_file_name, im)
print('Image detection output saved to {}'.format(out_file_name))
else:
print('(Skip)Image detection output saved to {}'.format(out_file_name))
def video_demo():
"""Detect videos."""
cap = cv2.VideoCapture(FLAGS.input_path)
# cap = cv2.VideoCapture(0)
print("FLAGS.input_path:", FLAGS.input_path)
if False == cap.isOpened():
print('open video failed')
return
else:
print('open video succeeded')
# Define the codec and create VideoWriter object
# fourcc = cv2.cv.CV_FOURCC(*'XVID')
fourcc = cv2.cv.CV_FOURCC('X', 'V', 'I', 'D')
# fourcc = cv2.cv.CV_FOURCC(*'MJPG')
# in_file_name = os.path.split(FLAGS.input_path)[1]
# out_file_name = os.path.join(FLAGS.out_dir, 'out_'+in_file_name)
out = cv2.VideoWriter("res.avi", fourcc, 30.0, (1242, 375), True)
# out = VideoWriter(out_file_name, frameSize=(1242, 375))
#out.open()
numFrames = cap.get(cv2.cv.CV_CAP_PROP_FRAME_COUNT)
with tf.Graph().as_default():
# Load model
mc = kitti_squeezeDet_config()
mc.BATCH_SIZE = 1
# model parameters will be restored from checkpoint
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDet(mc, FLAGS.gpu)
saver = tf.train.Saver(model.model_params)
cnt = 0
with tf.Session(config=tf.ConfigProto(
allow_soft_placement=True)) as sess:
saver.restore(sess, FLAGS.checkpoint)
times = {}
count = 0
while cap.isOpened():
#print ("open success!")
t_start = time.time()
count += 1
out_im_name = os.path.join(FLAGS.out_dir,
str(count).zfill(6) + '.jpg')
# Load images from video and crop
ret, frame = cap.read()
# print ("np.shape(frame):",np.shape(frame))
if ret == True:
# crop frames
#frame = frame[500:-205, 239:-439, :]
frame, ret = preprocess_frame(mc, frame)
if ret < 0:
print("preprocess_frame error")
return
im_input = frame.astype(np.float32) - mc.BGR_MEANS
else:
break
t_reshape = time.time()
times['reshape'] = t_reshape - t_start
# Detect
det_boxes, det_probs, det_class = sess.run(
[model.det_boxes, model.det_probs, model.det_class],
feed_dict={
model.image_input: [im_input],
model.keep_prob: 1.0
})
t_detect = time.time()
times['detect'] = t_detect - t_reshape
# Filter
final_boxes, final_probs, final_class = model.filter_prediction(
det_boxes[0], det_probs[0], det_class[0])
keep_idx = [idx for idx in range(len(final_probs)) \
if final_probs[idx] > mc.PLOT_PROB_THRESH]
final_boxes = [final_boxes[idx] for idx in keep_idx]
final_probs = [final_probs[idx] for idx in keep_idx]
final_class = [final_class[idx] for idx in keep_idx]
t_filter = time.time()
times['filter'] = t_filter - t_detect
# Draw boxes
# TODO(bichen): move this color dict to configuration file
cls2clr = {
'car': (255, 191, 0),
'cyclist': (0, 191, 255),
'pedestrian': (255, 0, 191)
}
_draw_box(
frame, final_boxes,
[mc.CLASS_NAMES[idx]+': (%.2f)'% prob \
for idx, prob in zip(final_class, final_probs)],
cdict=cls2clr
)
t_draw = time.time()
times['draw'] = t_draw - t_filter
cv2.imwrite(out_im_name, frame)
out.write(frame)
cv2.imshow("show", frame)
times['total'] = time.time() - t_start
# time_str = ''
# for t in times:
# time_str += '{} time: {:.4f} '.format(t[0], t[1])
# time_str += '\n'
time_str = 'Total time: {:.4f}, detection time: {:.4f}, filter time: '\
'{:.4f}'. \
format(times['total'], times['detect'], times['filter'])
#print (time_str)
if (cnt % 100 == 0):
print("cnt:", cnt, " allframes:", numFrames)
cnt = cnt + 1
if cv2.waitKey(1) & 0xFF == ord('q'):
break
# Release everything if job is finished
cap.release()
out.release()
cv2.destroyAllWindows()
def image_demo():
"""Detect image."""
assert FLAGS.demo_net == 'squeezeDet' or FLAGS.demo_net == 'squeezeDet+', \
'Selected nueral net architecture not supported: {}'.format(FLAGS.demo_net)
with tf.Graph().as_default():
# Load model
if FLAGS.demo_net == 'squeezeDet':
mc = kitti_squeezeDet_config()
mc.BATCH_SIZE = 1
# model parameters will be restored from checkpoint
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDet(mc, FLAGS.gpu)
elif FLAGS.demo_net == 'squeezeDet+':
mc = kitti_squeezeDetPlus_config()
mc.BATCH_SIZE = 1
mc.LOAD_PRETRAINED_MODEL = False
model = SqueezeDetPlus(mc, FLAGS.gpu)
saver = tf.train.Saver(model.model_params)
with tf.Session(config=tf.ConfigProto(
allow_soft_placement=True)) as sess:
saver.restore(sess, FLAGS.checkpoint)
fimgs = glob.glob(os.path.join(FLAGS.input_path, '*.jpg'))
for f in fimgs: #glob.iglob(FLAGS.input_path):
print(f)
im = cv2.imread(f)
im = im.astype(np.float32, copy=False)
im = cv2.resize(im, (mc.IMAGE_WIDTH, mc.IMAGE_HEIGHT))
input_image = im - mc.BGR_MEANS
# Detect
ss = time.time()
det_boxes, det_probs, det_class, det_pose, det_landmarks = sess.run(
[
model.det_boxes, model.det_probs, model.det_class,
model.det_pose, model.det_landmarks
],
feed_dict={model.image_input: [input_image]})
# Filter
final_boxes, final_landmarks, final_probs, final_pose, final_class = model.filter_prediction(
det_boxes[0], det_landmarks[0], det_probs[0], det_pose[0],
det_class[0])
keep_idx = [idx for idx in range(len(final_probs)) \
if final_probs[idx] > mc.PLOT_PROB_THRESH]
final_boxes = [final_boxes[idx] for idx in keep_idx]
final_probs = [final_probs[idx] for idx in keep_idx]
final_landmarks = [final_landmarks[idx] for idx in keep_idx]
final_pose = [final_pose[idx] for idx in keep_idx]
final_class = [final_class[idx] for idx in keep_idx]
# TODO(bichen): move this color dict to configuration file
cls2clr = {
'01cry': (255, 191, 0),
'02cute': (0, 191, 255),
'03normal': (255, 0, 191),
'04sleep': (191, 0, 255),
'05wakeup': (191, 255, 0)
}
pose2clr = {
'01frontal': (255, 191, 0),
'02nonfrontal': (0, 191, 255)
}
ee = time.time() - ss
print('elapsed time: ', ee)
# Draw boxes
_draw_box(
im, final_boxes,
[mc.CLASS_NAMES[idx] + ': (%.2f)' % prob \
for idx, prob in zip(final_class, final_probs)], [mc.POSE_NAMES[idx] for idx in final_pose],
cdict=cls2clr,
)
_draw_point(im, final_landmarks, (0, 0, 255))
file_name = os.path.split(f)[1]
out_file_name = os.path.join(FLAGS.out_dir, 'out_' + file_name)
cv2.imwrite(out_file_name, im)
print(
'Image detection output saved to {}'.format(out_file_name))
