Exemplo n.º 1
0
def main(argv):
    if len(argv) > 1:
        raise RuntimeError('Too many command-line arguments.')

    # Read list of images.
    print('Reading list of images...')
    image_paths = _ReadImageList(cmd_args.list_images_path)
    num_images = len(image_paths)
    print(f'done! Found {num_images} images')

    # Create output directories if necessary.
    if not tf.io.gfile.exists(cmd_args.output_dir):
        tf.io.gfile.makedirs(cmd_args.output_dir)
    if cmd_args.output_viz_dir and not tf.io.gfile.exists(
            cmd_args.output_viz_dir):
        tf.io.gfile.makedirs(cmd_args.output_viz_dir)

    detector_fn = detector.MakeDetector(cmd_args.detector_path)

    start = time.time()
    for i, image_path in enumerate(image_paths):
        # Report progress once in a while.
        if i == 0:
            print('Starting to detect objects in images...')
        elif i % _STATUS_CHECK_ITERATIONS == 0:
            elapsed = (time.time() - start)
            print(f'Processing image {i} out of {num_images}, last '
                  f'{_STATUS_CHECK_ITERATIONS} images took {elapsed} seconds')
            start = time.time()

        # If descriptor already exists, skip its computation.
        base_boxes_filename, _ = os.path.splitext(os.path.basename(image_path))
        out_boxes_filename = base_boxes_filename + _BOX_EXT
        out_boxes_fullpath = os.path.join(cmd_args.output_dir,
                                          out_boxes_filename)
        if tf.io.gfile.exists(out_boxes_fullpath):
            print(f'Skipping {image_path}')
            continue

        im = np.expand_dims(np.array(utils.RgbLoader(image_paths[i])), 0)

        # Extract and save boxes.
        (boxes_out, scores_out, class_indices_out) = detector_fn(im)
        (selected_boxes, selected_scores,
         selected_class_indices) = _FilterBoxesByScore(
             boxes_out[0], scores_out[0], class_indices_out[0],
             cmd_args.detector_thresh)

        box_io.WriteToFile(out_boxes_fullpath, selected_boxes, selected_scores,
                           selected_class_indices)
        if cmd_args.output_viz_dir:
            out_viz_filename = base_boxes_filename + _VIZ_SUFFIX
            out_viz_fullpath = os.path.join(cmd_args.output_viz_dir,
                                            out_viz_filename)
            _PlotBoxesAndSaveImage(im[0], selected_boxes, out_viz_fullpath)
Exemplo n.º 2
0
def ExtractBoxesAndFeaturesToFiles(image_names, image_paths, delf_config_path,
                                   detector_model_dir, detector_thresh,
                                   output_features_dir, output_boxes_dir,
                                   output_mapping):
    """Extracts boxes and features, saving them to files.

  Boxes are saved to <image_name>.boxes files. DELF features are extracted for
  the entire image and saved into <image_name>.delf files. In addition, DELF
  features are extracted for each high-confidence bounding box in the image, and
  saved into files named <image_name>_0.delf, <image_name>_1.delf, etc.

  It checks if descriptors/boxes already exist, and skips computation for those.

  Args:
    image_names: List of image names. These are used to compose output file
      names for boxes and features.
    image_paths: List of image paths. image_paths[i] is the path for the image
      named by image_names[i]. `image_names` and `image_paths` must have the
      same number of elements.
    delf_config_path: Path to DelfConfig proto text file.
    detector_model_dir: Directory where detector SavedModel is located.
    detector_thresh: Threshold used to decide if an image's detected box
      undergoes feature extraction.
    output_features_dir: Directory where DELF features will be written to.
    output_boxes_dir: Directory where detected boxes will be written to.
    output_mapping: CSV file which maps each .delf file name to the image ID and
      detected box ID.

  Raises:
    ValueError: If len(image_names) and len(image_paths) are different.
  """
    num_images = len(image_names)
    if len(image_paths) != num_images:
        raise ValueError(
            'image_names and image_paths have different number of items')

    # Parse DelfConfig proto.
    config = delf_config_pb2.DelfConfig()
    with tf.io.gfile.GFile(delf_config_path, 'r') as f:
        text_format.Merge(f.read(), config)

    # Create output directories if necessary.
    if not tf.io.gfile.exists(output_features_dir):
        tf.io.gfile.makedirs(output_features_dir)
    if not tf.io.gfile.exists(output_boxes_dir):
        tf.io.gfile.makedirs(output_boxes_dir)
    if not tf.io.gfile.exists(os.path.dirname(output_mapping)):
        tf.io.gfile.makedirs(os.path.dirname(output_mapping))

    names_ids_and_boxes = []
    with tf.Graph().as_default():
        with tf.compat.v1.Session() as sess:
            # Initialize variables, construct detector and DELF extractor.
            init_op = tf.compat.v1.global_variables_initializer()
            sess.run(init_op)
            detector_fn = detector.MakeDetector(sess,
                                                detector_model_dir,
                                                import_scope='detector')
            delf_extractor_fn = extractor.MakeExtractor(
                sess, config, import_scope='extractor_delf')

            start = time.clock()
            for i in range(num_images):
                if i == 0:
                    print('Starting to extract features/boxes...')
                elif i % _STATUS_CHECK_ITERATIONS == 0:
                    elapsed = (time.clock() - start)
                    print('Processing image %d out of %d, last %d '
                          'images took %f seconds' %
                          (i, num_images, _STATUS_CHECK_ITERATIONS, elapsed))
                    start = time.clock()

                image_name = image_names[i]
                output_feature_filename_whole_image = os.path.join(
                    output_features_dir, image_name + _DELF_EXTENSION)
                output_box_filename = os.path.join(output_boxes_dir,
                                                   image_name + _BOX_EXTENSION)

                pil_im = utils.RgbLoader(image_paths[i])
                width, height = pil_im.size

                # Extract and save boxes.
                if tf.io.gfile.exists(output_box_filename):
                    print('Skipping box computation for %s' % image_name)
                    (boxes_out, scores_out, class_indices_out
                     ) = box_io.ReadFromFile(output_box_filename)
                else:
                    (boxes_out, scores_out, class_indices_out) = detector_fn(
                        np.expand_dims(pil_im, 0))
                    # Using only one image per batch.
                    boxes_out = boxes_out[0]
                    scores_out = scores_out[0]
                    class_indices_out = class_indices_out[0]
                    box_io.WriteToFile(output_box_filename, boxes_out,
                                       scores_out, class_indices_out)

                # Select boxes with scores greater than threshold. Those will be the
                # ones with extracted DELF features (besides the whole image, whose DELF
                # features are extracted in all cases).
                num_delf_files = 1
                selected_boxes = []
                for box_ind, box in enumerate(boxes_out):
                    if scores_out[box_ind] >= detector_thresh:
                        selected_boxes.append(box)
                num_delf_files += len(selected_boxes)

                # Extract and save DELF features.
                for delf_file_ind in range(num_delf_files):
                    if delf_file_ind == 0:
                        box_name = image_name
                        output_feature_filename = output_feature_filename_whole_image
                    else:
                        box_name = image_name + '_' + str(delf_file_ind - 1)
                        output_feature_filename = os.path.join(
                            output_features_dir, box_name + _DELF_EXTENSION)

                    names_ids_and_boxes.append(
                        [box_name, i, delf_file_ind - 1])

                    if tf.io.gfile.exists(output_feature_filename):
                        print('Skipping DELF computation for %s' % box_name)
                        continue

                    if delf_file_ind >= 1:
                        bbox_for_cropping = selected_boxes[delf_file_ind - 1]
                        bbox_for_cropping_pil_convention = [
                            int(math.floor(bbox_for_cropping[1] * width)),
                            int(math.floor(bbox_for_cropping[0] * height)),
                            int(math.ceil(bbox_for_cropping[3] * width)),
                            int(math.ceil(bbox_for_cropping[2] * height))
                        ]
                        pil_cropped_im = pil_im.crop(
                            bbox_for_cropping_pil_convention)
                        im = np.array(pil_cropped_im)
                    else:
                        im = np.array(pil_im)

                    extracted_features = delf_extractor_fn(im)
                    locations_out = extracted_features['local_features'][
                        'locations']
                    descriptors_out = extracted_features['local_features'][
                        'descriptors']
                    feature_scales_out = extracted_features['local_features'][
                        'scales']
                    attention_out = extracted_features['local_features'][
                        'attention']

                    feature_io.WriteToFile(output_feature_filename,
                                           locations_out, feature_scales_out,
                                           descriptors_out, attention_out)

    # Save mapping from output DELF name to image id and box id.
    _WriteMappingBasenameToIds(names_ids_and_boxes, output_mapping)
Exemplo n.º 3
0
def main(argv):
    if len(argv) > 1:
        raise RuntimeError('Too many command-line arguments.')

    tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.INFO)

    # Read list of images.
    tf.compat.v1.logging.info('Reading list of images...')
    image_paths = _ReadImageList(cmd_args.list_images_path)
    num_images = len(image_paths)
    tf.compat.v1.logging.info('done! Found %d images', num_images)

    # Create output directories if necessary.
    if not tf.io.gfile.exists(cmd_args.output_dir):
        tf.io.gfile.makedirs(cmd_args.output_dir)
    if cmd_args.output_viz_dir and not tf.io.gfile.exists(
            cmd_args.output_viz_dir):
        tf.io.gfile.makedirs(cmd_args.output_viz_dir)

    # Tell TensorFlow that the model will be built into the default Graph.
    with tf.Graph().as_default():
        # Reading list of images.
        filename_queue = tf.compat.v1.train.string_input_producer(
            image_paths, shuffle=False)
        reader = tf.compat.v1.WholeFileReader()
        _, value = reader.read(filename_queue)
        image_tf = tf.io.decode_jpeg(value, channels=3)
        image_tf = tf.expand_dims(image_tf, 0)

        with tf.compat.v1.Session() as sess:
            init_op = tf.compat.v1.global_variables_initializer()
            sess.run(init_op)

            detector_fn = detector.MakeDetector(sess, cmd_args.detector_path)

            # Start input enqueue threads.
            coord = tf.train.Coordinator()
            threads = tf.compat.v1.train.start_queue_runners(sess=sess,
                                                             coord=coord)
            start = time.clock()
            for i, image_path in enumerate(image_paths):
                # Write to log-info once in a while.
                if i == 0:
                    tf.compat.v1.logging.info(
                        'Starting to detect objects in images...')
                elif i % _STATUS_CHECK_ITERATIONS == 0:
                    elapsed = (time.clock() - start)
                    tf.compat.v1.logging.info(
                        'Processing image %d out of %d, last %d '
                        'images took %f seconds', i, num_images,
                        _STATUS_CHECK_ITERATIONS, elapsed)
                    start = time.clock()

                # # Get next image.
                im = sess.run(image_tf)

                # If descriptor already exists, skip its computation.
                base_boxes_filename, _ = os.path.splitext(
                    os.path.basename(image_path))
                out_boxes_filename = base_boxes_filename + _BOX_EXT
                out_boxes_fullpath = os.path.join(cmd_args.output_dir,
                                                  out_boxes_filename)
                if tf.io.gfile.exists(out_boxes_fullpath):
                    tf.compat.v1.logging.info('Skipping %s', image_path)
                    continue

                # Extract and save boxes.
                (boxes_out, scores_out, class_indices_out) = detector_fn(im)
                (selected_boxes, selected_scores,
                 selected_class_indices) = _FilterBoxesByScore(
                     boxes_out[0], scores_out[0], class_indices_out[0],
                     cmd_args.detector_thresh)

                box_io.WriteToFile(out_boxes_fullpath, selected_boxes,
                                   selected_scores, selected_class_indices)
                if cmd_args.output_viz_dir:
                    out_viz_filename = base_boxes_filename + _VIZ_SUFFIX
                    out_viz_fullpath = os.path.join(cmd_args.output_viz_dir,
                                                    out_viz_filename)
                    _PlotBoxesAndSaveImage(im[0], selected_boxes,
                                           out_viz_fullpath)

            # Finalize enqueue threads.
            coord.request_stop()
            coord.join(threads)