def main(argv=None): assert FLAGS.base_dir, 'Directory base_dir not set.' assert os.path.exists(FLAGS.base_dir), \ 'Directory %s not exist.' % FLAGS.base_dir os.makedirs(FLAGS.output_dir, exist_ok=True) files = os.listdir(FLAGS.base_dir) json_files = filter(lambda f: f.endswith('.json'), files) json_file_paths = list( map(lambda f: os.path.join(FLAGS.base_dir, f), json_files)) feature_map = tf.get_variable(shape=FEATURE_MAP_SHAPE, name='feature_map') ssd_layers = create_ssd_layers(feature_map) file_count = len(json_file_paths) total_face_count = 0 total_match_count = 0 for index, json_path in enumerate(json_file_paths): for rotate in range(4): annotation = _load_annotation(json_path, rotate=rotate) image_file_name = annotation['file_name'] regions = annotation['regions'] face_regions = list( filter(lambda region: region['label'] == 0, regions)) boxes = create_boxes(ssd_layers) face_count, match_count = _assign_box(face_regions, boxes) total_face_count += face_count total_match_count += match_count image_path = os.path.join(FLAGS.base_dir, image_file_name) image = _create_image_data(image_path, rotate=rotate) name, _ = os.path.splitext(image_file_name) tfrecord_name = '%s_%d.tfrecord' % (name, rotate) tfrecord_path = os.path.join(FLAGS.output_dir, tfrecord_name) _create_tfrecord(image, image_file_name, boxes, tfrecord_path) if FLAGS.debug: debug_image_file_name = '%s_%d.jpg' % (name, rotate) debug_image_file_path = os.path.join(FLAGS.output_dir, debug_image_file_name) _create_debug_image(image, boxes, debug_image_file_path) percentage = total_match_count / float(total_face_count) print('%d/%d, %d/%d - %f' % (index + 1, file_count, total_match_count, total_face_count, percentage))
def _export_boxes_position(feature_map, output_dir): box_layers = create_ssd_layers(feature_map) boxes = create_boxes(box_layers) box_position = [] for index, box in enumerate(boxes): box_position.append({ 'index': index, 'left': box.left, 'top': box.top, 'right': box.right, 'bottom': box.bottom, }) file_name = '%s_boxes_position.json' % model.NAME box_position_path = os.path.join(output_dir, file_name) with open(box_position_path, mode='w') as fp: json.dump(box_position, fp, indent=4)
def main(argv=None): assert FLAGS.train_path, 'train_path not set' assert FLAGS.image_path, 'image_path not set' assert os.path.exists(FLAGS.image_path), '%s is not exist' % FLAGS.image_path print(FLAGS.image_path) os.makedirs(FLAGS.output_dir, exist_ok=True) raw_image, resized_image = _get_image_data(FLAGS.image_path, model.IMAGE_SIZE) image = np.array(resized_image.getdata()) \ .reshape(model.IMAGE_SIZE, model.IMAGE_SIZE, model.CHANNELS) \ .astype(np.float32) file_name = os.path.basename(FLAGS.image_path) name, _ = os.path.splitext(file_name) output_image_file_name = '%s.jpg' % name output_image_path = os.path.join(FLAGS.output_dir, output_image_file_name) output_result_image_file_name = '%s_result.jpg' % name output_result_image_path = os.path.join(FLAGS.output_dir, output_result_image_file_name) global_step = tf.Variable(0, trainable=False) image_ph = tf.placeholder(tf.float32, [model.IMAGE_SIZE, model.IMAGE_SIZE, model.CHANNELS]) normalized_image = image_ph * (1 / 255.0) image_batch = tf.expand_dims(normalized_image, axis=0) feature_map = model.base_layers(image_batch, is_train=False) ssd_logits = model.ssd_layers(feature_map, is_train=False) loc_logits = ssd_logits[:, :, :4] conf_logits = ssd_logits[:, :, 4:] loc_logits = tf.nn.tanh(loc_logits) conf_logits = tf.nn.sigmoid(conf_logits) ssd_logits = tf.concat([loc_logits, conf_logits], axis=2) saver = tf.train.Saver(var_list=tf.global_variables(), max_to_keep=3) with tf.Session() as sess: sess.run(tf.global_variables_initializer()) saver.restore(sess, FLAGS.train_path) start = time.time() ssd_logits_values = sess.run(ssd_logits, feed_dict={ image_ph: image }) elapsed = time.time() - start print('Inference Elapsed %f s' % elapsed) start = time.time() ssd_layers = create_ssd_layers(feature_map) boxes = create_boxes(ssd_layers) for index, box_value in enumerate(ssd_logits_values[0]): boxes[index].label = box_value[4:] boxes[index].offset = box_value[:4] filtered_boxes = filter( lambda box: box.label[0] > FLAGS.confidence_threshold, boxes) faces = list(map(lambda box: Face(box), filtered_boxes)) _merge_faces(faces, overlap_threshold=0.3) elapsed = time.time() - start print('Post-process Elapsed %f s' % elapsed) result_image = _draw_recognized_faces(raw_image.copy(), faces) result_image.save(output_result_image_path, format='jpeg') raw_image.save(output_image_path, format='jpeg')