def main(argv=None):
img = Image.open(FLAGS.input_img)
img_resized = img.resize(size=(FLAGS.size, FLAGS.size))
classes = load_coco_names(FLAGS.class_names)
# placeholder for detector inputs
inputs = tf.placeholder(tf.float32, [None, FLAGS.size, FLAGS.size, 3])
with tf.variable_scope('detector'):
#detections = yolo_v3(inputs, len(classes), data_format='NCHW')
detections = yolo_v3(inputs, len(classes), data_format='NHWC')
load_ops = load_weights(tf.global_variables(scope='detector'),
FLAGS.weights_file)
boxes = detections_boxes(detections)
with tf.Session() as sess:
sess.run(load_ops)
detected_boxes = sess.run(
boxes,
feed_dict={inputs: [np.array(img_resized, dtype=np.float32)]})
filtered_boxes = non_max_suppression(
detected_boxes,
confidence_threshold=FLAGS.conf_threshold,
iou_threshold=FLAGS.iou_threshold)
draw_boxes(filtered_boxes, img, classes, (FLAGS.size, FLAGS.size))
img.save(FLAGS.output_img)
def main(argv=None):
img = Image.open(FLAGS.input_img)
img_resized = img.resize(size=(FLAGS.size, FLAGS.size))
classes = load_coco_names(FLAGS.class_names)
# placeholder for detector inputs
inputs = tf.placeholder(tf.float32, [None, FLAGS.size, FLAGS.size, 3])
with tf.variable_scope('detector'):
detections = yolo_v3(inputs, len(classes), data_format='NHWC')
load_ops = load_weights(tf.global_variables(scope='detector'),
FLAGS.weights_file)
boxes = detections_boxes(detections)
with tf.Session() as sess:
sess.run(load_ops)
frozen = tf.graph_util.convert_variables_to_constants(
sess, sess.graph_def, ['concat_1'])
graph_io.write_graph(frozen, './', 'yolo_v3.pb', as_text=False)
detected_boxes = sess.run(
boxes,
feed_dict={inputs: [np.array(img_resized, dtype=np.float32)]})
filtered_boxes = non_max_suppression(
detected_boxes,
confidence_threshold=FLAGS.conf_threshold,
iou_threshold=FLAGS.iou_threshold)
draw_boxes(filtered_boxes, img, classes, (FLAGS.size, FLAGS.size))
img.save(FLAGS.output_img)
def init_yolo(sess, inputs, num_classes, weights, header_size=5):
with tf.variable_scope('detector'):
detections = yolo_v3(inputs, num_classes, data_format='NHWC')
load_ops = load_weights(tf.global_variables(scope='detector'),
weights,
header_size=header_size)
boxes = detections_boxes(detections)
sess.run(load_ops)
return detections, boxes
def main(argv=None):
img = Image.open(FLAGS.input_img)
img_resized = img.resize(size=(FLAGS.size, FLAGS.size))
classes = load_coco_names(FLAGS.class_names)
# placeholder for detector inputs
inputs = tf.placeholder(tf.float32, [1, FLAGS.size, FLAGS.size, 3],
name="input")
with tf.variable_scope('detector'):
detections = yolo_v3(inputs, len(classes), data_format='NHWC')
load_ops = load_weights(tf.global_variables(scope='detector'),
FLAGS.weights_file)
boxes = detections_boxes(detections)
graph = tf.get_default_graph()
output_graph = os.path.join(MODEL_DIR, MODEL_NAME) # PB模型保存路径
graph_def = graph.as_graph_def()
with tf.Session() as sess:
sess.run(load_ops)
detected_boxes = sess.run(
boxes,
feed_dict={inputs: [np.array(img_resized, dtype=np.float32)]})
output_graph_def = graph_util.convert_variables_to_constants( # 模型持久化，将变量值固定
sess,
graph_def,
["output"] # 如果有多个输出节点，以逗号隔开
)
with tf.gfile.GFile(output_graph, "wb") as f: # 保存模型
f.write(output_graph_def.SerializeToString()) # 序列化输出
print("%d ops in the final graph." %
len(output_graph_def.node)) # 得到当前图有几个操作节点
print("detected_boxes[0].shape:", detected_boxes[0].shape)
print("detected_boxes:", detected_boxes)
filtered_boxes = non_max_suppression(
detected_boxes,
confidence_threshold=FLAGS.conf_threshold,
iou_threshold=FLAGS.iou_threshold)
draw_boxes(filtered_boxes, img, classes, (FLAGS.size, FLAGS.size))
img.save(FLAGS.output_img)
print("done")
def main(argv=None):
BASE_PATH = 'images'
TEST_IMAGES = os.listdir(BASE_PATH)
TEST_IMAGES.sort()
print(TEST_IMAGES)
# img = Image.open(FLAGS.input_img)
# w,h = img.size
# img_resized = img.resize(size=(FLAGS.size, FLAGS.size))
classes = load_coco_names(FLAGS.class_names)
# placeholder for detector inputs
inputs = tf.placeholder(tf.float32, [None, FLAGS.size, FLAGS.size, 3])
with tf.variable_scope('detector'):
detections = yolo_v3(inputs, len(classes), data_format='NHWC')#Tensor("detector/yolo-v3/concat:0", shape=(?, 10647, 85), dtype=float32)
load_ops = load_weights(tf.global_variables(scope='detector'), FLAGS.weights_file)
boxes = detections_boxes(detections)#shape=(?, 10647, 85), dtype=float32)
#coordinates of top left and bottom right points+num_class_confidence
saver = tf.train.Saver()
with tf.Session() as sess:
sess.run(load_ops)
writer =tf.summary.FileWriter("logs/",graph = sess.graph)
writer.close()
saver.save(sess,"models/yolov3.ckpt")
for img in TEST_IMAGES:
image_path = os.path.join(BASE_PATH, img)
image = Image.open(image_path)
w,h = image.size
img_resized = image.resize(size=(FLAGS.size, FLAGS.size))
detected_boxes = sess.run(boxes, feed_dict={inputs: [np.array(img_resized, dtype=np.float32)]})
filtered_boxes = non_max_suppression(detected_boxes, confidence_threshold=FLAGS.conf_threshold,
iou_threshold=FLAGS.iou_threshold)
draw_boxes(filtered_boxes, image, classes, (FLAGS.size, FLAGS.size))
plt.imshow(image)
plt.show()
image.save(FLAGS.output_img)
def main(argv=None):
# placeholder for detector inputs
classes = load_coco_names(FLAGS.class_names)
inputs = tf.placeholder(tf.float32, [None, None, None, 3], name="input")
resize_method = tf.image.ResizeMethod.BILINEAR
inputs_resized = tf.image.resize_images(images=inputs,
size=[FLAGS.size, FLAGS.size],
method=resize_method,
align_corners=True)
with tf.variable_scope('detector'):
# format = 'NCHW'
format = 'NHWC'
detections = yolo_v3(inputs_resized, len(classes), data_format=format)
load_ops = load_weights(tf.global_variables(scope='detector'),
FLAGS.weights_file)
boxes = detections_boxes(detections, FLAGS.size, FLAGS.size)
boxes = tf.identity(boxes, "output")
config = tf.ConfigProto(device_count={'GPU': 1})
config.gpu_options.allow_growth = True
config.gpu_options.per_process_gpu_memory_fraction = 0.1
if FLAGS.use_xla:
jit_level = tf.OptimizerOptions.ON_1
config.graph_options.optimizer_options.global_jit_level = jit_level
with tf.Session(config=config) as sess:
sess.run(tf.global_variables_initializer())
sess.run(tf.local_variables_initializer())
sess.run(load_ops)
print("Finished loading graph")
print("%d ops in the base graph." % len(sess.graph.get_operations()))
# save the (trained) model
# We use a built-in TF helper to export variables to constants
output_node_names = 'output'
output_graph_def = tf.graph_util.convert_variables_to_constants(
sess, # The session is used to retrieve the weights
tf.get_default_graph().as_graph_def(
), # The graph_def is used to retrieve the nodes
output_node_names.split(
","
) # The output node names are used to select the usefull nodes
)
tf.train.write_graph(output_graph_def, '', FLAGS.output_file, False)
print("%d ops in the final graph." % len(output_graph_def.node))
def main(argv=None):
parser = argparse.ArgumentParser()
parser.add_argument('camid', type=int, help='source webcam id')
args = parser.parse_args()
classes = load_coco_names(FLAGS.class_names)
np.random.seed(2018)
colors = [np.random.randint(0, 255, 3) for _ in range(len(classes))]
# placeholder for detector inputs
inputs = tf.placeholder(tf.float32, [None, FLAGS.size, FLAGS.size, 3])
with tf.variable_scope('detector'):
detections = yolo_v3(inputs, len(classes), data_format='NCHW')
load_ops = load_weights(tf.global_variables(scope='detector'),
FLAGS.weights_file)
boxes = detections_boxes(detections)
vc = cv2.VideoCapture()
vc.open(args.camid)
with tf.Session() as sess:
sess.run(load_ops)
while True:
_, img = vc.read()
img_resized = cv2.resize(img, dsize=(FLAGS.size, FLAGS.size))
img_resized = cv2.cvtColor(img_resized, cv2.COLOR_BGR2RGB)
detected_boxes = sess.run(
boxes,
feed_dict={inputs: [np.array(img_resized, dtype=np.float32)]})
filtered_boxes = non_max_suppression(
detected_boxes,
confidence_threshold=FLAGS.conf_threshold,
iou_threshold=FLAGS.iou_threshold)
img = cv2.resize(img, (1920, 1080))
img = draw_boxes(filtered_boxes,
img,
classes, (FLAGS.size, FLAGS.size),
colors=colors)
# img_resized = draw_boxes(filtered_boxes, img_resized, classes, (FLAGS.size, FLAGS.size))
cv2.imshow('detections', img)
# cv2.imshow('img_resized', img_resized)
key = cv2.waitKey(1)
if key == ord('q') or key & 0xFFFF == 27:
break
def train(sess,
X_data,
y_data,
validation_set=None,
initialize=True,
epochs=20,
shuffle=True,
dropout=0.5,
random_seed=None):
training_loss = []
# placeholder for detector inputs
X = tf.placeholder(tf.float32, [None, FLAGS.size, FLAGS.size, 3], name='X')
y = tf.placeholder(tf.float32, [None, 1, 5], name='y')
# Loss and optimizer
cost = calculate_loss(X, y)
##TODO: change optimizer
##TODO:add learning_rate=learning_rate
optimizer = tf.train.RMSPropOptimizer(learning_rate=0.5).minimize(3)
if initialize:
sess.run(tf.global_variables_initializer())
with tf.variable_scope('detector'):
load_ops = load_weights(tf.global_variables(scope='detector'),
FLAGS.weights_file)
sess.run(load_ops)
np.random.seed(random_seed) # for shuffling in batch_generator
for epoch in range(1, epochs + 1):
batch_gen = batch_generator(X_data, y_data, shuffle=shuffle)
avg_loss = 0.0
for batch_idx, (batch_x, batch_y) in enumerate(batch_gen):
feed = {'X': batch_x, 'y': batch_y}
_, loss = sess.run(optimizer, cost, feed_dict=feed)
avg_loss += loss
print('Batch Loss : %7.3f' % loss, end=' ')
training_loss.append(avg_loss / (batch_idx + 1))
print('Epoch %02d Training Avg. Loss: %7.3f' % (epoch, avg_loss),
end=' ')
def main(argv=None):
classes = load_coco_names(FLAGS.class_names)
# placeholder for detector inputs
inputs = tf.placeholder(tf.float32, [None, FLAGS.size, FLAGS.size, 3])
with tf.variable_scope('detector'):
detections = yolo_v3(inputs, len(classes), data_format='NHWC')
load_ops = load_weights(tf.global_variables(scope='detector'),
FLAGS.weights_file)
boxes = detections_boxes(detections)
with tf.Session() as sess:
sess.run(load_ops)
from tensorflow.python.framework import graph_io
frozen = tf.graph_util.convert_variables_to_constants(
sess, sess.graph_def, ['concat_1'])
graph_io.write_graph(frozen,
'../build/YOLOv3/',
'yolo_v3.pb',
as_text=False)
def detect_obj(file_path):
img = Image.open(file_path)
img_resized = img.resize(size=(FLAGS.size, FLAGS.size))
classes = load_coco_names(FLAGS.class_names)
# placeholder for detector inputs
inputs = tf.placeholder(tf.float32, [None, FLAGS.size, FLAGS.size, 3])
with tf.variable_scope('detector'):
#detections = yolo_v3(inputs, len(classes), data_format='NCHW')
detections = yolo_v3(inputs, len(classes), data_format='NHWC')
#detections = yolo_v3_tiny(inputs, len(classes), data_format='NHWC')
load_ops = load_weights(tf.global_variables(scope='detector'),
FLAGS.weights_file)
boxes = detections_boxes(detections)
#saver = tf.train.Saver()
with tf.Session() as sess:
sess.run(load_ops)
#saver.save(sess, "./yolov3-coco")
detected_boxes = sess.run(
boxes,
feed_dict={inputs: [np.array(img_resized, dtype=np.float32)]})
filtered_boxes = non_max_suppression(
detected_boxes,
confidence_threshold=FLAGS.conf_threshold,
iou_threshold=FLAGS.iou_threshold)
draw_boxes(filtered_boxes, img, classes, (FLAGS.size, FLAGS.size))
img.save(os.path.join(FLAGS.output_dir, os.path.basename(file_path)))
tf.reset_default_graph()
