text: from __future__ import print_function
  import argparse
  import json
  import os
  import pickle
  
  import cv2
  import tensorflow as tf
  from utils.helpers import load_module
  from ssd_detector.readers.object_detector_json import ObjectDetectorJson
  from ssd_detector.trainer import create_session, detection_model, InputInferData, InputValData
  
  
  def parse_args():
    parser = argparse.ArgumentParser(
      description='Perform inference of a detection model on a video file or an annotation file in JSON format')
  
    parser.add_argument('path_to_config', help='Path to a config.py')
  
    input_type_group = parser.add_mutually_exclusive_group(required=True)
    input_type_group.add_argument('--json', dest='input_type', action='store_const', const='json',
                                  help='Get images from annotation file in JSON format')
    input_type_group.add_argument('--video', dest='input_type', action='store_const', const='video',
                                  help='Get images from video file')
    input_type_group.set_defaults(input_type='json')
  
    parser.add_argument('--input', help='Path to the input file', required=True)
    parser.add_argument('--conf_threshold', type=float, help='Confidence threshold for detector', default=0.1)
  
    dump_to_json_group = parser.add_mutually_exclusive_group()
    dump_to_json_group.add_argument('--dump-to-json', dest='dump_to_json', action='store_true')
    dump_to_json_group.add_argument('--no-dump-to-json', dest='dump_to_json', action='store_false')
    dump_to_json_group.set_defaults(dump_to_json=True)
  
    show_group = parser.add_mutually_exclusive_group()
    show_group.add_argument('--show', dest='show', action='store_true')
    show_group.add_argument('--no-show', dest='show', action='store_false')
    show_group.set_defaults(show=False)
  
    dump_output_video_group = parser.add_mutually_exclusive_group()
    dump_output_video_group.add_argument('--dump-to-video', dest='dump_output_video', action='store_true')
    dump_output_video_group.add_argument('--no-dump-to-video', dest='dump_output_video', action='store_false')
    dump_output_video_group.set_defaults(dump_output_video=False)
  
    parser.add_argument('--path_to_output_video', help='Path to output video with predictions', default='output.avi')
    return parser.parse_args()
  
  
  # pylint: disable=too-many-locals
  def process_image(predictions, img_size, img_id, conf_threshold, classes):
    img_width, img_height = img_size
  
    coco_detections = []
    for prediction in predictions:
      det_label = int(prediction[0])
      det_conf = float(prediction[1])
  
      if conf_threshold and det_conf < conf_threshold:
        continue
  
      top_left_x = float(prediction[2] * img_width)
      top_left_y = float(prediction[3] * img_height)
      bottom_right_x = float(prediction[4] * img_width)
      bottom_right_y = float(prediction[5] * img_height)
      if det_label >= len(classes):
        print('Wrong label: {0}'.format(det_label))
        exit(1)
  
      obj_width = round(bottom_right_x - top_left_x, 1)
      obj_height = round(bottom_right_y - top_left_y, 1)
  
      coco_det = dict()
      coco_det['image_id'] = img_id
      coco_det['category_id'] = det_label
      coco_det['bbox'] = [round(top_left_x, 1), round(top_left_y, 1), obj_width, obj_height]
      coco_det['score'] = det_conf
      coco_detections.append(coco_det)
  
    return coco_detections
  
  
  def draw_detections(img, coco_detections):
    for det in coco_detections:
      x, y, width, height = det['bbox'] # pylint: disable=invalid-name
      top_left = (int(x), int(y))
      bottom_right = (int(x + width), int(y + height))
      cv2.rectangle(img, top_left, bottom_right, (255, 0, 0), 2)
      cv2.putText(img, '{0}: {1}'.format(det['category_id'], det['score']), top_left,
                  cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 0, 255), 2)
    return img
  
  
  # pylint: disable=too-many-arguments
  def predict_on_video(predictions, path_to_video, classes, show=False, conf_threshold=None, dump_output_video=False,
                       path_to_output_video='output.avi'):
    output = []
  
    cap = cv2.VideoCapture(path_to_video)
    fourcc = cv2.VideoWriter_fourcc(*'XVID')
    width = cap.get(cv2.CAP_PROP_FRAME_WIDTH)
    height = cap.get(cv2.CAP_PROP_FRAME_HEIGHT)
    fps = cap.get(cv2.CAP_PROP_FPS)
    if dump_output_video:
      out = cv2.VideoWriter(path_to_output_video, fourcc, fps, (int(width), int(height)))
    for i, pred in enumerate(predictions):
      if cap.isOpened():
        _, frame = cap.read()
        det = process_image(pred[:, 1:], (width, height), i, conf_threshold, classes)
        output.extend(det)
        img = draw_detections(frame, det)
        if show:
          cv2.imshow('detections', img)
          key = cv2.waitKey(10)
          if key == 27:
            break
        if dump_output_video:
          out.write(img)
      else:
        break
  
    cap.release()
    if dump_output_video:
      out.release()
  
    return output
  
  
  def predict_on_json(predictions, annotation_path, classes, show=False, conf_threshold=None, dump_output_video=False,
                      path_to_output_video='output.avi', width=640, height=480, fps=30):
    annotation_generator, _ = ObjectDetectorJson.json_iterator(annotation_path, classes)
    annotation_data = [pickle.loads(x) for x in annotation_generator()]
  
    output = []
    if dump_output_video:
      fourcc = cv2.VideoWriter_fourcc(*'XVID')
      out = cv2.VideoWriter(path_to_output_video, fourcc, fps, (int(width), int(height)))
    for i, pred in enumerate(predictions):
      annotation = annotation_data[i]
      image_size = annotation['image_size']
      img_path = annotation['image']
      img_id = annotation['image_id']
      det = process_image(pred[:, 1:], image_size, img_id, conf_threshold, classes)
      output.extend(det)
      frame = cv2.imread(img_path)
      frame = cv2.resize(frame, tuple(image_size))
      img = draw_detections(frame, det)
      if show:
        cv2.imshow('detections', img)
        key = cv2.waitKey(10)
        if key == 27:
          break
      if dump_output_video:
        img_resized = cv2.resize(img, (width, height))
        out.write(img_resized)
    if dump_output_video:
      out.release()
  
    return output
  
  
  def infer(config, source, path, conf_threshold=None, dump_to_json=False,
            show=False, dump_output_video=False, path_to_output_video='output.avi'):
    session_config = create_session(config, 'infer')
  
    out_dir = os.path.join(config.MODEL_DIR, config.infer.out_subdir)
    if not os.path.exists(out_dir):
      os.mkdir(out_dir)
  
    run_config = tf.estimator.RunConfig(session_config=session_config)
  
    config.detector_params['log_dir'] = config.MODEL_DIR
  
    predictor = tf.estimator.Estimator(model_fn=detection_model,
                                       params=config.detector_params, model_dir=config.MODEL_DIR, config=run_config)
  
    checkpoint_path = tf.train.latest_checkpoint(config.MODEL_DIR)
  
    basename = os.path.basename(path)
    filename = os.path.splitext(basename)[0]
    name = '{0}.json'.format(filename)
    output_json_path = os.path.join(out_dir, name)
  
    if source == 'video':
      input_data = InputInferData(path, config.input_shape, config.infer.batch_size)
    elif source == 'json':
      input_data = InputValData(batch_size=config.infer.batch_size, input_shape=config.input_shape, json_path=path,
                                classes=config.classes,
                                num_parallel_calls=config.infer.execution.transformer_parallel_calls,
                                prefetch_size=config.infer.execution.transformer_prefetch_size)
  
    predictions = predictor.predict(input_fn=input_data.input_fn, checkpoint_path=checkpoint_path)
    if source == 'video':
      predictions = predict_on_video(predictions, path, config.classes, show, conf_threshold, dump_output_video,
                                     path_to_output_video)
    elif source == 'json':
      predictions = predict_on_json(predictions, path, config.classes, show, conf_threshold, dump_output_video,
                                    path_to_output_video)
  
    if dump_to_json:
      with open(output_json_path, 'w') as output_file:
        json.dump(predictions, output_file, sort_keys=True, indent=4)
  
  
  def main(_):
    args = parse_args()
    cfg = load_module(args.path_to_config)
  
    infer(cfg, args.input_type, args.input, args.conf_threshold, args.dump_to_json, args.show,
          args.dump_output_video, args.path_to_output_video)
  
  
  if __name__ == '__main__':
    tf.logging.set_verbosity(tf.logging.INFO)
    tf.app.run(main)
  